Merge branch

Former-commit-id: daf0684aff6d0f12ea49dbf402ef90d9a3be46d5

.gitignore

@@ -1,3 +1,4 @@
 /downloads
 /results
-.DS_Store
+.DS_Store
+/__pycache__

README.md

@@ -6,13 +6,18 @@
 pip install -r requirement.txt
 ```
 
-## Usage
-```
-usage: pipeline.py [-h] [--link LINK] [--local_path LOCAL_PATH] [--download DOWNLOAD] [--result RESULT] [--video_name VIDEO_NAME] [--model_name]
+## Quick Start:
+
+example online: 
 
-quick start:
+```
+python3 pipeline.py --link "https://www.youtube.com/watch?v=61c4dn6851g" 
+```
 
-example online: python3 pipeline.py --link https://www.youtube.com/watch?v=61c4dn6851g --download ./downloads --result ./results --video_name SO_I_CHOSE_RANDOM
+example offline: 
+```
+python3 pipeline.py --audio_file test_translation.m4a --result ./results --video_name test_translation
+```
 
 python3 pipeline.py --link https://www.youtube.com/watch?v=VrigMmXt9A0 --video_name Ukraine_and_its_Global_Impact
 
@@ -20,21 +25,30 @@ python3 pipeline.py --video_file '/home/jiaenliu/project-t/downloads/audio/Ukrai
 
 example offline: python3 pipeline.py --local_path test_translation.m4a --result ./results --video_name test_translation
 
-example text input: python pipeline.py --text_file "/home/jiaenliu/project-t/results/huanghe_translation_en.txt" --result "/home/jiaenliu/project-t/results" --video_name "huanghe_test"
+## Usage
+```
+usage: pipeline.py [-h] [--link LINK] [--video_file VIDEO_FILE] [--audio_file AUDIO_FILE] [--srt_file SRT_FILE] [--download DOWNLOAD]
+                   [--output_dir OUTPUT_DIR] [--video_name VIDEO_NAME] [--model_name MODEL_NAME] [-only_srt] [-v]
 
 options:
   -h, --help            show this help message and exit
   --link LINK           youtube video link here
-  --local_path LOCAL_PATH
+  --video_file VIDEO_FILE
                         local video path here
+  --audio_file AUDIO_FILE
+                        local audio path here
+  --srt_file SRT_FILE   srt file input path here
   --download DOWNLOAD   download path
-  --result RESULT       translate result path
+  --output_dir OUTPUT_DIR
+                        translate result path
   --video_name VIDEO_NAME
-                        video name
-  --model_name MODEL_NAME 
-                        model name
+                        video name, if use video link as input, the name will auto-filled by youtube video name
+  --model_name MODEL_NAME
+                        model name only support text-davinci-003 and gpt-3.5-turbo
+  -only_srt             set script output to only .srt file
+  -v                    auto encode script with video
+```
 
+## Notice
 if you cannot download youtube video, please follow the link below.
-https://github.com/pytube/pytube/issues/1498
-
-```
+https://github.com/pytube/pytube/issues/1498

pipeline.py

@@ -2,9 +2,8 @@ import openai
 from pytube import YouTube
 import argparse
 import os
-import io
 import whisper
-import ffmpeg
+from tqdm import tqdm
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
@@ -65,14 +64,12 @@ if args.link is not None and args.video_file is None:
         exit()
     
     video_path = f'{DOWNLOAD_PATH}/video/{video.default_filename}'
-    # video_file = open(video_path, "rb")
     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_file = open(args.video_file, "rb")
     video_path = args.video_file
     if args.audio_file is not None:
         audio_file= open(args.audio_file, "rb")
@@ -134,7 +131,7 @@ if script.strip():
     script_arr.append(script.strip())
 
 # Translate and save
-for s in script_arr:
+for s in tqdm(script_arr):
     # using chatgpt model
     if model_name == "gpt-3.5-turbo":
         # print(s + "\n")