Merge pull request #23 from project-kxkg/eason/refactor

Eason/refactor to test the log function

Former-commit-id: c3a243de07855712925578c1aa4749153ef05bdd

SRT.py

@@ -420,7 +420,7 @@ class SRT_script():
             n = 0
         return real_word, len(word)+n
     
-    def realtime_write_srt(self,path,range,length,idx):
+    def realtime_write_srt(self,path,range,length, idx):
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
@@ -428,19 +428,19 @@ class SRT_script():
             #     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
+                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
 
-    def realtime_bilingual_write_srt(self,path,range,length,idx):
+    def realtime_bilingual_write_srt(self,path,range, length,idx):
         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):
-                if i<range[0]-1: continue
-                if i>=range[1] + length:break
+                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

@@ -5,6 +5,7 @@ import os
 from tqdm import tqdm
 from SRT import SRT_script
 import stable_whisper
+import subprocess
 
 import time
 
@@ -22,7 +23,7 @@ parser.add_argument("-v", help="auto encode script with video", action='store_tr
 args = parser.parse_args()
 
 # input should be either video file or youtube video link.
-if args.link is None and args.video_file is None and args.srt_file is None:
+if args.link is None and args.video_file is None and args.srt_file is None and args.audio_file is None:
     print("need video source or srt file")
     exit()
 
@@ -49,6 +50,8 @@ if args.video_name == 'placeholder' :
         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
 
@@ -87,13 +90,22 @@ if args.link is not None and args.video_file is None:
 elif args.video_file is not None:
     # Read from local
     video_path = args.video_file
+    # audio_path = "{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3".format(DOWNLOAD_PATH, VIDEO_NAME)
     if args.audio_file is not None:
         audio_file= open(args.audio_file, "rb")
         audio_path = args.audio_file
     else:
-        os.system(f'ffmpeg -i {args.video_file} -f mp3 -ab 192000 -vn {DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3')
-        audio_file= open(f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3', "rb")
-        audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
+        # escaped_video_path = args.video_file.replace('(', '\(').replace(')', '\)').replace(' ', '\ ')
+        # print(escaped_video_path)
+        # os.system(f'ffmpeg -i {escaped_video_path} -f mp3 -ab 192000 -vn {DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3')
+        # audio_file= open(f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3', "rb")
+        # audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
+        output_audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
+        # print(video_path)
+        # print(output_audio_path)
+        subprocess.run(['ffmpeg', '-i', video_path, '-f', 'mp3', '-ab', '192000', '-vn', output_audio_path])
+        audio_file = open(output_audio_path, "rb")
+        audio_path = output_audio_path
 
 if not os.path.exists(f'{RESULT_PATH}/{VIDEO_NAME}'):
     os.mkdir(f'{RESULT_PATH}/{VIDEO_NAME}')
@@ -247,7 +259,11 @@ def get_response(model_name, sentence):
         
         
 # Translate and save
+previous_length = 0
 for sentence, range in tqdm(zip(script_arr, range_arr)):
+    # update the range based on previous length
+    range = (range[0]+previous_length, range[1]+previous_length)
+
     # using chatgpt model
     print(f"now translating sentences {range}")
     flag = True
@@ -261,14 +277,12 @@ for sentence, range in tqdm(zip(script_arr, range_arr)):
             time.sleep(30)
             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,args.link)
+    srt.set_translation(translate, range, model_name)
 
-    add_length = srt.check_len_and_split_range(range, threshold)
-    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: