Add Gradio app for video subtitling

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+from utils import process_video  # Import your backend logic
+
+# Define supported languages
+language_map = {
+    "English": None,
+    "Hindi": "Helsinki-NLP/opus-mt-en-hi",
+    "Spanish": "Helsinki-NLP/opus-mt-en-es",
+    "French": "Helsinki-NLP/opus-mt-en-fr",
+    "German": "Helsinki-NLP/opus-mt-en-de",
+    "Telugu": "facebook/nllb-200-distilled-600M",
+    "Portuguese": "Helsinki-NLP/opus-mt-en-pt",
+    "Russian": "Helsinki-NLP/opus-mt-en-ru",
+    "Chinese": "Helsinki-NLP/opus-mt-en-zh",
+    "Arabic": "Helsinki-NLP/opus-mt-en-ar",
+    "Japanese": "Helsinki-NLP/opus-mt-en-jap"
+}
+
+def generate_subtitles(video_file, language):
+    """
+    Process the uploaded video and generate subtitles.
+    """
+    try:
+        srt_path = process_video(video_file, language)
+        return srt_path  # Return the path to the generated SRT file
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Define Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# AI-Powered Video Subtitling")
+    gr.Markdown("Upload a video and select a language to generate subtitles.")
+
+    with gr.Row():
+        video_input = gr.File(label="Upload Video File", file_types=["mp4", "mkv", "avi"])
+        language_dropdown = gr.Dropdown(
+            choices=list(language_map.keys()),
+            label="Select Subtitle Language",
+            value="English"
+        )
+
+    generate_button = gr.Button("Generate Subtitles")
+    output_srt = gr.File(label="Download Subtitles")
+
+    generate_button.click(
+        generate_subtitles,
+        inputs=[video_input, language_dropdown],
+        outputs=output_srt
+    )
+
+# Launch Gradio App
+demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+gradio==3.41.2
+transformers==4.35.2
+whisper==20230314
+torch==2.0.1

utils.py

@@ -0,0 +1,66 @@
+import whisper
+from transformers import MarianMTModel, MarianTokenizer, AutoTokenizer, AutoModelForSeq2SeqLM
+import os
+
+# Load Whisper model
+model = whisper.load_model("base")
+
+def process_video(video_file, language):
+    # Save uploaded video locally
+    video_path = "/tmp/video.mp4"
+    with open(video_path, "wb") as f:
+        f.write(video_file.read())
+
+    try:
+        print("Transcribing video to English...")
+        result = model.transcribe(video_path, language="en")
+
+        segments = []
+        if language == "English":
+            segments = result["segments"]
+        else:
+            if language == "Telugu":
+                model_name = "facebook/nllb-200-distilled-600M"
+                tokenizer = AutoTokenizer.from_pretrained(model_name)
+                translation_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+                tgt_lang = "tel_Telu"
+                print(f"Translating to Telugu using NLLB-200 Distilled...")
+                for segment in result["segments"]:
+                    inputs = tokenizer(segment["text"], return_tensors="pt", padding=True)
+                    translated_tokens = translation_model.generate(**inputs, forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang))
+                    translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+                    segments.append({"text": translated_text, "start": segment["start"], "end": segment["end"]})
+            else:
+                model_map = {
+                    "Hindi": "Helsinki-NLP/opus-mt-en-hi",
+                    "Spanish": "Helsinki-NLP/opus-mt-en-es",
+                    "French": "Helsinki-NLP/opus-mt-en-fr",
+                    "German": "Helsinki-NLP/opus-mt-en-de",
+                    "Portuguese": "Helsinki-NLP/opus-mt-en-pt",
+                    "Russian": "Helsinki-NLP/opus-mt-en-ru",
+                    "Chinese": "Helsinki-NLP/opus-mt-en-zh",
+                    "Arabic": "Helsinki-NLP/opus-mt-en-ar",
+                    "Japanese": "Helsinki-NLP/opus-mt-en-jap"
+                }
+                model_name = model_map[language]
+                tokenizer = MarianTokenizer.from_pretrained(model_name)
+                translation_model = MarianMTModel.from_pretrained(model_name)
+                print(f"Translating to {language}...")
+                for segment in result["segments"]:
+                    inputs = tokenizer(segment["text"], return_tensors="pt", padding=True)
+                    translated = translation_model.generate(**inputs)
+                    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+                    segments.append({"text": translated_text, "start": segment["start"], "end": segment["end"]})
+
+        # Create SRT file
+        srt_path = "/tmp/subtitles.srt"
+        with open(srt_path, "w", encoding="utf-8") as f:
+            for i, segment in enumerate(segments, 1):
+                start = f"{segment['start']:.3f}".replace(".", ",")
+                end = f"{segment['end']:.3f}".replace(".", ",")
+                text = segment["text"].strip()
+                f.write(f"{i}\n00:00:{start} --> 00:00:{end}\n{text}\n\n")
+        return srt_path
+
+    except Exception as e:
+        return f"Error: {str(e)}"