solve split problem; add punctuation remove

Former-commit-id: 8125cf63dcc38aa00e2eb10c513daa1d2ca3aeb5

SRT.py

@@ -51,9 +51,18 @@ class SRT_segment(object):
         self.source_text += seg.source_text
         self.translation += seg.translation
         self.end_time_str = seg.end_time_str
+        self.end = seg.end
+        self.end_ms = seg.end_ms
         self.duration = f"{self.start_time_str} --> {self.end_time_str}"
         pass
-    
+
+    def remove_trans_punc(self):
+        # remove punctuations in translation text
+        self.translation = self.translation.replace('，', ' ')
+        self.translation = self.translation.replace('。', ' ')
+        self.translation = self.translation.replace('！', ' ')
+        self.translation = self.translation.replace('？', ' ')
+
     def __str__(self) -> str:
         return  f'{self.duration}\n{self.source_text}\n\n'
     
@@ -96,7 +105,7 @@ class SRT_script():
         merge_list = [] # a list of indices that should be merged e.g. [[0], [1, 2, 3, 4], [5, 6], [7]]
         sentence = []
         for i, seg in enumerate(self.segments):
-            if seg.source_text[-1] == '.':
+            if seg.source_text[-1] in ['.', '!', '?']:
                 sentence.append(i)
                 merge_list.append(sentence)
                 sentence = []
@@ -108,7 +117,11 @@ class SRT_script():
             segments.append(self.merge_segs(idx_list))
 
         self.segments = segments # need memory release?
-        
+    
+    def remove_trans_punctuation(self):
+        # Post-process: remove all punc after translation and split
+        for i, seg in enumerate(self.segments):
+            seg.remove_trans_punc()
 
     def set_translation(self, translate:str, id_range:tuple, model, video_name, video_link=None):
         start_seg_id = id_range[0]
@@ -197,10 +210,12 @@ class SRT_script():
         for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
             # naive way to due with merge translation problem
             # TODO: need a smarter solution
-
-            if i < len(lines):
+            max_num = len(lines)
+            if i < max_num:
                 if "(Note:" in lines[i]: # to avoid note
+                    print(f'remove: {lines[i]}')
                     lines.remove(lines[i])
+                    max_num -= 1
                     if i == len(lines) - 1:
                         break
                 try:
@@ -209,7 +224,7 @@ class SRT_script():
                     seg.translation = lines[i]
 
     
-    def split_seg(self, seg, threshold):
+    def split_seg(self, seg, text_threshold, time_threshold):
         # evenly split seg to 2 parts and add new seg into self.segments
 
         # ignore the initial comma to solve the recursion problem
@@ -221,6 +236,8 @@ class SRT_script():
 
         source_text = seg.source_text
         translation = seg.translation
+
+        # split the text based on commas
         src_commas = [m.start() for m in re.finditer(',', source_text)]
         trans_commas = [m.start() for m in re.finditer('，', translation)]
         if len(src_commas) != 0:
@@ -237,13 +254,18 @@ class SRT_script():
         else:
             trans_split_idx = len(translation)//2
         
+        # split the time duration based on text length
+        time_split_ratio = trans_split_idx/(len(seg.translation) - 1)
+
         src_seg1 = source_text[:src_split_idx]
         src_seg2 = source_text[src_split_idx:]
         trans_seg1 = translation[:trans_split_idx]
         trans_seg2 = translation[trans_split_idx:]
+
         start_seg1 = seg.start
-        end_seg1 = start_seg2 = seg.start + (seg.end - seg.start)/2
+        end_seg1 = start_seg2 = seg.start + (seg.end - seg.start)*time_split_ratio
         end_seg2 = seg.end
+
         seg1_dict = {}
         seg1_dict['text'] = src_seg1
         seg1_dict['start'] = start_seg1
@@ -259,26 +281,26 @@ class SRT_script():
         seg2.translation = trans_seg2
 
         result_list = []
-        if len(seg1.translation) > threshold:
-            result_list += self.split_seg(seg1, threshold)
+        if len(seg1.translation) > text_threshold and (seg1.end - seg1.start) > time_threshold:
+            result_list += self.split_seg(seg1, text_threshold, time_threshold)
         else:
             result_list.append(seg1)
         
-        if len(seg2.translation) > threshold:
-            result_list += self.split_seg(seg2, threshold)
+        if len(seg2.translation) > text_threshold and (seg2.end - seg2.start) > time_threshold:
+            result_list += self.split_seg(seg2, text_threshold, time_threshold)
         else:
             result_list.append(seg2)
 
         return result_list
 
 
-    def check_len_and_split(self, threshold=30):
+    def check_len_and_split(self, text_threshold=30, time_threshold=1.0):
         # DEPRECATED
-        # if sentence length >= threshold, split this segments to two
+        # if sentence length >= threshold and sentence duration > time_threshold, split this segments to two
         segments = []
         for seg in self.segments:
-            if len(seg.translation) > threshold:
-                seg_list = self.split_seg(seg, threshold)
+            if len(seg.translation) > text_threshold and (seg.end - seg.start) > time_threshold:
+                seg_list = self.split_seg(seg, text_threshold, time_threshold)
                 segments += seg_list
             else:
                 segments.append(seg)
@@ -287,73 +309,25 @@ class SRT_script():
 
         pass
 
-    def check_len_and_split_range(self, range, threshold=30):
-        # if sentence length >= threshold, split this segments to two
+    def check_len_and_split_range(self, range, text_threshold=30, time_threshold=1.0):
+        # if sentence length >= text_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)
+            if len(seg.translation) > text_threshold and (seg.end - seg.start) > time_threshold:
+                seg_list = self.split_seg(seg, text_threshold, time_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 extra_len
-    
-    def get_source_only(self):
-        # return a string with pure source text
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result+=f'SENTENCE {i+1}: {seg.source_text}\n\n\n'
-        
-        return result
-    
-    def reform_src_str(self):
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result += f'{i+1}\n'
-            result += str(seg)
-        return result
-
-    def reform_trans_str(self):
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result += f'{i+1}\n'
-            result += seg.get_trans_str()
-        return result
-    
-    def form_bilingual_str(self):
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result += f'{i+1}\n'
-            result += seg.get_bilingual_str()
-        return result
-
-    def write_srt_file_src(self, path:str):
-        # write srt file to path
-        with open(path, "w", encoding='utf-8') as f:
-            f.write(self.reform_src_str())
-        pass
-
-    def write_srt_file_translate(self, path:str):
-        with open(path, "w", encoding='utf-8') as f:
-            f.write(self.reform_trans_str())
-        pass
-
-    def write_srt_file_bilingual(self, path:str):
-        with open(path, "w", encoding='utf-8') as f:
-            f.write(self.form_bilingual_str())
-        pass
 
     def correct_with_force_term(self):
         ## force term correction
-        # TODO: shortcut translation i.e. VA, ob
-        # TODO: variety of translation
 
         # load term dictionary
         with open("./finetune_data/dict_enzh.csv",'r', encoding='utf-8') as f:
@@ -442,8 +416,57 @@ class SRT_script():
             real_word = word.lower()
             n = 0
         return real_word, len(word)+n
+
+
+    ## WRITE AND READ FUNCTIONS ##
+
+    def get_source_only(self):
+        # return a string with pure source text
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result+=f'SENTENCE {i+1}: {seg.source_text}\n\n\n'
+        
+        return result
     
+    def reform_src_str(self):
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result += f'{i+1}\n'
+            result += str(seg)
+        return result
+
+    def reform_trans_str(self):
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result += f'{i+1}\n'
+            result += seg.get_trans_str()
+        return result
+    
+    def form_bilingual_str(self):
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result += f'{i+1}\n'
+            result += seg.get_bilingual_str()
+        return result
+
+    def write_srt_file_src(self, path:str):
+        # write srt file to path
+        with open(path, "w", encoding='utf-8') as f:
+            f.write(self.reform_src_str())
+        pass
+
+    def write_srt_file_translate(self, path:str):
+        with open(path, "w", encoding='utf-8') as f:
+            f.write(self.reform_trans_str())
+        pass
+
+    def write_srt_file_bilingual(self, path:str):
+        with open(path, "w", encoding='utf-8') as f:
+            f.write(self.form_bilingual_str())
+        pass
+
     def realtime_write_srt(self,path,range,length, idx):
+        # DEPRECATED
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
@@ -458,6 +481,7 @@ class SRT_script():
         pass
 
     def realtime_bilingual_write_srt(self,path,range, length,idx):
+        # DEPRECATED
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
@@ -466,4 +490,6 @@ class SRT_script():
                 if i>=range[1] + length :break
                 f.write(f'{i+idx}\n')
                 f.write(seg.get_bilingual_str())
-        pass
+        pass
+    
+

pipeline.py

@@ -49,7 +49,7 @@ if args.video_name == 'placeholder' :
     elif args.audio_file is not None:
         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].split("_")[:-1]
+        VIDEO_NAME = args.srt_file.split('/')[-1].split('.')[0].split("_")[0]
     else:
         VIDEO_NAME = args.video_name
 else:
@@ -257,6 +257,7 @@ for sentence, range in tqdm(zip(script_arr, range_arr)):
     # srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
 
 srt.check_len_and_split()
+srt.remove_trans_punctuation()
 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")