Merge branch 'eason/main' into Macrodove/suggestionFunc

SRT.py

@@ -402,7 +402,7 @@ class SRT_script():
         for word in keywords:
             for i, seg in enumerate(self.segments):
                 if word in seg.source_text.lower():
-                    seg.source_text = seg.source_text.lower().replace(word, term_enzh_dict.get(word))
+                    seg.source_text = re.sub(fr"({word}es|{word}s?)\b", "{}".format(term_enzh_dict.get(word)), seg.source_text, flags=re.IGNORECASE)
                     logging.info("replace term: " + word + " --> " + term_enzh_dict.get(word) + " in time stamp {}".format(i+1))
                     logging.info("source text becomes: " + seg.source_text)
                     
@@ -436,13 +436,13 @@ class SRT_script():
             for i in range(len(ready_words)):
                 word = ready_words[i]
                 [real_word, pos] = self.get_real_word(word)
-                if not dict.check(word[:pos]):
+                if not dict.check(word[:pos]) and not term_spellDict.check(real_word):
                     new_word = word.replace(word[:pos],self.fetchfunc(word[:pos],0.5))
                     
-                    logging.info(word + "\t" + self.fetchfunc(word[:pos],0.5) + "\t" + str(enchant.utils.levenshtein(word, self.fetchfunc(word[:pos],0.5)))+'\n')
+                    logging.info(real_word + "\t" + self.fetchfunc(word[:pos],0.5) + "\t" + str(enchant.utils.levenshtein(real_word, self.fetchfunc(word[:pos],0.5)))+'\n')
                                 
                     #suggest = term_spellDict.suggest(real_word)
-                    #if suggest and enchant.utils.levenshtein(word, suggest[0]) < (len(word)+len(suggest[0]))/4:  # relax spell check
+                    #if suggest and enchant.utils.levenshtein(real_word, suggest[0]) < (len(real_word)+len(suggest[0]))/4:  # relax spell check
                         
                         # with open("dislog.log","a") as log:
                         #     if not os.path.exists("dislog.log"):

dict_util.py

@@ -7,11 +7,12 @@ import pickle
 # 1_2_3, 1 is action, 2 is supply object, 3 is accept object
 def update_dict_csv(term_dict, f):
     for rows in csv.reader(f):
-        if rows[0] in term_dict:
-            if rows[1] not in term_dict[rows[0]]:
-                term_dict[rows[0]] = term_dict[rows[0]]+[rows[1]]
+        word = rows[0].lower()
+        if word in term_dict:
+            if rows[1] not in term_dict[word]:
+                term_dict[word] = term_dict[word]+[rows[1]]
         else:
-            term_dict[rows[0]]=[rows[1]]
+            term_dict[word]=[rows[1]]
     pass
 
 def export_dict_csv(term_dict, f):
@@ -21,17 +22,14 @@ def export_dict_csv(term_dict, f):
 
 def save_dict_pickle(term_dict, f):
     pickle.dump(term_dict, f, pickle.HIGHEST_PROTOCOL)
+    pass
 
 def update_csv_pickle(pickle_f, csv_f):
     term_dict = pickle.load(pickle_f)
-    for rows in csv.reader(csv_f):
-        if rows[0] in term_dict:
-            if rows[1] not in term_dict[rows[0]]:
-                term_dict[rows[0]] = term_dict[rows[0]]+[rows[1]]
-        else:
-            term_dict[rows[0]]=[rows[1]]
+    update_dict_csv(term_dict, csv_f)
     #save to pickle file, highest protocal to get better performance
     pickle.dump(term_dict, pickle_f, pickle.HIGHEST_PROTOCOL)
+    pass
 
 
 #demo

finetune_data/dict_enzh.csv

@@ -179,4 +179,18 @@ stalker,追猎
 disruptor,自爆球
 zerg,虫族
 protross,神族
-terran,人族
+terran,人族
+gas,气矿
+cannon rush,野炮台
+supply block,卡人口
+macro,运营
+natural expansion,开矿
+roach warren,蟑螂虫巢
+proxy,野
+cyber core,控制芯核
+prismatic alignment,充能射线
+flooding,余钱
+chrono boost,星空加速
+cheese,狗
+supply,人口
+war prism,棱镜

finetune_data/dict_freq.txt

@@ -23,8 +23,8 @@ gateway
 warpgate
 immortal
 zealot
-nydus network
-nydus worm
+nydus
+worm
 hydralisk
 grooved spines
 muscular augments
@@ -173,12 +173,12 @@ concussive shells
 stalker
 disruptor
 zerg
-protross
+protoss
 terran
 starcraft
 TvT
 Maxpax
-showtime
+ShowTime
 PvP
 ZvZ
 TvZ
@@ -187,7 +187,7 @@ ZvP
 PvZ
 PvT
 ZvT
-Florencio
 cybercore
 nest
-follow-up
+follow-up
+robo

finetune_data/test.csv

@@ -0,0 +1,193 @@
+barracks,['兵营']
+zerg,['虫族']
+protoss,['神族']
+terran,['人族']
+engineering bay,['工程站']
+forge,['锻炉']
+blink,['闪现']
+evolution chamber,['进化腔']
+cybernetics core,['控制芯核']
+cybercore,['控制芯核']
+enhanced shockwaves,['EMP范围']
+gravitic boosters,['ob速度']
+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,['大龙']
+broodling,"['巢虫', '小虫子']"
+greater spire,['大龙塔']
+anabolic synthesis,['大牛速度']
+cyclone,['导弹车']
+bunker,['地堡']
+lurker,['地刺']
+seismic spines,['地刺射程']
+adaptive talons,['地刺速埋']
+lurker den,['地刺塔']
+widow mine,['地雷']
+ground carapace,['地面单位甲壳等级']
+high templar,['电兵']
+shield battery,['电池']
+observer,"['叮当', 'OB']"
+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,['炮台']
+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,['先知']
+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,['自爆球']
+protross,['神族']
+gas,['气矿']
+hatch,['基地']
+supply,['人口']
+macro,['运营']
+natural,['二矿']
+proxy,['前置']
+kite,['甩']
+cyber core,['控制芯核']
+saturate,['满采']
+nydusworm,['虫洞']
+float,['余钱']
+transfuse,['加血']
+chrono boost,['加速']
+war prism,['棱镜']

pipeline.py

@@ -9,6 +9,7 @@ import whisper
 from srt2ass import srt2ass
 import logging
 from datetime import datetime
+import torch
 
 import subprocess
 
@@ -109,7 +110,10 @@ def get_srt_class(srt_file_en, result_path, video_name, audio_path, audio_file =
 
             # use stable-whisper
             elif method == "stable":
-                model = stable_whisper.load_model(whisper_model)
+
+                # use cuda if available
+                devices = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
+                model = stable_whisper.load_model(whisper_model, device = devices)
                 transcript = model.transcribe(audio_path, regroup = False, initial_prompt="Hello, welcome to my lecture. Are you good my friend?")
                 (
                     transcript
@@ -265,7 +269,9 @@ def main():
 
     audio_path, audio_file, video_path, VIDEO_NAME = get_sources(args, DOWNLOAD_PATH, RESULT_PATH, VIDEO_NAME)
 
-    logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler("{}/{}_{}.log".format(args.log_dir, VIDEO_NAME, datetime.now().strftime("%m%d%Y_%H%M%S")))], encoding='utf-8')
+    if not os.path.exists(args.log_dir):
+        os.makedirs(args.log_dir)
+    logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler("{}/{}_{}.log".format(args.log_dir, VIDEO_NAME, datetime.now().strftime("%m%d%Y_%H%M%S")), 'w', encoding='utf-8')])
     logging.info("---------------------Video Info---------------------")
     logging.info("Video name: {}, translation model: {}, video link: {}".format(VIDEO_NAME, args.model_name, args.link))
 
@@ -313,7 +319,7 @@ def main():
             os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.ass" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
 
     end_time = time.time()
-    logging.info("Pipeline finished, time duration:{}".format(start_time - end_time))
+    logging.info("Pipeline finished, time duration:{}".format(time.strftime("%H:%M:%S", time.gmtime(end_time - start_time))))
 
 if __name__ == "__main__":
     main()