Merge pull request #68 from project-kxkg/gradio-interface

debug& add gradio web interface

Former-commit-id: cb64790fdb61a3e73ca6abd93ca50db0cbee6f7f

configs/task_config.yaml

@@ -29,7 +29,7 @@ post_process:
 # output type that user receive
 output_type: 
   subtitle: srt
-  video: False
-  bilingal: True
+  video: True
+  bilingual: True
 
 

entries/app.py

@@ -0,0 +1,90 @@
+import __init_lib_path
+import gradio as gr
+from src.task import Task
+import logging
+from yaml import Loader, Dumper, load, dump
+import os
+from pathlib import Path
+from datetime import datetime
+import shutil
+from uuid import uuid4
+
+launch_config = "./configs/local_launch.yaml"
+task_config = './configs/task_config.yaml'
+
+def init(output_type, src_lang, tgt_lang, domain):
+    launch_cfg = load(open(launch_config), Loader=Loader)
+    task_cfg = load(open(task_config), Loader=Loader)
+
+    # overwrite config file
+    task_cfg["source_lang"] = src_lang
+    task_cfg["target_lang"] = tgt_lang
+    task_cfg["field"] = domain
+
+    if "Video File" in output_type:
+        task_cfg["output_type"]["video"] = True
+    else:
+        task_cfg["output_type"]["video"] = False
+    
+    if "Bilingual" in output_type:
+        task_cfg["output_type"]["bilingual"] = True
+    else:
+        task_cfg["output_type"]["bilingual"] = False
+    
+    if ".ass output" in output_type:
+        task_cfg["output_type"]["subtitle"] = "ass"
+    else:
+        task_cfg["output_type"]["subtitle"] = "srt"
+
+    # initialize dir
+    local_dir = Path(launch_cfg['local_dump'])
+    if not local_dir.exists():
+        local_dir.mkdir(parents=False, exist_ok=False)
+
+    # get task id
+    task_id = str(uuid4())
+
+    # create locak dir for the task
+    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)
+
+    # logging setting
+    logfmt = "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
+    logging.basicConfig(level=logging.INFO, format=logfmt, handlers=[
+        logging.FileHandler(
+            "{}/{}_{}.log".format(task_dir, f"task_{task_id}", datetime.now().strftime("%m%d%Y_%H%M%S")),
+            'w', encoding='utf-8')])
+    return task_id, task_dir, task_cfg
+
+def process_input(video_file, youtube_link, src_lang, tgt_lang, domain, output_type):
+    task_id, task_dir, task_cfg = init(output_type, src_lang, tgt_lang, domain)
+    if youtube_link:
+        task = Task.fromYoutubeLink(youtube_link, task_id, task_dir, task_cfg)
+        task.run()
+        return task.result
+    elif video_file is not None:
+        task = Task.fromVideoFile(video_file, task_id, task_dir, task_cfg)
+        task.run()
+        return task.result
+    else:
+        return None
+
+demo = gr.Interface(fn=process_input,
+    inputs=[
+        gr.components.Video(label="Upload a video"),
+        gr.components.Textbox(label="Or enter a YouTube URL"), 
+        gr.components.Dropdown(choices=["EN", "ZH"], label="Select Source Language"),
+        gr.components.Dropdown(choices=["ZH", "EN"], label="Select Target Language"),
+        gr.components.Dropdown(choices=["General", "SC2"], label="Select Domain"),
+        gr.CheckboxGroup(["Video File", "Bilingual", ".ass output"], label="Output Settings", info="What do you want?"),
+    ],
+    outputs=[
+        gr.components.Video(label="Processed Video")
+    ],
+    title="ViDove: video translation toolkit demo",
+    description="Upload a video or enter a YouTube URL."
+    )
+
+if __name__ == "__main__":
+    demo.launch()

src/task.py

@@ -3,7 +3,7 @@ import time
 
 import openai
 from pytube import YouTube
-from os import getenv
+from os import getenv, getcwd
 from pathlib import Path
 from enum import Enum, auto
 import logging
@@ -15,6 +15,7 @@ from src.translators.translation import get_translation, prompt_selector
 
 import torch
 import stable_whisper
+import shutil
 
 """
 Youtube link
@@ -97,7 +98,7 @@ class Task:
         logging.info(f"Translation Model: {self.translation_model}")
         logging.info(f"subtitle_type: {self.output_type['subtitle']}")
         logging.info(f"video_ouput: {self.output_type['video']}")
-        logging.info(f"bilingal_ouput: {self.output_type['bilingal']}")
+        logging.info(f"bilingual_ouput: {self.output_type['bilingual']}")
         logging.info("Pre-process setting:")
         for key in self.pre_setting:
             logging.info(f"{key}: {self.pre_setting[key]}")
@@ -176,7 +177,7 @@ class Task:
 
         if self.output_type["subtitle"] == "ass":
             logging.info("write English .srt file to .ass")
-            assSub_src = srt2ass(processed_srt_path_src)
+            assSub_src = srt2ass(processed_srt_path_src, "default", "No", "Modest")
             logging.info('ASS subtitle saved as: ' + assSub_src)
         self.script_input = self.SRT_Script.get_source_only()
         pass
@@ -207,13 +208,13 @@ class Task:
         self.status = TaskStatus.OUTPUT_MODULE
         video_out = self.output_type["video"]
         subtitle_type = self.output_type["subtitle"]
-        is_bilingal = self.output_type["bilingal"]
+        is_bilingual = self.output_type["bilingual"]
 
-        results_dir = Path(self.task_local_dir)/ "results"
+        results_dir =f"{self.task_local_dir}/results"
 
         subtitle_path = f"{results_dir}/{self.task_id}_{self.target_lang}.srt"
         self.SRT_Script.write_srt_file_translate(subtitle_path)
-        if is_bilingal:
+        if is_bilingual:
             subtitle_path = f"{results_dir}/{self.task_id}_{self.source_lang}_{self.target_lang}.srt"
             self.SRT_Script.write_srt_file_bilingual(subtitle_path)
 
@@ -227,8 +228,12 @@ class Task:
         # encode to .mp4 video file
         if video_out and self.video_path is not None:
             logging.info("encoding video file")
+            logging.info(f'ffmpeg comand: \nffmpeg -i {self.video_path} -vf "subtitles={subtitle_path}" {results_dir}/{self.task_id}.mp4')
             subprocess.run(
-                f'ffmpeg -i {self.video_path} -vf "subtitles={subtitle_path}" {results_dir}/{self.task_id}.mp4')
+                ["ffmpeg",
+                    "-i", self.video_path,
+                    "-vf", f"subtitles={subtitle_path}",
+                    f"{results_dir}/{self.task_id}.mp4"])
             final_res = f"{results_dir}/{self.task_id}.mp4"
 
         self.t_e = time()
@@ -242,7 +247,7 @@ class Task:
         self.translation()
         self.postprocess()
         self.result = self.output_render()
-        print(self.result)
+        # print(self.result)
 
 class YoutubeTask(Task):
     def __init__(self, task_id, task_local_dir, task_cfg, youtube_url):
@@ -286,8 +291,8 @@ class AudioTask(Task):
         self.video_path = None
 
     def run(self):
-        logging.info(f" Video File Dir: {self.video_path}")
-        logging.info(f" Audio File Dir: {self.audio_path}")
+        logging.info(f"Video File Dir: {self.video_path}")
+        logging.info(f"Audio File Dir: {self.audio_path}")
         logging.info("Data Prep Complete. Start pipeline")
         super().run_pipeline()
 
@@ -295,7 +300,11 @@ class VideoTask(Task):
     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
+        new_video_path = f"{task_local_dir}/task_{self.task_id}.mp4"
+        print(new_video_path)
+        logging.info(f"Copy video file to: {new_video_path}")
+        shutil.copyfile(video_path, new_video_path)
+        self.video_path = new_video_path
 
     def run(self):
         logging.info("using ffmpeg to extract audio")