Create app.py

app.py

@@ -0,0 +1,101 @@
+import gradio as gr
+import os
+import ffmpeg
+import pysrt
+import pandas as pd
+import requests
+import io
+from transformers import MarianMTModel, MarianTokenizer
+
+def fetch_languages(url):
+    response = requests.get(url)
+    if response.status_code == 200:
+        csv_content = response.content.decode('utf-8')
+        df = pd.read_csv(io.StringIO(csv_content), delimiter="|", skiprows=2, header=None).dropna(axis=1, how='all')
+        df.columns = ['ISO 639-1', 'ISO 639-2', 'Language Name', 'Native Name']
+        df['ISO 639-1'] = df['ISO 639-1'].str.strip()
+        language_options = [(row['ISO 639-1'], f"{row['ISO 639-1']} - {row['Language Name']}") for index, row in df.iterrows()]
+        return language_options
+    else:
+        return []
+
+def text_to_srt(text):
+    lines = text.split('\n')
+    srt_content = ""
+    for i, line in enumerate(lines):
+        if line.strip() == "":
+            continue
+        try:
+            times, content = line.split(']', 1)
+            start, end = times[1:].split(' -> ')
+            if start.count(":") == 1:
+                start = "00:" + start
+            if end.count(":") == 1:
+                end = "00:" + end
+            srt_content += f"{i+1}\n{start.replace('.', ',')} --> {end.replace('.', ',')}\n{content.strip()}\n\n"
+        except ValueError:
+            continue
+    temp_file_path = '/tmp/output.srt'
+    with open(temp_file_path, 'w', encoding='utf-8') as file:
+        file.write(srt_content)
+    return temp_file_path
+
+def translate_text(text, source_language_code, target_language_code):
+    model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}"
+    try:
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+    except Exception as e:
+        return f"Failed to load model for {source_language_code} to {target_language_code}: {str(e)}"
+    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
+    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+    return translated_text
+
+def translate_srt(input_file, source_language_code, target_language_code):
+    subs = pysrt.open(input_file)
+    for sub in subs:
+        sub.text = translate_text(sub.text, source_language_code, target_language_code)
+    translated_srt_path = input_file.replace(".srt", f"_{target_language_code}.srt")
+    subs.save(translated_srt_path)
+    return translated_srt_path
+
+def add_subtitle_to_video(input_video, subtitle_file, soft_subtitle=True):
+    video_input_stream = ffmpeg.input(input_video)
+    subtitle_input_stream = ffmpeg.input(subtitle_file)
+    input_video_name = os.path.splitext(os.path.basename(input_video))[0]
+    output_video = f"/tmp/{input_video_name}_subtitled.mp4"
+
+    if soft_subtitle:
+        stream = ffmpeg.output(
+            video_input_stream, subtitle_input_stream, output_video,
+            **{"c": "copy", "c:s": "mov_text"}
+        )
+    else:
+        stream = ffmpeg.output(
+            video_input_stream, output_video,
+            vf=f"subtitles={subtitle_file}"
+        )
+
+    ffmpeg.run(stream, overwrite_output=True)
+    return output_video
+
+def process_video(input_video, text_transcription, video_language, target_language):
+    srt_path = text_to_srt(text_transcription)
+    translated_srt_path = translate_srt(srt_path, video_language, target_language)
+    output_video = add_subtitle_to_video(input_video.name, translated_srt_path)
+    return output_video
+
+language_url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
+video_language_options = fetch_languages(language_url)
+
+with gr.Blocks() as app:
+    with gr.Row():
+        input_video = gr.Video(label="Video File")
+        text_transcription = gr.TextArea(label="Text Transcription")
+        video_language = gr.Dropdown(choices=video_language_options, label="Language of the Video")
+        target_language = gr.Dropdown(choices=video_language_options, label="Language Translated")
+    output_video = gr.Video(label="Video with Translated Subtitles")
+
+    input_video.change(fn=process_video, inputs=[input_video, text_transcription, video_language, target_language], outputs=output_video)
+
+app.launch()