Merge branch 'eason/refactor' into MergeFix

Former-commit-id: 5cf3410f6825636ac35226749d95ec63b4af6072

SRT.py

@@ -1,7 +1,6 @@
 from datetime import timedelta
-import os
-import whisper
 from csv import reader
+from datetime import datetime
 import re
 import openai
 
@@ -9,27 +8,41 @@ class SRT_segment(object):
     def __init__(self, *args) -> None:
         if isinstance(args[0], dict):
             segment = args[0]
-            start_ms = int((segment['start']*100)%100*10)
-            end_ms = int((segment['end']*100)%100*10)
-            start_time = str(timedelta(seconds=int(segment['start']), milliseconds=start_ms))
-            end_time = str(timedelta(seconds=int(segment['end']), milliseconds=end_ms))
-            if start_ms == 0:
-                self.start_time_str = str(0)+start_time.split('.')[0]+',000'
+            self.start = segment['start']
+            self.end = segment['end']
+            self.start_ms = int((segment['start']*100)%100*10)
+            self.end_ms = int((segment['end']*100)%100*10)
+
+            if self.start_ms == self.end_ms and int(segment['start']) == int(segment['end']): # avoid empty time stamp
+                self.end_ms+=500
+
+            self.start_time = timedelta(seconds=int(segment['start']), milliseconds=self.start_ms)
+            self.end_time = timedelta(seconds=int(segment['end']), milliseconds=self.end_ms)
+            if self.start_ms == 0:
+                self.start_time_str = str(0)+str(self.start_time).split('.')[0]+',000'
             else:
-                self.start_time_str = str(0)+start_time.split('.')[0]+','+start_time.split('.')[1][:3]
-            if end_ms == 0:
-                self.end_time_str = str(0)+end_time.split('.')[0]+',000'
+                self.start_time_str = str(0)+str(self.start_time).split('.')[0]+','+str(self.start_time).split('.')[1][:3]
+            if self.end_ms == 0:
+                self.end_time_str = str(0)+str(self.end_time).split('.')[0]+',000'
             else:
-                self.end_time_str = str(0)+end_time.split('.')[0]+','+end_time.split('.')[1][:3]
-            self.source_text = segment['text'][1:]
+                self.end_time_str = str(0)+str(self.end_time).split('.')[0]+','+str(self.end_time).split('.')[1][:3]
+            self.source_text = segment['text']
             self.duration = f"{self.start_time_str} --> {self.end_time_str}"
             self.translation = ""
 
         elif isinstance(args[0], list):
-            self.source_text = args[0][2][:-1]
+            self.source_text = args[0][2]
             self.duration = args[0][1]
             self.start_time_str = self.duration.split(" --> ")[0]
             self.end_time_str = self.duration.split(" --> ")[1]
+
+            # parse the time to float
+            self.start_ms = int(self.start_time_str.split(',')[1])/10
+            self.end_ms = int(self.end_time_str.split(',')[1])/10
+            start_list = self.start_time_str.split(',')[0].split(':')
+            self.start = int(start_list[0])*3600 + int(start_list[1])*60 + int(start_list[2]) + self.start_ms/100
+            end_list = self.end_time_str.split(',')[0].split(':')
+            self.end = int(end_list[0])*3600 + int(end_list[1])*60 + int(end_list[2]) + self.end_ms/100
             self.translation = ""
     
     def merge_seg(self, seg):
@@ -64,7 +77,7 @@ class SRT_script():
         for i in range(len(script_lines)):
             if i % 4 == 0:
                 segments.append(list(script_lines[i:i+4]))
-        
+
         return cls(segments)
 
     def merge_segs(self, idx_list) -> SRT_segment:
@@ -152,12 +165,70 @@ class SRT_script():
                     #print(lines[i])
         pass
     
-    def split_seg(self, seg_id):
+    def split_seg(self, seg, threshold):
         # TODO: evenly split seg to 2 parts and add new seg into self.segments
-        pass
+        source_text = seg.source_text
+        translation = seg.translation
+        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:
+            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(trans_commas) != 0:
+            trans_split_idx = trans_commas[len(src_commas)//2] if len(trans_commas) % 2 == 1 else trans_commas[len(trans_commas)//2 - 1]
+        else:
+            trans_split_idx = len(translation)//2
+        
+        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_seg2 = seg.end
+        seg1_dict = {}
+        seg1_dict['text'] = src_seg1
+        seg1_dict['start'] = start_seg1
+        seg1_dict['end'] = end_seg1
+        seg1 = SRT_segment(seg1_dict)
+        seg1.translation = trans_seg1
 
-    def check_len_and_split(self, threshold):
+        seg2_dict = {}
+        seg2_dict['text'] = src_seg2
+        seg2_dict['start'] = start_seg2
+        seg2_dict['end'] = end_seg2
+        seg2 = SRT_segment(seg2_dict)
+        seg2.translation = trans_seg2
+
+        result_list = []
+        if len(seg1.translation) > threshold:
+            result_list += self.split_seg(seg1, threshold)
+        else:
+            result_list.append(seg1)
+        
+        if len(seg2.translation) > threshold:
+            result_list += self.split_seg(seg2, threshold)
+        else:
+            result_list.append(seg2)
+
+        return result_list
+
+
+    def check_len_and_split(self, threshold=30):
         # TODO: if sentence length >= threshold, split this segments to two
+        segments = []
+        for seg in self.segments:
+            if len(seg.translation) > threshold:
+                seg_list = self.split_seg(seg, threshold)
+                segments += seg_list
+            else:
+                segments.append(seg)
+        
+        self.segments = segments
+
         pass
 
     def get_source_only(self):
@@ -211,29 +282,90 @@ class SRT_script():
         # TODO: variety of translation
 
         # load term dictionary
-        with open("finetune_data/dict.csv",'r', encoding='utf-8') as f:
-            csv_reader = reader(f)
-            term_dict = {rows[0]:rows[1] for rows in csv_reader}
+        with open("./finetune_data/dict_enzh.csv",'r', encoding='utf-8') as f:
+            term_enzh_dict = {rows[0]:rows[1] for rows in reader(f)}
 
         # change term
         for seg in self.segments:
-            ready_words = re.sub('\n', '\n ', seg.source_text).split(" ")
+            ready_words = seg.source_text.split(" ")
             for i in range(len(ready_words)):
                 word = ready_words[i]
-                if word[-2:] == ".\n" :
-                    if word[:-2].lower() in term_dict :
-                        new_word = word.replace(word[:-2], term_dict.get(word[:-2].lower())) + ' '
-                        ready_words[i] = new_word
-                    else:
-                        ready_words[i] = word + ' '
-                elif word.lower() in term_dict :
-                    new_word = word.replace(word,term_dict.get(word.lower())) + ' '
-                    ready_words[i] = new_word
-                else :
-                    ready_words[i]= word + ' '
-            seg.source_text = re.sub('\n ', '\n', "".join(ready_words))
+                [real_word, pos] = self.get_real_word(word)
+                if real_word in term_enzh_dict:
+                    new_word = word.replace(word[:pos], term_enzh_dict.get(real_word))
+                else:
+                    new_word = word
+                ready_words[i] = new_word
+            seg.source_text = " ".join(ready_words)
         pass
 
-    
+    def spell_check_term(self):
+        ## known bug: I've will be replaced because i've is not in the dict
+
+        
+        import enchant
+        dict = enchant.Dict('en_US')
+        term_spellDict = enchant.PyPWL('./finetune_data/dict_freq.txt')
 
-    
+        for seg in self.segments:
+            ready_words = seg.source_text.split(" ")
+            for i in range(len(ready_words)):
+                word = ready_words[i]
+                [real_word, pos] = self.get_real_word(word)
+                if not dict.check(real_word):
+                    suggest = term_spellDict.suggest(real_word)
+                    if suggest:  # relax spell check
+                        new_word = word.replace(word[:pos],suggest[0])
+                else: 
+                    new_word = word
+                ready_words[i] = new_word
+            seg.source_text = " ".join(ready_words)
+        pass
+
+    def spell_correction(self, word:str, arg:int):
+        try:
+            arg in [0,1]
+        except ValueError:
+            print('only 0 or 1 for argument')
+
+
+        def uncover(word:str):
+            if word[-2:] == ".\n":
+                real_word = word[:-2].lower()
+                n = -2
+            elif word[-1:] in [".", "\n", ",", "!", "?"]:
+                real_word = word[:-1].lower()
+                n = -1
+            else:
+                real_word = word.lower()
+                n = 0
+            return real_word, len(word)+n
+        
+        real_word = uncover(word)[0]
+        pos = uncover(word)[1]
+        new_word = word
+        if arg == 0:  # term translate mode
+            with open("finetune_data/dict_enzh.csv",'r', encoding='utf-8') as f:
+                term_enzh_dict = {rows[0]:rows[1] for rows in reader(f)}
+            if real_word in term_enzh_dict:
+                new_word = word.replace(word[:pos], term_enzh_dict.get(real_word))
+        elif arg == 1:  # term spell check mode
+            import enchant
+            dict = enchant.Dict('en_US')
+            term_spellDict = enchant.PyPWL('./finetune_data/dict_freq.txt')
+            if not dict.check(real_word):
+                if term_spellDict.suggest(real_word):  # relax spell check
+                    new_word = word.replace(word[:pos],term_spellDict.suggest(real_word)[0])
+        return new_word
+    
+    def get_real_word(self, word:str):
+        if word[-2:] == ".\n":
+            real_word = word[:-2].lower()
+            n = -2
+        elif word[-1:] in [".", "\n", ",", "!", "?"]:
+            real_word = word[:-1].lower()
+            n = -1
+        else:
+            real_word = word.lower()
+            n = 0
+        return real_word, len(word)+n

finetune_data/{dict.csv → dict_enzh.csv}

@@ -1,4 +1,4 @@
-barracks,兵营
+barracks,兵营
 engineering bay,工程站
 forge,锻炉
 blink,闪现

finetune_data/dict_freq.csv

@@ -0,0 +1,176 @@
+barracks,1
+engineering bay,1
+forge,1
+blink,1
+evolution chamber,1
+cybernetics core,1
+enhanced shockwaves,1
+gravitic boosters,1
+armory,1
+robotics bay,1
+twilight council,1
+fusion core,1
+fleet beacon,1
+factory,1
+ghost academy,1
+infestation pit,1
+robotics facility,1
+stargate,1
+starport,1
+archon,1
+smart servos,1
+gateway,1
+warpgate,1
+immortal,1
+zealot,1
+nydus network,1
+nydus worm,1
+hydralisk,1
+grooved spines,1
+muscular augments,1
+hydralisk den,1
+planetary fortress,1
+battle cruiser,1
+weapon refit,1
+brood lord,1
+greater spire,1
+anabolic synthesis,1
+cyclone,1
+bunker,1
+lurker,1
+seismic spines,1
+adaptive talons,1
+lurker den,1
+widow mine,1
+ground carapace,1
+high templar,1
+shield battery,1
+observer,1
+baneling,1
+centrifugal hooks,1
+baneling nest,1
+raven,1
+combat shield,1
+shield,1
+lair,1
+missile turret,1
+spore crawler,1
+supply depot,1
+overlord,1
+pneumatized carapace,1
+mutalisk,1
+spire,1
+viper,1
+flyer attacks,1
+flyer carapace,1
+tempest,1
+tectonic destabilizers,1
+phoenix,1
+anion pulse-crystals,1
+corruptor,1
+infestor,1
+pathogen glands,1
+zergling,1
+spawning pool,1
+metabolic boost,1
+spine crawler,1
+marauder,1
+ghost,1
+arm silo with nuke,1
+carrier,1
+hellion,1
+hellbat,1
+ravager,1
+nexus,1
+hatchery,1
+command center,1
+neosteel armor,1
+hi-sec auto tracking,1
+ship weapons,1
+charge,1
+liberator,1
+advanced ballistics,1
+melee attacks,1
+colossus,1
+extended thermal lance,1
+creep tumor,1
+tech lab,1
+air armor,1
+air weapons,1
+adrenal glands,1
+mule,1
+infernal pre-igniter,1
+thor,1
+warp prism,1
+gravitic drive,1
+dragoon,1
+cocoon,1
+larva,1
+mothership,1
+burrow,1
+changeling,1
+ultralisk,1
+chitinous plating,1
+ultralisk cavern,1
+drone,1
+scv,1
+queen,1
+banshee,1
+hyperflight rotors,1
+photon cannon,1
+missile attacks,1
+assimilator,1
+extractor,1
+refinery,1
+roach,1
+marine,1
+sensor tower,1
+infantry armor,1
+infantry weapons,1
+hive,1
+psionic storm,1
+templar archives,1
+sentry,1
+ground armor,1
+ground weapons,1
+adept,1
+resonating glaives,1
+reactor,1
+pylon,1
+reaper,1
+drilling claws,1
+swarm host,1
+mag-field accelerator,1
+siege tank,1
+probe,1
+corvid reactor,1
+neural parasite,1
+viking,1
+oracle,1
+broodling,1
+locust,1
+mothership core,1
+orbital command,1
+stimpack,1
+void ray,1
+flux vanes,1
+overseer,1
+ignite afterburners,1
+dark templar,1
+shadow stride,1
+dark shrine,1
+cloaking field,1
+personal cloaking,1
+medivac dropship,1
+vehicle and ship plating,1
+vehicle weapons,1
+war hound,1
+roach warren,1
+tunneling claws,1
+glial reconstitution,1
+concussive shells,1
+stalker,1
+disruptor,1
+zerg,1
+protross,1
+terran,1

finetune_data/dict_freq.txt

@@ -0,0 +1,177 @@
+barracks
+engineering bay
+forge
+blink
+evolution chamber
+cybernetics core
+enhanced shockwaves
+gravitic boosters
+armory
+robotics bay
+twilight council
+fusion core
+fleet beacon
+factory
+ghost academy
+infestation pit
+robotics facility
+stargate
+starport
+archon
+smart servos
+gateway
+warpgate
+immortal
+zealot
+nydus network
+nydus worm
+hydralisk
+grooved spines
+muscular augments
+hydralisk den
+planetary fortress
+battle cruiser
+weapon refit
+brood lord
+greater spire
+anabolic synthesis
+cyclone
+bunker
+lurker
+seismic spines
+adaptive talons
+lurker den
+widow mine
+ground carapace
+high templar
+shield battery
+observer
+baneling
+centrifugal hooks
+baneling nest
+raven
+combat shield
+shield
+lair
+missile turret
+spore crawler
+supply depot
+overlord
+pneumatized carapace
+mutalisk
+spire
+viper
+flyer attacks
+flyer carapace
+tempest
+tectonic destabilizers
+phoenix
+anion pulse-crystals
+corruptor
+infestor
+pathogen glands
+zergling
+spawning pool
+metabolic boost
+spine crawler
+marauder
+ghost
+arm silo with nuke
+carrier
+hellion
+hellbat
+ravager
+nexus
+hatchery
+command center
+neosteel armor
+hi-sec auto tracking
+ship weapons
+charge
+liberator
+advanced ballistics
+melee attacks
+colossus
+extended thermal lance
+creep tumor
+tech lab
+air armor
+air weapons
+adrenal glands
+mule
+infernal pre-igniter
+thor
+warp prism
+gravitic drive
+dragoon
+cocoon
+larva
+mothership
+burrow
+changeling
+ultralisk
+chitinous plating
+ultralisk cavern
+drone
+scv
+queen
+banshee
+hyperflight rotors
+photon cannon
+missile attacks
+assimilator
+extractor
+refinery
+roach
+marine
+sensor tower
+infantry armor
+infantry weapons
+hive
+psionic storm
+templar archives
+sentry
+ground armor
+ground weapons
+adept
+resonating glaives
+reactor
+pylon
+reaper
+drilling claws
+swarm host
+mag-field accelerator
+siege tank
+probe
+corvid reactor
+neural parasite
+viking
+oracle
+broodling
+locust
+mothership core
+orbital command
+stimpack
+void ray
+flux vanes
+overseer
+ignite afterburners
+dark templar
+shadow stride
+dark shrine
+cloaking field
+personal cloaking
+medivac dropship
+vehicle and ship plating
+vehicle weapons
+war hound
+roach warren
+tunneling claws
+glial reconstitution
+concussive shells
+stalker
+disruptor
+zerg
+protross
+terran
+starcraft

pipeline.py

@@ -2,7 +2,6 @@ import openai
 from pytube import YouTube
 import argparse
 import os
-import whisper
 from tqdm import tqdm
 from SRT import SRT_script
 import stable_whisper
@@ -126,6 +125,7 @@ else:
 
 # srt preprocess
 srt.form_whole_sentence()
+srt.spell_check_term()
 srt.correct_with_force_term()
 srt.write_srt_file_src(srt_file_en)
 script_input = srt.get_source_only()
@@ -198,7 +198,9 @@ for s, range in tqdm(zip(script_arr, range_arr)):
 
     srt.set_translation(translate, range)
 
+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:
     assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")