Spaces:
Running
Running
| # app.py - Main Gradio application | |
| import gradio as gr | |
| import whisper | |
| import torch | |
| from transformers import MarianMTModel, MarianTokenizer | |
| import yt_dlp | |
| import os | |
| import tempfile | |
| import subprocess | |
| from pathlib import Path | |
| import re | |
| class SubtitleTranslator: | |
| def __init__(self): | |
| # Use the smallest Whisper model for speed | |
| self.whisper_model = whisper.load_model("tiny") | |
| # Translation model cache | |
| self.translation_models = {} | |
| self.tokenizers = {} | |
| def download_youtube_audio(self, url): | |
| """Download audio from YouTube video""" | |
| try: | |
| ydl_opts = { | |
| 'format': 'bestaudio/best', | |
| 'outtmpl': 'temp_audio.%(ext)s', | |
| 'postprocessors': [{ | |
| 'key': 'FFmpegExtractAudio', | |
| 'preferredcodec': 'mp3', | |
| 'preferredquality': '192', | |
| }], | |
| } | |
| with yt_dlp.YoutubeDL(ydl_opts) as ydl: | |
| ydl.download([url]) | |
| # Find the downloaded file | |
| for file in os.listdir('.'): | |
| if file.startswith('temp_audio') and file.endswith('.mp3'): | |
| return file | |
| return None | |
| except Exception as e: | |
| return None | |
| def extract_audio_from_video(self, video_path): | |
| """Extract audio from uploaded video file""" | |
| try: | |
| audio_path = "temp_extracted_audio.wav" | |
| cmd = [ | |
| 'ffmpeg', '-i', video_path, | |
| '-acodec', 'pcm_s16le', | |
| '-ac', '1', | |
| '-ar', '16000', | |
| audio_path, '-y' | |
| ] | |
| subprocess.run(cmd, check=True, capture_output=True) | |
| return audio_path | |
| except Exception as e: | |
| return None | |
| def transcribe_audio(self, audio_path): | |
| """Transcribe audio using Whisper""" | |
| result = self.whisper_model.transcribe(audio_path) | |
| return result | |
| def get_translation_model(self, source_lang, target_lang="en"): | |
| """Load translation model for language pair""" | |
| model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}" | |
| try: | |
| if model_name not in self.translation_models: | |
| self.tokenizers[model_name] = MarianTokenizer.from_pretrained(model_name) | |
| self.translation_models[model_name] = MarianMTModel.from_pretrained(model_name) | |
| return self.translation_models[model_name], self.tokenizers[model_name] | |
| except: | |
| # Fallback to multilingual model | |
| fallback_model = "Helsinki-NLP/opus-mt-mul-en" | |
| if fallback_model not in self.translation_models: | |
| self.tokenizers[fallback_model] = MarianTokenizer.from_pretrained(fallback_model) | |
| self.translation_models[fallback_model] = MarianMTModel.from_pretrained(fallback_model) | |
| return self.translation_models[fallback_model], self.tokenizers[fallback_model] | |
| def translate_text(self, text, source_lang, target_lang="en"): | |
| """Translate text using MarianMT""" | |
| if source_lang == target_lang: | |
| return text | |
| try: | |
| model, tokenizer = self.get_translation_model(source_lang, target_lang) | |
| inputs = tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=512) | |
| translated = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True) | |
| return tokenizer.decode(translated[0], skip_special_tokens=True) | |
| except: | |
| return text # Return original if translation fails | |
| def format_timestamp(self, seconds): | |
| """Convert seconds to SRT timestamp format""" | |
| hours = int(seconds // 3600) | |
| minutes = int((seconds % 3600) // 60) | |
| secs = int(seconds % 60) | |
| millisecs = int((seconds % 1) * 1000) | |
| return f"{hours:02d}:{minutes:02d}:{secs:02d},{millisecs:03d}" | |
| def create_srt(self, segments, source_lang): | |
| """Create SRT subtitle content""" | |
| srt_content = "" | |
| for i, segment in enumerate(segments, 1): | |
| start_time = self.format_timestamp(segment['start']) | |
| end_time = self.format_timestamp(segment['end']) | |
| original_text = segment['text'].strip() | |
| translated_text = self.translate_text(original_text, source_lang, "en") | |
| srt_content += f"{i}\n" | |
| srt_content += f"{start_time} --> {end_time}\n" | |
| srt_content += f"{translated_text}\n\n" | |
| return srt_content | |
| def process_video(self, video_input, youtube_url): | |
| """Main processing function""" | |
| try: | |
| # Determine input source | |
| if youtube_url and youtube_url.strip(): | |
| audio_path = self.download_youtube_audio(youtube_url.strip()) | |
| if not audio_path: | |
| return "Error: Could not download YouTube video", None | |
| elif video_input: | |
| audio_path = self.extract_audio_from_video(video_input) | |
| if not audio_path: | |
| return "Error: Could not extract audio from video", None | |
| else: | |
| return "Please provide either a video file or YouTube URL", None | |
| # Transcribe audio | |
| result = self.transcribe_audio(audio_path) | |
| # Detect language | |
| detected_lang = result.get('language', 'unknown') | |
| # Language code mapping for translation models | |
| lang_mapping = { | |
| 'spanish': 'es', 'french': 'fr', 'german': 'de', 'italian': 'it', | |
| 'portuguese': 'pt', 'russian': 'ru', 'chinese': 'zh', 'japanese': 'ja', | |
| 'korean': 'ko', 'arabic': 'ar', 'hindi': 'hi', 'dutch': 'nl', | |
| 'swedish': 'sv', 'norwegian': 'no', 'danish': 'da', 'finnish': 'fi' | |
| } | |
| source_lang_code = lang_mapping.get(detected_lang, detected_lang) | |
| # Create SRT content | |
| srt_content = self.create_srt(result['segments'], source_lang_code) | |
| # Save SRT file | |
| srt_filename = "translated_subtitles.srt" | |
| with open(srt_filename, 'w', encoding='utf-8') as f: | |
| f.write(srt_content) | |
| # Clean up temporary files | |
| if os.path.exists(audio_path): | |
| os.remove(audio_path) | |
| status_msg = f"✅ Processing complete!\n" | |
| status_msg += f"🔍 Detected language: {detected_lang}\n" | |
| status_msg += f"📝 Generated {len(result['segments'])} subtitle segments\n" | |
| status_msg += f"🌍 Translated to English" | |
| return status_msg, srt_filename | |
| except Exception as e: | |
| return f"Error during processing: {str(e)}", None | |
| # Initialize the translator | |
| translator = SubtitleTranslator() | |
| # Create Gradio interface | |
| def process_video_interface(video_file, youtube_url, progress=gr.Progress()): | |
| progress(0.1, desc="Starting processing...") | |
| progress(0.3, desc="Extracting audio...") | |
| result = translator.process_video(video_file, youtube_url) | |
| progress(0.7, desc="Transcribing and translating...") | |
| progress(1.0, desc="Complete!") | |
| return result | |
| # Custom CSS for better UI | |
| css = """ | |
| .gradio-container { | |
| max-width: 900px !important; | |
| } | |
| .title { | |
| text-align: center; | |
| color: #2563eb; | |
| font-size: 2.5rem; | |
| font-weight: bold; | |
| margin-bottom: 1rem; | |
| } | |
| .subtitle { | |
| text-align: center; | |
| color: #64748b; | |
| font-size: 1.2rem; | |
| margin-bottom: 2rem; | |
| } | |
| .feature-box { | |
| background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); | |
| color: white; | |
| padding: 1rem; | |
| border-radius: 10px; | |
| margin: 1rem 0; | |
| } | |
| """ | |
| # Create the Gradio app | |
| with gr.Blocks(css=css, title="Video Subtitle Translator") as app: | |
| gr.HTML(""" | |
| <div class="title">🎬 Video Subtitle Translator</div> | |
| <div class="subtitle">Generate English subtitles from any language video using AI</div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.HTML(""" | |
| <div class="feature-box"> | |
| <h3>🚀 Features:</h3> | |
| <ul> | |
| <li>📹 Upload video files or paste YouTube links</li> | |
| <li>🎯 Automatic speech recognition with Whisper AI</li> | |
| <li>🌍 Auto-detect source language</li> | |
| <li>📝 Generate accurate English subtitles</li> | |
| <li>⏱️ Perfect timing synchronization</li> | |
| <li>💾 Download ready-to-use SRT files</li> | |
| </ul> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| video_input = gr.File( | |
| label="📁 Upload Video File", | |
| file_types=[".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v"], | |
| type="filepath" | |
| ) | |
| youtube_input = gr.Textbox( | |
| label="🔗 Or paste YouTube URL", | |
| placeholder="https://www.youtube.com/watch?v=...", | |
| lines=1 | |
| ) | |
| process_btn = gr.Button( | |
| "🚀 Generate Subtitles", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| with gr.Column(scale=1): | |
| status_output = gr.Textbox( | |
| label="📊 Processing Status", | |
| lines=6, | |
| interactive=False | |
| ) | |
| srt_output = gr.File( | |
| label="💾 Download SRT File", | |
| interactive=False | |
| ) | |
| gr.HTML(""" | |
| <div style="text-align: center; margin-top: 2rem; color: #64748b;"> | |
| <p>⚡ Powered by Whisper AI & MarianMT | 🤗 Running on Hugging Face Spaces</p> | |
| <p>💡 Tip: For best results, use videos with clear audio and minimal background noise</p> | |
| </div> | |
| """) | |
| # Connect the processing function | |
| process_btn.click( | |
| fn=process_video_interface, | |
| inputs=[video_input, youtube_input], | |
| outputs=[status_output, srt_output], | |
| show_progress=True | |
| ) | |
| if __name__ == "__main__": | |
| app.launch() |