solving split bug TODO: bilingual issue

Former-commit-id: ca894904a230ea50f15ead328500ffb8a3ae6452

SRT.py

@@ -193,8 +193,8 @@ class SRT_script():
     
     def split_seg(self, seg, threshold):
         # evenly split seg to 2 parts and add new seg into self.segments
-        if seg.source_text[0] == ',':
-            seg.source_text = seg.source_text[1:]
+        if seg.source_text[:2] == ', ':
+            seg.source_text = seg.source_text[2:]
         if seg.translation[0] == '，':
             seg.translation = seg.translation[1:]
         source_text = seg.source_text
@@ -205,7 +205,10 @@ class SRT_script():
             src_split_idx = src_commas[len(src_commas)//2] if len(src_commas) % 2 == 1 else src_commas[len(src_commas)//2 - 1]
         else:
             src_space = [m.start() for m in re.finditer(' ', source_text)]
-            src_split_idx = src_space[len(src_space)//2] if len(src_space) % 2 == 1 else src_space[len(src_space)//2 - 1]
+            if len(src_space) > 0: 
+                src_split_idx = src_space[len(src_space)//2] if len(src_space) % 2 == 1 else src_space[len(src_space)//2 - 1]
+            else:
+                src_split_idx = 0
 
         if len(trans_commas) != 0:
             trans_split_idx = trans_commas[len(trans_commas)//2] if len(trans_commas) % 2 == 1 else trans_commas[len(trans_commas)//2 - 1]
@@ -248,7 +251,8 @@ class SRT_script():
 
 
     def check_len_and_split(self, threshold=30):
-        # TODO: if sentence length >= threshold, split this segments to two
+        # DEPRECATED
+        # if sentence length >= threshold, split this segments to two
         segments = []
         for seg in self.segments:
             if len(seg.translation) > threshold:
@@ -262,20 +266,22 @@ class SRT_script():
         pass
 
     def check_len_and_split_range(self, range, threshold=30):
-        # TODO: if sentence length >= threshold, split this segments to two
+        # if sentence length >= threshold, split this segments to two
         start_seg_id = range[0]
         end_seg_id = range[1]
+        extra_len = 0
         segments = []
         for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
             if len(seg.translation) > threshold:
                 seg_list = self.split_seg(seg, threshold)
                 segments += seg_list
+                extra_len += len(seg_list) - 1 
             else:
                 segments.append(seg)
         
         self.segments[start_seg_id-1:end_seg_id] = segments
 
-        return len(segments)
+        return extra_len
     
     def get_source_only(self):
         # return a string with pure source text
@@ -419,7 +425,12 @@ class SRT_script():
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
-            for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id+length]):
+            # for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id+length]):
+            #     f.write(f'{i+idx}\n')
+            #     f.write(seg.get_trans_str())
+            for i, seg in enumerate(self.segments):
+                if i<range[0]-1: continue
+                if i>=range[1] + length:break
                 f.write(f'{i+idx}\n')
                 f.write(seg.get_trans_str())
         pass
@@ -428,7 +439,9 @@ class SRT_script():
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
-            for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id+length]):
+            for i, seg in enumerate(self.segments):
+                if i<range[0]-1: continue
+                if i>=range[1] + length:break
                 f.write(f'{i+idx}\n')
                 f.write(seg.get_bilingual_str())
         pass

pipeline.py

@@ -47,8 +47,8 @@ if args.video_name == 'placeholder' :
         VIDEO_NAME = args.audio_file.split('/')[-1].split('.')[0]
     elif args.srt_file is not None:
         VIDEO_NAME = args.srt_file.split('/')[-1].split('.')[0]
-else:
-    VIDEO_NAME = args.video_name
+    else:
+        VIDEO_NAME = args.video_name
 
 model_name = args.model_name
 
@@ -260,12 +260,12 @@ for sentence, range in tqdm(zip(script_arr, range_arr)):
             flag = True
     # add read-time output back and modify the post-processing by using one batch as an unit.
     srt.set_translation(translate, range, model_name)
-    # add_length = srt.check_len_and_split_range(range)
-    # srt.realtime_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt",range, add_length,segidx)
-    # srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
+    add_length = srt.check_len_and_split_range(range)
+    srt.realtime_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt",range, add_length,segidx)
+    srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
 
-srt.check_len_and_split()
-srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")
+# srt.check_len_and_split()
+# srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")
 # srt.write_srt_file_bilingual(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt")
 
 if not args.only_srt: