重构 fixer.go 的设置 V2 版本

Signed-off-by: allan716 <[email protected]>

internal/pkg/sub_timeline_fixer/fixer.go

@@ -26,13 +26,13 @@ import (
 )
 
 type SubTimelineFixer struct {
-	fixerConfig  sub_timeline_fiexer.SubTimelineFixerConfig
+	FixerConfig  sub_timeline_fiexer.SubTimelineFixerConfig
 	ffmpegHelper *ffmpeg_helper.FFMPEGHelper
 }
 
 func NewSubTimelineFixer(fixerConfig sub_timeline_fiexer.SubTimelineFixerConfig) *SubTimelineFixer {
 	return &SubTimelineFixer{
-		fixerConfig:  fixerConfig,
+		FixerConfig:  fixerConfig,
 		ffmpegHelper: ffmpeg_helper.NewFFMPEGHelper(),
 	}
 }
@@ -143,7 +143,7 @@ func (s *SubTimelineFixer) GetOffsetTimeV1(infoBase, infoSrc *subparser.FileInfo
 	/*
 		确认两个字幕间的偏移，暂定的方案是两边都连续匹配上 5 个索引，再抽取一个对话的时间进行修正计算
 	*/
-	maxCompareDialogue := s.fixerConfig.MaxCompareDialogue
+	maxCompareDialogue := s.FixerConfig.MaxCompareDialogue
 	// 基线的长度
 	_, docsLength := tfidf.Dims()
 	var matchIndexList = make([]MatchIndex, 0)
@@ -337,14 +337,14 @@ func (s *SubTimelineFixer) GetOffsetTimeV1(infoBase, infoSrc *subparser.FileInfo
 	matchIndexLineCount := len(matchIndexList) * maxCompareDialogue
 	//perMatch := float64(matchIndexLineCount) / float64(len(infoSrc.DialoguesEx))
 	perMatch := float64(matchIndexLineCount) / float64(len(baseCorpus))
-	if perMatch < s.fixerConfig.MinMatchedPercent {
-		tmpContent := infoSrc.Name + fmt.Sprintf(" Sequence match %d dialogues (< %f%%), Skip,", s.fixerConfig.MaxCompareDialogue, s.fixerConfig.MinMatchedPercent*100) + fmt.Sprintf(" %f%% ", perMatch*100)
+	if perMatch < s.FixerConfig.MinMatchedPercent {
+		tmpContent := infoSrc.Name + fmt.Sprintf(" Sequence match %d dialogues (< %f%%), Skip,", s.FixerConfig.MaxCompareDialogue, s.FixerConfig.MinMatchedPercent*100) + fmt.Sprintf(" %f%% ", perMatch*100)
 
 		debugInfos = append(debugInfos, tmpContent)
 
 		log_helper.GetLogger().Infoln(tmpContent)
 	} else {
-		tmpContent := infoSrc.Name + fmt.Sprintf(" Sequence match %d dialogues,", s.fixerConfig.MaxCompareDialogue) + fmt.Sprintf(" %f%% ", perMatch*100)
+		tmpContent := infoSrc.Name + fmt.Sprintf(" Sequence match %d dialogues,", s.FixerConfig.MaxCompareDialogue) + fmt.Sprintf(" %f%% ", perMatch*100)
 
 		debugInfos = append(debugInfos, tmpContent)
 
@@ -359,7 +359,7 @@ func (s *SubTimelineFixer) GetOffsetTimeV1(infoBase, infoSrc *subparser.FileInfo
 		}
 	}
 	// 虽然有条件判断是认为有问题的，但是返回值还是要填写除去的
-	if perMatch < s.fixerConfig.MinMatchedPercent {
+	if perMatch < s.FixerConfig.MinMatchedPercent {
 		return false, newMean, newSd, nil
 	}
 
@@ -396,12 +396,12 @@ func (s *SubTimelineFixer) GetOffsetTimeV2(baseUnit, srcUnit *sub_helper.SubUnit
 
 	srcVADLen := len(srcUnit.VADList)
 	// 滑动窗口的长度
-	srcWindowLen := int(float64(srcVADLen) * MatchPer)
+	srcWindowLen := int(float64(srcVADLen) * s.FixerConfig.MatchPer)
 	srcSlideLen := srcVADLen - srcWindowLen
 	// 窗口可以滑动的长度
 	srcSlideLenHalf := srcSlideLen / 2
 	//
-	oneStep := srcSlideLenHalf / CompareParts
+	oneStep := srcSlideLenHalf / s.FixerConfig.CompareParts
 	if srcSlideLen <= 0 {
 		srcSlideLen = 1
 	}
@@ -425,7 +425,7 @@ func (s *SubTimelineFixer) GetOffsetTimeV2(baseUnit, srcUnit *sub_helper.SubUnit
 		if bUseSubOrAudioAsBase == false {
 			// 使用 音频 来进行匹配
 			// 去掉头和尾，具体百分之多少，见 FrontAndEndPerBase
-			audioCutLen := int(float64(len(inData.AudioVADList)) * FrontAndEndPerBase)
+			audioCutLen := int(float64(len(inData.AudioVADList)) * s.FixerConfig.FrontAndEndPerBase)
 
 			offsetIndex, score = fffAligner.Fit(inData.AudioVADList[audioCutLen:len(inData.AudioVADList)-audioCutLen], inData.SrcUnit.GetVADFloatSlice()[inData.OffsetIndex:srcWindowLen+inData.OffsetIndex])
 			realOffsetIndex := offsetIndex + audioCutLen
@@ -466,10 +466,10 @@ func (s *SubTimelineFixer) GetOffsetTimeV2(baseUnit, srcUnit *sub_helper.SubUnit
 		return nil
 	}
 	// -------------------------------------------------
-	antPool, err := ants.NewPoolWithFunc(FixThreads, func(inData interface{}) {
+	antPool, err := ants.NewPoolWithFunc(s.FixerConfig.FixThreads, func(inData interface{}) {
 		data := inData.(InputData)
 		defer data.Wg.Done()
-		ctx, cancel := context.WithTimeout(context.Background(), SubOneUnitProcessTimeOut)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(s.FixerConfig.SubOneUnitProcessTimeOut)*time.Second)
 		defer cancel()
 
 		done := make(chan error, 1)
@@ -523,7 +523,7 @@ func (s *SubTimelineFixer) GetOffsetTimeV2(baseUnit, srcUnit *sub_helper.SubUnit
 	}
 	wg.Wait()
 	// 这里可能遇到匹配的时候没有能够执行够 CompareParts 次，有可能是负数跳过或者时间转换失败导致，前者为主（可能是这两个就是一个东西的时候，或者说没有时间轴偏移的时候）
-	if insertIndex < CompareParts/2 {
+	if insertIndex < s.FixerConfig.CompareParts/2 {
 		return false, 0, 0, nil
 	}
 	outCorrelationFixResult := s.calcMeanAndSD(tmpStartDiffTimeListEx, tmpStartDiffTimeList)
@@ -601,14 +601,6 @@ func (s *SubTimelineFixer) calcMeanAndSD(startDiffTimeList stat.Float64Slice, tm
 }
 
 const FixMask = "-fix"
-const SubOneUnitProcessTimeOut = 60 * 5 * time.Second // 字幕时间轴校正一个单元的超时时间
-
-// 字幕匹配字幕 GetVADInfoFeatureFromSubNew
-const FrontAndEndPerBase = 0.15 // 前百分之 15 和后百分之 15 都不进行识别
-const FrontAndEndPerSrc = 0.0   // 前百分之 20 和后百分之 20 都不进行识别
-const MatchPer = 0.7
-const CompareParts = 5
-const FixThreads = 3 // 字幕校正的并发线程
 
 var mutexFixV2 sync.Mutex
 

internal/pkg/sub_timeline_fixer/fixer_test.go

@@ -324,13 +324,6 @@ func TestGetOffsetTimeV1(t *testing.T) {
 			want: 0, wantErr: false},
 	}
 
-	timelineFixer := NewSubTimelineFixer(sub_timeline_fiexer.SubTimelineFixerConfig{
-		MaxCompareDialogue: 3,
-		MaxStartTimeDiffSD: 0.1,
-		MinMatchedPercent:  0.1,
-		MinOffset:          0.1,
-	})
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 
@@ -567,13 +560,6 @@ func TestGetOffsetTimeV2_BaseSub(t *testing.T) {
 			want: 0, wantErr: false},
 	}
 
-	timelineFixer := NewSubTimelineFixer(sub_timeline_fiexer.SubTimelineFixerConfig{
-		MaxCompareDialogue: 3,
-		MaxStartTimeDiffSD: 0.1,
-		MinMatchedPercent:  0.1,
-		MinOffset:          0.1,
-	})
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 
@@ -609,7 +595,7 @@ func TestGetOffsetTimeV2_BaseSub(t *testing.T) {
 			//	t.Fatal(err)
 			//}
 			//baseUnitOld := baseUnitListOld[0]
-			baseUnitNew, err := sub_helper.GetVADInfoFeatureFromSubNew(infoBase, FrontAndEndPerBase)
+			baseUnitNew, err := sub_helper.GetVADInfoFeatureFromSubNew(infoBase, timelineFixer.FixerConfig.FrontAndEndPerBase)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -620,7 +606,7 @@ func TestGetOffsetTimeV2_BaseSub(t *testing.T) {
 			//	t.Fatal(err)
 			//}
 			//srcUnitOld := srcUnitListOld[0]
-			srcUnitNew, err := sub_helper.GetVADInfoFeatureFromSubNew(infoSrc, FrontAndEndPerSrc)
+			srcUnitNew, err := sub_helper.GetVADInfoFeatureFromSubNew(infoSrc, timelineFixer.FixerConfig.FrontAndEndPerSrc)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -745,9 +731,6 @@ func TestGetOffsetTimeV2_BaseAudio(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			s := &SubTimelineFixer{
-				fixerConfig: tt.fields.fixerConfig,
-			}
 
 			bFind, infoSrc, err := subParserHub.DetermineFileTypeFromFile(tt.args.subFilePath)
 			if err != nil {
@@ -762,7 +745,7 @@ func TestGetOffsetTimeV2_BaseAudio(t *testing.T) {
 			*/
 			//sub_helper.MergeMultiDialogue4EngSubtitle(infoSrc)
 			// Src，截取的部分要小于 Base 的部分
-			srcUnitNew, err := sub_helper.GetVADInfoFeatureFromSubNew(infoSrc, FrontAndEndPerSrc)
+			srcUnitNew, err := sub_helper.GetVADInfoFeatureFromSubNew(infoSrc, timelineFixer.FixerConfig.FrontAndEndPerSrc)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -776,7 +759,7 @@ func TestGetOffsetTimeV2_BaseAudio(t *testing.T) {
 			}
 
 			println("-------New--------")
-			got, got1, sd, err := s.GetOffsetTimeV2(nil, srcUnitNew, audioVADInfos)
+			got, got1, sd, err := timelineFixer.GetOffsetTimeV2(nil, srcUnitNew, audioVADInfos)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("GetOffsetTimeV3() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -795,7 +778,7 @@ func TestGetOffsetTimeV2_BaseAudio(t *testing.T) {
 				t.Errorf("GetOffsetTimeV1() got = %v, want %v", got, tt.want)
 			}
 
-			_, err = s.FixSubTimeline(infoSrc, got1, tt.args.subFilePath+FixMask+infoSrc.Ext)
+			_, err = timelineFixer.FixSubTimeline(infoSrc, got1, tt.args.subFilePath+FixMask+infoSrc.Ext)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -804,3 +787,18 @@ func TestGetOffsetTimeV2_BaseAudio(t *testing.T) {
 		})
 	}
 }
+
+var timelineFixer = NewSubTimelineFixer(sub_timeline_fiexer.SubTimelineFixerConfig{
+	// V1
+	MaxCompareDialogue: 3,
+	MaxStartTimeDiffSD: 0.1,
+	MinMatchedPercent:  0.1,
+	MinOffset:          0.1,
+	// V2
+	SubOneUnitProcessTimeOut: 5 * 60,
+	FrontAndEndPerBase:       0.15,
+	FrontAndEndPerSrc:        0.0,
+	MatchPer:                 0.7,
+	CompareParts:             5,
+	FixThreads:               3,
+})

internal/types/sub_timeline_fiexer/sub_timeline_fixer_config.go

@@ -1,10 +1,18 @@
 package sub_timeline_fiexer
 
 type SubTimelineFixerConfig struct {
+	// V1 的设置
 	MaxCompareDialogue int     // 最大需要匹配的连续对白，默认3
 	MaxStartTimeDiffSD float64 // 对白开始时间的统计 SD 最大误差，超过则不进行修正
 	MinMatchedPercent  float64 // 两个文件的匹配百分比（src/base），高于这个才比例进行修正
 	MinOffset          float64 // 超过这个(+-)偏移的时间轴才校正，否则跳过，单位秒
+	// V2 的设置
+	SubOneUnitProcessTimeOut int     // 字幕时间轴校正一个单元的超时时间，单位秒
+	FrontAndEndPerBase       float64 // 前百分之 15 和后百分之 15 都不进行识别
+	FrontAndEndPerSrc        float64 // 前百分之 20 和后百分之 20 都不进行识别
+	MatchPer                 float64 // SrcSub 滑动窗体的占比
+	CompareParts             int     // 滑动窗体分段次数
+	FixThreads               int     // 字幕校正的并发线程
 }
 
 // CheckDefault 检测默认值（比如某些之默认不能为0），不对就重置到默认值上