Merge pull request #31 from project-kxkg/eason/code_refine

Eason/code refine

Former-commit-id: 47a9c9e7b77eb9434c76654683dc6c9e5788cdb3

pipeline.py

@@ -6,155 +6,120 @@ from tqdm import tqdm
 from SRT import SRT_script
 import stable_whisper
 import whisper
+from srt2ass import srt2ass
 
 import subprocess
 
 import time
 
-parser = argparse.ArgumentParser()
-parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
-parser.add_argument("--video_file", help="local video path here", default=None, type=str, required=False)
-parser.add_argument("--audio_file", help="local audio path here", default=None, type=str, required=False)
-parser.add_argument("--srt_file", help="srt file input path here", default=None, type=str, required=False)  # New argument
-parser.add_argument("--download", help="download path", default='./downloads', type=str, required=False)
-parser.add_argument("--output_dir", help="translate result path", default='./results', type=str, required=False)
-parser.add_argument("--video_name", help="video name, if use video link as input, the name will auto-filled by youtube video name", default='placeholder', type=str, required=False)
-parser.add_argument("--model_name", help="model name only support gpt-4 and gpt-3.5-turbo", type=str, required=False, default="gpt-4") # default change to gpt-4
-parser.add_argument("-only_srt", help="set script output to only .srt file", action='store_true')
-parser.add_argument("-v", help="auto encode script with video", action='store_true')
-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 and args.audio_file is None:
-    print("need video source or srt file")
-    exit()
-
-# set up
-openai.api_key = os.getenv("OPENAI_API_KEY")
-DOWNLOAD_PATH = args.download
-if not os.path.exists(DOWNLOAD_PATH):
-    os.mkdir(DOWNLOAD_PATH)
-    os.mkdir(f'{DOWNLOAD_PATH}/audio')
-    os.mkdir(f'{DOWNLOAD_PATH}/video')
-
-RESULT_PATH = args.output_dir
-if not os.path.exists(RESULT_PATH):
-    os.mkdir(RESULT_PATH)
-
-# set video name as the input file name if not specified
-if args.video_name == 'placeholder' :
-    # set video name to upload file name
-    if args.video_file is not None:
-        VIDEO_NAME = args.video_file.split('/')[-1].split('.')[0]
-    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("_")[0]
-    else:
-        VIDEO_NAME = args.video_name
-else:
-    VIDEO_NAME = args.video_name
-
-model_name = args.model_name
-
-threshold = 30
-
-# get source audio
-if args.link is not None and args.video_file is None:
-    # Download audio from YouTube
-    video_link = args.link
-    video = None
-    audio = None
-    try:
-        yt = YouTube(video_link)
-        video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
-        if video:
-            video.download(f'{DOWNLOAD_PATH}/video')
-            print('Video download completed!')
-        else:
-            print("Error: Video stream not found")
-        audio = yt.streams.filter(only_audio=True, file_extension='mp4').first()
-        if audio:
-            audio.download(f'{DOWNLOAD_PATH}/audio')
-            print('Audio download completed!')
-        else:
-            print("Error: Audio stream not found")
-    except Exception as e:
-        print("Connection Error")
-        print(e) 
-        exit()
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
+    parser.add_argument("--video_file", help="local video path here", default=None, type=str, required=False)
+    parser.add_argument("--audio_file", help="local audio path here", default=None, type=str, required=False)
+    parser.add_argument("--srt_file", help="srt file input path here", default=None, type=str, required=False)  # New argument
+    parser.add_argument("--download", help="download path", default='./downloads', type=str, required=False)
+    parser.add_argument("--output_dir", help="translate result path", default='./results', type=str, required=False)
+    parser.add_argument("--video_name", help="video name, if use video link as input, the name will auto-filled by youtube video name", default='placeholder', type=str, required=False)
+    parser.add_argument("--model_name", help="model name only support gpt-4 and gpt-3.5-turbo", type=str, required=False, default="gpt-4") # default change to gpt-4
+    parser.add_argument("-only_srt", help="set script output to only .srt file", action='store_true')
+    parser.add_argument("-v", help="auto encode script with video", action='store_true')
+    args = parser.parse_args()
     
-    video_path = f'{DOWNLOAD_PATH}/video/{video.default_filename}'
-    audio_path = '{}/audio/{}'.format(DOWNLOAD_PATH, audio.default_filename)
-    audio_file = open(audio_path, "rb")
-    if VIDEO_NAME == 'placeholder':
-        VIDEO_NAME = audio.default_filename.split('.')[0]
-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)
+    return args
+
+def get_sources(args, download_path, result_path, video_name):
+    # get source audio
+    if args.link is not None and args.video_file is None:
+        # Download audio from YouTube
+        video_link = args.link
+        video = None
+        audio = None
+        try:
+            yt = YouTube(video_link)
+            video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
+            if video:
+                video.download(f'{download_path}/video')
+                print('Video download completed!')
+            else:
+                print("Error: Video stream not found")
+            audio = yt.streams.filter(only_audio=True, file_extension='mp4').first()
+            if audio:
+                audio.download(f'{download_path}/audio')
+                print('Audio download completed!')
+            else:
+                print("Error: Audio stream not found")
+        except Exception as e:
+            print("Connection Error")
+            print(e) 
+            exit()
+        
+        video_path = f'{download_path}/video/{video.default_filename}'
+        audio_path = '{}/audio/{}'.format(download_path, audio.default_filename)
+        audio_file = open(audio_path, "rb")
+        if video_name == 'placeholder':
+            video_name = audio.default_filename.split('.')[0]
+    elif args.video_file is not None:
+        # Read from local
+        video_path = args.video_file
+
+        if args.audio_file is not None:
+            audio_file= open(args.audio_file, "rb")
+            audio_path = args.audio_file
+        else:
+            output_audio_path = f'{download_path}/audio/{video_name}.mp3'
+            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}')
+
     if args.audio_file is not None:
         audio_file= open(args.audio_file, "rb")
         audio_path = args.audio_file
-    else:
-        output_audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
-        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}')
-
-if args.audio_file is not None:
-    audio_file= open(args.audio_file, "rb")
-    audio_path = args.audio_file
-
-# Instead of using the script_en variable directly, we'll use script_input
-srt_file_en = args.srt_file 
-
-if srt_file_en is not None: 
-    srt = SRT_script.parse_from_srt_file(srt_file_en)
-else:
-    # using whisper to perform speech-to-text and save it in <video name>_en.txt under RESULT PATH.
-    srt_file_en = "{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
-    if not os.path.exists(srt_file_en):
-
-        # use OpenAI API for transcribe
-        # transcript = openai.Audio.transcribe("whisper-1", audio_file) 
-
-        # use local whisper model 
-        # model = whisper.load_model("base") # using base model in local machine (may use large model on our server)
-        # transcript = model.transcribe(audio_path)
-
-        # use stable-whisper
-        model = stable_whisper.load_model('base')
-        transcript = model.transcribe(audio_path, regroup = False, initial_prompt="Hello, welcome to my lecture. Are you good my friend?")
-        (
-            transcript
-            .split_by_punctuation(['.', '。', '?'])
-            .merge_by_gap(.15, max_words=3)
-            .merge_by_punctuation([' '])
-            .split_by_punctuation(['.', '。', '?'])
-        )
+        pass
 
-        transcript = transcript.to_dict()
-        srt = SRT_script(transcript['segments']) # read segments to SRT class
+    return audio_path, audio_file, video_path, video_name
 
-    else:
+def get_srt_class(srt_file_en, result_path, video_name, audio_path, audio_file = None, whisper_model = 'base', method = "stable"):
+    # Instead of using the script_en variable directly, we'll use script_input
+    if srt_file_en is not None: 
         srt = SRT_script.parse_from_srt_file(srt_file_en)
+    else:
+        # using whisper to perform speech-to-text and save it in <video name>_en.txt under RESULT PATH.
+        srt_file_en = "{}/{}/{}_en.srt".format(result_path, video_name, video_name)
+        if not os.path.exists(srt_file_en):
+
+            # use OpenAI API for transcribe
+            if method == "api":
+                transcript = openai.Audio.transcribe("whisper-1", audio_file) 
+
+            # use local whisper model 
+            elif method == "basic":
+                model = whisper.load_model(whisper_model) # using base model in local machine (may use large model on our server)
+                transcript = model.transcribe(audio_path)
+
+            # use stable-whisper
+            elif method == "stable":
+                model = stable_whisper.load_model(whisper_model)
+                transcript = model.transcribe(audio_path, regroup = False, initial_prompt="Hello, welcome to my lecture. Are you good my friend?")
+                (
+                    transcript
+                    .split_by_punctuation(['.', '。', '?'])
+                    .merge_by_gap(.15, max_words=3)
+                    .merge_by_punctuation([' '])
+                    .split_by_punctuation(['.', '。', '?'])
+                )
+                transcript = transcript.to_dict()
+            else:
+                raise ValueError("invalid speech to text method")
+
+            srt = SRT_script(transcript['segments']) # read segments to SRT class
 
-# srt class 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()
-
-if not args.only_srt:
-    from srt2ass import srt2ass
-    assSub_en = srt2ass(srt_file_en, "default", "No", "Modest")
-    print('ASS subtitle saved as: ' + assSub_en)
-
+        else:
+            srt = SRT_script.parse_from_srt_file(srt_file_en)
+    return srt_file_en, srt
 
 # Split the video script by sentences and create chunks within the token limit
 def script_split(script_in, chunk_size = 1000):
@@ -181,43 +146,38 @@ def script_split(script_in, chunk_size = 1000):
     assert len(script_arr) == len(range_arr)
     return script_arr, range_arr
 
-script_arr, range_arr = script_split(script_input)
-# print(script_arr, range_arr)
-
 # check whether previous translation is done
-zh_file = "{}/{}/{}_zh.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
-segidx = 1
-if os.path.exists(zh_file):
-    temp_file = "{}/{}/temp.srt".format(RESULT_PATH, VIDEO_NAME)
-    if os.path.exists(temp_file):
-        os.remove(temp_file)
-    with open(zh_file, "r") as f0:
-        for count, _ in enumerate(f0):
-            pass
-        count += 1
-        segidx = int(count/4)+1
-    en_file = "{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
-    if args.srt_file is not None:
-        en_file = args.srt_file
-    with open(en_file, "r") as f1, open(temp_file, "a") as f2:
-        x = f1.readlines()
-        #print(len(x))
-        if count >= len(x):
-            print('Work already done! Please delete {}_zh.srt files in result directory first in order to rework'.format(VIDEO_NAME))
-            exit()
-        for i, line in enumerate(x):
-            if i >= count:
-                #print(i)
-                f2.write(line)
+# zh_file = "{}/{}/{}_zh.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
+# segidx = 1
+# if os.path.exists(zh_file):
+#     temp_file = "{}/{}/temp.srt".format(RESULT_PATH, VIDEO_NAME)
+#     if os.path.exists(temp_file):
+#         os.remove(temp_file)
+#     with open(zh_file, "r") as f0:
+#         for count, _ in enumerate(f0):
+#             pass
+#         count += 1
+#         segidx = int(count/4)+1
+#     en_file = "{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
+#     if args.srt_file is not None:
+#         en_file = args.srt_file
+#     with open(en_file, "r") as f1, open(temp_file, "a") as f2:
+#         x = f1.readlines()
+#         #print(len(x))
+#         if count >= len(x):
+#             print('Work already done! Please delete {}_zh.srt files in result directory first in order to rework'.format(VIDEO_NAME))
+#             exit()
+#         for i, line in enumerate(x):
+#             if i >= count:
+#                 f2.write(line)
                 
-    srt = SRT_script.parse_from_srt_file(temp_file)
-    print('temp_contents')
-    print(srt.get_source_only())  
+#     srt = SRT_script.parse_from_srt_file(temp_file)
+#     print('temp_contents')
+#     print(srt.get_source_only())  
 
 
 def get_response(model_name, sentence):
     if model_name == "gpt-3.5-turbo" or model_name == "gpt-4":
-        # print(s + "\n")
         response = openai.ChatCompletion.create(
             model=model_name,
             messages = [
@@ -234,39 +194,99 @@ 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
-    while flag:
-        flag = False
-        try:
-            translate = get_response(model_name, sentence)
-        except Exception as e:
-            print("An error has occurred during translation:",e)
-            print("Retrying... the script will continue after 30 seconds.")
-            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, VIDEO_NAME, args.link)
-
-    # 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")
-
-if not args.only_srt:
-    assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")
-    print('ASS subtitle saved as: ' + assSub_zh)
-
-if args.v:
-    if args.only_srt:
-        os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
+def translate(srt, script_arr, range_arr, model_name, video_name, video_link):
+    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
+        while flag:
+            flag = False
+            try:
+                translate = get_response(model_name, sentence)
+            except Exception as e:
+                print("An error has occurred during translation:",e)
+                print("Retrying... the script will continue after 30 seconds.")
+                time.sleep(30)
+                flag = True
+        srt.set_translation(translate, range, model_name, video_name, video_link)
+
+
+def main():
+    args = parse_args()
+
+    # input check: 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 and args.audio_file is None:
+        print("need video source or srt file")
+        exit()
+
+    # set up
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    DOWNLOAD_PATH = args.download
+    if not os.path.exists(DOWNLOAD_PATH):
+        os.mkdir(DOWNLOAD_PATH)
+        os.mkdir(f'{DOWNLOAD_PATH}/audio')
+        os.mkdir(f'{DOWNLOAD_PATH}/video')
+
+    RESULT_PATH = args.output_dir
+    if not os.path.exists(RESULT_PATH):
+        os.mkdir(RESULT_PATH)
+
+    # set video name as the input file name if not specified
+    if args.video_name == 'placeholder' :
+        # set video name to upload file name
+        if args.video_file is not None:
+            VIDEO_NAME = args.video_file.split('/')[-1].split('.')[0]
+        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("_")[0]
+        else:
+            VIDEO_NAME = args.video_name
     else:
-        os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.ass" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
+        VIDEO_NAME = args.video_name
+
+    audio_path, audio_file, video_path, VIDEO_NAME = get_sources(args, DOWNLOAD_PATH, RESULT_PATH, VIDEO_NAME)
+
+    srt_file_en, srt = get_srt_class(args.srt_file, RESULT_PATH, VIDEO_NAME, audio_path, audio_file)
+
+    # SRT class 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()
+
+    # write ass
+    if not args.only_srt:
+        assSub_en = srt2ass(srt_file_en, "default", "No", "Modest")
+        print('ASS subtitle saved as: ' + assSub_en)
+
+    script_arr, range_arr = script_split(script_input)
+
+    translate(srt, script_arr, range_arr, args.model_name, VIDEO_NAME, args.link)
+
+    # SRT post-processing
+    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")
+
+    # write ass
+    if not args.only_srt:
+        assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")
+        print('ASS subtitle saved as: ' + assSub_zh)
+
+    # encode to .mp4 video file
+    if args.v:
+        if args.only_srt:
+            os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
+        else:
+            os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.ass" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
+
+
+if __name__ == "__main__":
+    main()