task initialization done

Former-commit-id: c84c9e1172e01e6a6cdb30d3d1f79f6f38e76a04

configs/task_config.yaml

@@ -6,5 +6,5 @@ output_type:
   video: False
   bilingal: False
 source_lang: EN
-target_lang: CN
+target_lang: ZH
 field: SC2

entries/run.py

@@ -18,7 +18,7 @@ def parse_args():
     parser.add_argument("--srt_file", help="srt file input path here", default=None, type=str, required=False)
     parser.add_argument("--continue", help="task_id that need to continue", default=None, type=str, required=False) # need implement
     parser.add_argument("--launch_cfg", help="launch config path", default='./configs/local_launch.yaml', type=str, required=False)
-    # task config
+    parser.add_argument("--task_cfg", help="task config path", default='./configs/task_config.yaml', type=str, required=False)
     args = parser.parse_args()
 
     return args
@@ -27,7 +27,7 @@ if __name__ == "__main__":
     # read args and configs
     args = parse_args()
     launch_cfg = load(open(args.launch_cfg), Loader=Loader)
-    # TODO: task config
+    task_cfg = load(open(args.task_cfg), Loader=Loader)
 
     # initialize dir
     local_dir = Path(launch_cfg['local_dump'])
@@ -41,7 +41,6 @@ if __name__ == "__main__":
     task_dir = local_dir.joinpath(f"task_{task_id}")
     task_dir.mkdir(parents=False, exist_ok=False)
     task_dir.joinpath("results").mkdir(parents=False, exist_ok=False)
-    task_dir.joinpath("logs").mkdir(parents=False, exist_ok=False)
 
     # logging
     logging.basicConfig(level=logging.INFO, handlers=[
@@ -49,12 +48,22 @@ if __name__ == "__main__":
             "{}/{}_{}.log".format(task_dir, f"task_{task_id}", datetime.now().strftime("%m%d%Y_%H%M%S")),
             'w', encoding='utf-8')])
 
-    # TODO: write task info into log
-
     # Task create
     if args.link is not None:
         try:
-            task = Task.fromYoutubeLink(args.link, task_id, launch_cfg)
+            task = Task.fromYoutubeLink(args.link, task_id, task_dir, task_cfg)
+        except:
+            shutil.rmtree(task_dir)
+            raise RuntimeError("failed to create task from youtube link")
+    elif args.video_file is not None:
+        try:
+            task = Task.fromVideoFile(args.video_file, task_id, task_dir, task_cfg)
+        except:
+            shutil.rmtree(task_dir)
+            raise RuntimeError("failed to create task from youtube link")
+    elif args.audio_file is not None:
+        try:
+            task = Task.fromVideoFile(args.audio_file, task_id, task_dir, task_cfg)
         except:
             shutil.rmtree(task_dir)
             raise RuntimeError("failed to create task from youtube link")

src/task.py

@@ -64,36 +64,47 @@ class Task:
         with self.__status_lock:
             self.__status = new_status
 
-    def __init__(self, task_id, task_local_dir, launch_info):
+    def __init__(self, task_id, task_local_dir, task_cfg):
         self.__status_lock = threading.Lock()
         self.__status = TaskStatus.CREATED
         openai.api_key = getenv("OPENAI_API_KEY")
-        self.launch_info = launch_info
+        self.launch_info = task_cfg # do not use, just for fallback
         self.task_local_dir = task_local_dir
-        self.model = launch_info["model"]
+        self.model = task_cfg["model"]
         self.gpu_status = 0
-        self.output_type = launch_info["output_type"]
+        self.output_type = task_cfg["output_type"]
+        self.target_lang = task_cfg["target_lang"]
+        self.source_lang = task_cfg["source_lang"]
+        self.field = task_cfg["field"]
         self.task_id = task_id
         self.audio_path = None
         self.SRT_Script = None
         self.result = None
+        print(f" Task ID: {self.task_id}")
+        logging.info(f" Task ID: {self.task_id}")
+        logging.info(f" {self.source_lang} -> {self.target_lang} task in {self.field}")
+        logging.info(f" Model: {self.model}")
+        logging.info(f" subtitle_type: {self.output_type['subtitle']}")
+        logging.info(f" video_ouput: \t{self.output_type['video']}")
+        logging.info(f" bilingal_ouput: \t{self.output_type['bilingal']}")
 
     @staticmethod
-    def fromYoutubeLink(youtube_url, task_id, launch_info):
-        # get audio path
+    def fromYoutubeLink(youtube_url, task_id, task_dir, task_cfg):
+        # convert to audio
         logging.info("Task Creation method: Youtube Link")
-        local_dump = Path(launch_info['local_dump'])  # should get from launch config
-        return YoutubeTask(task_id, local_dump.joinpath(f"task_{task_id}"), launch_info, youtube_url)
+        return YoutubeTask(task_id, task_dir, task_cfg, youtube_url)
 
     @staticmethod
-    def fromAudioFile():
+    def fromAudioFile(audio_path, task_id, task_dir, task_cfg):
         # get audio path
-        return Task(...)
+        logging.info("Task Creation method: Audio File")
+        return AudioTask(task_id, task_dir, task_cfg, audio_path)
     
     @staticmethod
-    def fromVideoFile():
+    def fromVideoFile(video_path, task_id, task_dir, task_cfg):
         # get audio path
-        return Task(...)
+        logging.info("Task Creation method: Video File")
+        return VideoTask(task_id, task_dir, task_cfg, video_path)
     
     # Module 1 ASR: audio --> SRT_script
     def get_srt_class(self, whisper_model='tiny', method="stable"):
@@ -103,7 +114,8 @@ class Task:
         time.sleep(5)
         pass
 
-    # Module 2: SRT preprocess: perform preprocess steps   
+    # Module 2: SRT preprocess: perform preprocess steps
+    # TODO: multi-lang and multi-field support according to task_cfg   
     def preprocess(self):
         self.status = TaskStatus.PRE_PROCESSING
         logging.info("--------------------Start Preprocessing SRT class--------------------")
@@ -113,7 +125,7 @@ class Task:
         processed_srt_path_en = str(Path(self.task_local_dir).with_suffix('')) + '_processed.srt'
         self.SRT_Script.write_srt_file_src(processed_srt_path_en)
 
-        if self.output_type == "ass":
+        if self.output_type["subtitle"] == "ass":
             logging.info("write English .srt file to .ass")
             assSub_en = srt2ass(processed_srt_path_en)
             logging.info('ASS subtitle saved as: ' + assSub_en)
@@ -140,7 +152,6 @@ class Task:
     def output_render(self):
         self.status = TaskStatus.OUTPUT_MODULE
         return "TODO"
-        pass
     
     def run_pipeline(self):
         self.get_srt_class()
@@ -150,8 +161,8 @@ class Task:
         self.result = self.output_render()
 
 class YoutubeTask(Task):
-    def __init__(self, task_id, task_local_dir, launch_info, youtube_url):
-        super().__init__(task_id, task_local_dir, launch_info)
+    def __init__(self, task_id, task_local_dir, task_cfg, youtube_url):
+        super().__init__(task_id, task_local_dir, task_cfg)
         self.youtube_url = youtube_url
 
     def run(self):
@@ -176,26 +187,37 @@ class YoutubeTask(Task):
                  '-ab', '192000', '-vn', self.task_local_dir.joinpath(f"task_{self.task_id}.mp3")])
             logging.info("audio extraction finished")
         
+        self.video_path = self.task_local_dir.joinpath(f"task_{self.task_id}.mp4")
         self.audio_path = self.task_local_dir.joinpath(f"task_{self.task_id}.mp3")
 
-        logging.info("Task Creation Complete.")
+        logging.info("Data Prep Complete. Start pipeline")
 
         super().run_pipeline()
 
 class AudioTask(Task):
-    def __init__(self, task_id, task_local_dir, launch_info, audio_path):
-        super().__init__(task_id, task_local_dir, launch_info)
+    def __init__(self, task_id, task_local_dir, task_cfg, audio_path):
+        super().__init__(task_id, task_local_dir, task_cfg)
+        # TODO: check audio format
         self.audio_path = audio_path
+        self.video_path = None
 
     def run(self):
+        logging.info("Data Prep Complete. Start pipeline")
         super().run_pipeline()
 
 class VideoTask(Task):
-    def __init__(self, task_id, task_local_dir, launch_info, video_path):
-        super().__init__(task_id, task_local_dir, launch_info)
-        self.audio_path = video_path
+    def __init__(self, task_id, task_local_dir, task_cfg, video_path):
+        super().__init__(task_id, task_local_dir, task_cfg)
+        # TODO: check video format {.mp4}
+        self.video_path = video_path
 
     def run(self):
-        
+        logging.info("using ffmpeg to extract audio")
+        subprocess.run(
+                ['ffmpeg', '-i', self.video_path, '-f', 'mp3',
+                 '-ab', '192000', '-vn', self.task_local_dir.joinpath(f"task_{self.task_id}.mp3")])
+        logging.info("audio extraction finished")
 
+        self.audio_path = self.task_local_dir.joinpath(f"task_{self.task_id}.mp3")
+        logging.info("Data Prep Complete. Start pipeline")
         super().run_pipeline()