Add class Task skeleton

Co-authored-by: Eason Lu <[email protected]>
Co-authored-by: Ivanfangsc <[email protected]>

Former-commit-id: b9dbcb50abb28983b81a5eb04dfc460d234d84cf

src/Pigeon.py

@@ -317,7 +317,7 @@ class Pigeon(object):
         logging.info("--------------------Start Preprocessing SRT class--------------------")
         self.srt.write_srt_file_src(self.srt_path)
         self.srt.form_whole_sentence()
-        self.srt.spell_check_term()
+        # self.srt.spell_check_term()
         self.srt.correct_with_force_term()
         processed_srt_file_en = str(Path(self.srt_path).with_suffix('')) + '_processed.srt'
         self.srt.write_srt_file_src(processed_srt_file_en)

src/preprocess/audio_extract.py

@@ -0,0 +1,13 @@
+import pathlib
+import os
+import subprocess
+
+
+def extract_audio(local_video_path: str, save_dir_path: str = "./downloads/audio") -> str:
+    if os.name == 'nt':
+        NotImplementedError("Filename extraction on Windows not yet implemented")
+
+    out_file_name = os.path.basename(local_video_path)
+    audio_path_out = save_dir_path.join("/").join(out_file_name)
+    subprocess.run(['ffmpeg', '-i', local_video_path, '-f', 'mp3', '-ab', '192000', '-vn', audio_path_out])
+    return audio_path_out

src/preprocess/video_download.py

@@ -0,0 +1,20 @@
+from pytube import YouTube
+import logging
+
+def download_youtube_to_local_file(youtube_url: str, local_dir_path: str = "./downloads") -> str:
+    yt = YouTube(youtube_url)
+    try:
+        audio = yt.streams.filter(only_audio=True, file_extension='mp4').order_by('abr').desc().first()
+        # video = yt.streams.filter(file_extension='mp4').order_by('resolution').asc().first()
+        if audio:
+            saved_audio = audio.download(output_path=local_dir_path.join("/audio"))
+            logging.info(f"Audio download successful: {saved_audio}")
+            return saved_audio
+        else:
+            logging.error(f"Audio stream not found in {youtube_url}")
+            raise f"Audio stream not found in {youtube_url}"
+    except Exception as e:
+        # print("Connection Error: ", end='')
+        print(e)
+        raise e
+

src/task.py

@@ -0,0 +1,111 @@
+from src.srt_util.srt import SrtScript
+from src.srt_util.srt2ass import srt2ass
+import openai
+import stable_whisper
+import torch
+import whisper
+from pytube import YouTube
+from os import getenv
+from enum import Enum
+from pathlib import Path
+from enum import Enum, auto
+
+
+"""
+Youtube link
+    - link
+    - model
+    - output type
+
+Video file
+    - path
+    - model
+    - output type
+
+Audio file
+    - path
+    - model
+    - output type
+
+""" 
+"""
+TaskID
+Progress: Enum
+Computing resrouce status 
+SRT_Script : SrtScript
+    -  input module -> initialize (ASR module)
+    -  Pre-process
+    -  Translation  (%)
+    -  Post process (time stamp)
+    -  Output module: SRT_Script --> output(.srt)
+    -  (Optional) mp4
+"""
+
+class TaskStatus(Enum):
+    INITIALIZING_ASR = (auto(), None)
+    PRE_PROCESSING = (auto(), None)
+    TRANSLATING = (auto(), 0.0)
+    POST_PROCESSING = (auto(), None)
+    OUTPUT_MODULE = (auto(), None)
+
+
+class Task:
+    def __init__(self, task_id, audio_path, model, output_type):
+        # openai.api_key = getenv("OPENAI_API_KEY")
+        self.audio_path = audio_path
+        self.model = model
+        self.gpu_status = 0
+        self.output_type = output_type
+        self.task_id = task_id
+        self.progress = NotImplemented
+        self.SRT_Script = None
+        self.local_dump = Path()
+
+    @staticmethod
+    def fromYoutubeLink(youtube_url):
+        # convert to audio
+        
+        return Task(...)
+
+    @staticmethod
+    def fromAudioFile():
+        # 
+        return Task(...)
+    
+    @staticmethod
+    def fromVideoFile():
+        # convert to audio
+        return Task(...)
+    
+    # Module 1 ASR: audio --> SRT_script
+    def get_srt_class(self, whisper_model='tiny', method="stable"):
+        # Instead of using the script_en variable directly, we'll use script_input
+        pass
+
+    # Module 2: SRT preprocess: perform preprocess steps   
+    def preprocess(self):
+        pass
+    
+    def update_translation_progress(self, new_progress):
+        if self.progress == TaskStatus.TRANSLATING:
+            self.progress = TaskStatus.TRANSLATING.value[0], new_progress
+
+    # Module 3: perform srt translation
+    def translation(self):
+        pass
+    
+    # Module 4: perform srt post process steps
+    def postprocess(self):
+        pass
+
+    # Module 5: output module
+    def output_render(self):
+        pass
+    
+    def run_pipeline(self):
+        self.get_srt_class()
+        self.preprocess()
+        self.translation()
+        self.postprocess()
+        out = self.output_render()
+        return out