Create app.py

app.py

@@ -0,0 +1,180 @@
+import gradio as gr
+import whisper
+import moviepy.editor as mp
+from moviepy.video.tools.subtitles import SubtitlesClip
+from datetime import timedelta
+import os
+from transformers import MarianMTModel, MarianTokenizer
+import torch
+import spaces
+
+# Dictionary of supported languages and their codes for MarianMT
+LANGUAGE_CODES = {
+    "English": "en",
+    "Spanish": "es",
+    "French": "fr",
+    "German": "de",
+    "Italian": "it",
+    "Portuguese": "pt",
+    "Russian": "ru",
+    "Chinese": "zh",
+    "Japanese": "ja",
+    "Korean": "ko"
+}
+
+# Mapping of language pairs to MarianMT model names
+def get_model_name(source_lang, target_lang):
+    return f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
+
+def format_timestamp(seconds):
+    """Convert seconds to SRT timestamp format"""
+    td = timedelta(seconds=seconds)
+    hours = td.seconds//3600
+    minutes = (td.seconds//60)%60
+    seconds = td.seconds%60
+    milliseconds = td.microseconds//1000
+    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
+
+def translate_text(text, source_lang, target_lang):
+    """Translate text using MarianMT"""
+    if source_lang == target_lang:
+        return text
+    
+    try:
+        model_name = get_model_name(source_lang, target_lang)
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        
+        # Tokenize and translate
+        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        translated = model.generate(**inputs)
+        translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
+        
+        return translated_text
+    except Exception as e:
+        print(f"Translation error: {e}")
+        return text  # Return original text if translation fails
+
+def create_srt(segments, target_lang="en"):
+    """Convert whisper segments to SRT format with optional translation"""
+    srt_content = ""
+    for i, segment in enumerate(segments, start=1):
+        start_time = format_timestamp(segment['start'])
+        end_time = format_timestamp(segment['end'])
+        text = segment['text'].strip()
+        
+        # Translate if target language is different
+        if 'language' in segment and segment['language'] != target_lang:
+            text = translate_text(text, segment['language'], target_lang)
+        
+        srt_content += f"{i}\n{start_time} --> {end_time}\n{text}\n\n"
+    return srt_content
+
+def create_subtitle_clips(segments, videosize, target_lang="en"):
+    """Create subtitle clips for moviepy with translation support"""
+    subtitle_clips = []
+    
+    for segment in segments:
+        start_time = segment['start']
+        end_time = segment['end']
+        duration = end_time - start_time
+        text = segment['text'].strip()
+        
+        # Translate if target language is different
+        if 'language' in segment and segment['language'] != target_lang:
+            text = translate_text(text, segment['language'], target_lang)
+        
+        text_clip = mp.TextClip(
+            text,
+            font='Arial',
+            fontsize=24,
+            color='white',
+            stroke_color='black',
+            stroke_width=1,
+            size=videosize,
+            method='caption'
+        ).set_position(('center', 'bottom'))
+        
+        text_clip = text_clip.set_start(start_time).set_duration(duration)
+        subtitle_clips.append(text_clip)
+    
+    return subtitle_clips
+
+@spaces.GPU
+def process_video(video_path, target_lang="en"):
+    """Main function to process video and add subtitles with translation"""
+    # Load Whisper model
+    model = whisper.load_model("base")
+    
+    # Extract audio from video
+    video = mp.VideoFileClip(video_path)
+    audio = video.audio
+    
+    # Save audio temporarily
+    temp_audio_path = "temp_audio.wav"
+    audio.write_audiofile(temp_audio_path)
+    
+    # Transcribe audio
+    result = model.transcribe(temp_audio_path)
+    
+    # Add detected language to segments
+    for segment in result["segments"]:
+        segment['language'] = result["language"]
+    
+    # Create SRT content with translation
+    srt_content = create_srt(result["segments"], target_lang)
+    
+    # Save SRT file
+    video_name = os.path.splitext(os.path.basename(video_path))[0]
+    srt_path = f"{video_name}_subtitles_{target_lang}.srt"
+    with open(srt_path, "w", encoding="utf-8") as f:
+        f.write(srt_content)
+    
+    # Create subtitle clips with translation
+    subtitle_clips = create_subtitle_clips(result["segments"], video.size, target_lang)
+    
+    # Combine video with subtitles
+    final_video = mp.CompositeVideoClip([video] + subtitle_clips)
+    
+    # Save final video
+    output_video_path = f"{video_name}_with_subtitles_{target_lang}.mp4"
+    final_video.write_videofile(output_video_path)
+    
+    # Clean up
+    os.remove(temp_audio_path)
+    video.close()
+    final_video.close()
+    
+    return output_video_path, srt_path
+
+def gradio_interface(video_file, target_language):
+    """Gradio interface function with language selection"""
+    try:
+        video_path = video_file.name
+        target_lang = LANGUAGE_CODES[target_language]
+        output_video, srt_file = process_video(video_path, target_lang)
+        return output_video, srt_file
+    except Exception as e:
+        return str(e), None
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Video(label="Upload Video"),
+        gr.Dropdown(
+            choices=list(LANGUAGE_CODES.keys()),
+            value="English",
+            label="Target Language"
+        )
+    ],
+    outputs=[
+        gr.Video(label="Video with Subtitles"),
+        gr.File(label="SRT Subtitle File")
+    ],
+    title="Video Subtitler with Translation",
+    description="Upload a video to generate subtitles, translate them to your chosen language, and embed them directly in the video."
+)
+
+if __name__ == "__main__":
+    iface.launch()