Nick021402 commited on
Commit
8a5a458
Β·
verified Β·
1 Parent(s): 6e5f830

Create App.py

Browse files
Files changed (1) hide show
  1. App.py +289 -0
App.py ADDED
@@ -0,0 +1,289 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py - Main Gradio application
2
+ import gradio as gr
3
+ import whisper
4
+ import torch
5
+ from transformers import MarianMTModel, MarianTokenizer
6
+ import yt_dlp
7
+ import os
8
+ import tempfile
9
+ import subprocess
10
+ from pathlib import Path
11
+ import re
12
+
13
+ class SubtitleTranslator:
14
+ def __init__(self):
15
+ # Use the smallest Whisper model for speed
16
+ self.whisper_model = whisper.load_model("tiny")
17
+
18
+ # Translation model cache
19
+ self.translation_models = {}
20
+ self.tokenizers = {}
21
+
22
+ def download_youtube_audio(self, url):
23
+ """Download audio from YouTube video"""
24
+ try:
25
+ ydl_opts = {
26
+ 'format': 'bestaudio/best',
27
+ 'outtmpl': 'temp_audio.%(ext)s',
28
+ 'postprocessors': [{
29
+ 'key': 'FFmpegExtractAudio',
30
+ 'preferredcodec': 'mp3',
31
+ 'preferredquality': '192',
32
+ }],
33
+ }
34
+
35
+ with yt_dlp.YoutubeDL(ydl_opts) as ydl:
36
+ ydl.download([url])
37
+
38
+ # Find the downloaded file
39
+ for file in os.listdir('.'):
40
+ if file.startswith('temp_audio') and file.endswith('.mp3'):
41
+ return file
42
+ return None
43
+ except Exception as e:
44
+ return None
45
+
46
+ def extract_audio_from_video(self, video_path):
47
+ """Extract audio from uploaded video file"""
48
+ try:
49
+ audio_path = "temp_extracted_audio.wav"
50
+ cmd = [
51
+ 'ffmpeg', '-i', video_path,
52
+ '-acodec', 'pcm_s16le',
53
+ '-ac', '1',
54
+ '-ar', '16000',
55
+ audio_path, '-y'
56
+ ]
57
+ subprocess.run(cmd, check=True, capture_output=True)
58
+ return audio_path
59
+ except Exception as e:
60
+ return None
61
+
62
+ def transcribe_audio(self, audio_path):
63
+ """Transcribe audio using Whisper"""
64
+ result = self.whisper_model.transcribe(audio_path)
65
+ return result
66
+
67
+ def get_translation_model(self, source_lang, target_lang="en"):
68
+ """Load translation model for language pair"""
69
+ model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
70
+
71
+ try:
72
+ if model_name not in self.translation_models:
73
+ self.tokenizers[model_name] = MarianTokenizer.from_pretrained(model_name)
74
+ self.translation_models[model_name] = MarianMTModel.from_pretrained(model_name)
75
+
76
+ return self.translation_models[model_name], self.tokenizers[model_name]
77
+ except:
78
+ # Fallback to multilingual model
79
+ fallback_model = "Helsinki-NLP/opus-mt-mul-en"
80
+ if fallback_model not in self.translation_models:
81
+ self.tokenizers[fallback_model] = MarianTokenizer.from_pretrained(fallback_model)
82
+ self.translation_models[fallback_model] = MarianMTModel.from_pretrained(fallback_model)
83
+ return self.translation_models[fallback_model], self.tokenizers[fallback_model]
84
+
85
+ def translate_text(self, text, source_lang, target_lang="en"):
86
+ """Translate text using MarianMT"""
87
+ if source_lang == target_lang:
88
+ return text
89
+
90
+ try:
91
+ model, tokenizer = self.get_translation_model(source_lang, target_lang)
92
+ inputs = tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=512)
93
+ translated = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
94
+ return tokenizer.decode(translated[0], skip_special_tokens=True)
95
+ except:
96
+ return text # Return original if translation fails
97
+
98
+ def format_timestamp(self, seconds):
99
+ """Convert seconds to SRT timestamp format"""
100
+ hours = int(seconds // 3600)
101
+ minutes = int((seconds % 3600) // 60)
102
+ secs = int(seconds % 60)
103
+ millisecs = int((seconds % 1) * 1000)
104
+ return f"{hours:02d}:{minutes:02d}:{secs:02d},{millisecs:03d}"
105
+
106
+ def create_srt(self, segments, source_lang):
107
+ """Create SRT subtitle content"""
108
+ srt_content = ""
109
+
110
+ for i, segment in enumerate(segments, 1):
111
+ start_time = self.format_timestamp(segment['start'])
112
+ end_time = self.format_timestamp(segment['end'])
113
+
114
+ original_text = segment['text'].strip()
115
+ translated_text = self.translate_text(original_text, source_lang, "en")
116
+
117
+ srt_content += f"{i}\n"
118
+ srt_content += f"{start_time} --> {end_time}\n"
119
+ srt_content += f"{translated_text}\n\n"
120
+
121
+ return srt_content
122
+
123
+ def process_video(self, video_input, youtube_url):
124
+ """Main processing function"""
125
+ try:
126
+ # Determine input source
127
+ if youtube_url and youtube_url.strip():
128
+ audio_path = self.download_youtube_audio(youtube_url.strip())
129
+ if not audio_path:
130
+ return "Error: Could not download YouTube video", None
131
+ elif video_input:
132
+ audio_path = self.extract_audio_from_video(video_input)
133
+ if not audio_path:
134
+ return "Error: Could not extract audio from video", None
135
+ else:
136
+ return "Please provide either a video file or YouTube URL", None
137
+
138
+ # Transcribe audio
139
+ result = self.transcribe_audio(audio_path)
140
+
141
+ # Detect language
142
+ detected_lang = result.get('language', 'unknown')
143
+
144
+ # Language code mapping for translation models
145
+ lang_mapping = {
146
+ 'spanish': 'es', 'french': 'fr', 'german': 'de', 'italian': 'it',
147
+ 'portuguese': 'pt', 'russian': 'ru', 'chinese': 'zh', 'japanese': 'ja',
148
+ 'korean': 'ko', 'arabic': 'ar', 'hindi': 'hi', 'dutch': 'nl',
149
+ 'swedish': 'sv', 'norwegian': 'no', 'danish': 'da', 'finnish': 'fi'
150
+ }
151
+
152
+ source_lang_code = lang_mapping.get(detected_lang, detected_lang)
153
+
154
+ # Create SRT content
155
+ srt_content = self.create_srt(result['segments'], source_lang_code)
156
+
157
+ # Save SRT file
158
+ srt_filename = "translated_subtitles.srt"
159
+ with open(srt_filename, 'w', encoding='utf-8') as f:
160
+ f.write(srt_content)
161
+
162
+ # Clean up temporary files
163
+ if os.path.exists(audio_path):
164
+ os.remove(audio_path)
165
+
166
+ status_msg = f"βœ… Processing complete!\n"
167
+ status_msg += f"πŸ” Detected language: {detected_lang}\n"
168
+ status_msg += f"πŸ“ Generated {len(result['segments'])} subtitle segments\n"
169
+ status_msg += f"🌍 Translated to English"
170
+
171
+ return status_msg, srt_filename
172
+
173
+ except Exception as e:
174
+ return f"Error during processing: {str(e)}", None
175
+
176
+ # Initialize the translator
177
+ translator = SubtitleTranslator()
178
+
179
+ # Create Gradio interface
180
+ def process_video_interface(video_file, youtube_url, progress=gr.Progress()):
181
+ progress(0.1, desc="Starting processing...")
182
+
183
+ progress(0.3, desc="Extracting audio...")
184
+ result = translator.process_video(video_file, youtube_url)
185
+
186
+ progress(0.7, desc="Transcribing and translating...")
187
+ progress(1.0, desc="Complete!")
188
+
189
+ return result
190
+
191
+ # Custom CSS for better UI
192
+ css = """
193
+ .gradio-container {
194
+ max-width: 900px !important;
195
+ }
196
+ .title {
197
+ text-align: center;
198
+ color: #2563eb;
199
+ font-size: 2.5rem;
200
+ font-weight: bold;
201
+ margin-bottom: 1rem;
202
+ }
203
+ .subtitle {
204
+ text-align: center;
205
+ color: #64748b;
206
+ font-size: 1.2rem;
207
+ margin-bottom: 2rem;
208
+ }
209
+ .feature-box {
210
+ background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
211
+ color: white;
212
+ padding: 1rem;
213
+ border-radius: 10px;
214
+ margin: 1rem 0;
215
+ }
216
+ """
217
+
218
+ # Create the Gradio app
219
+ with gr.Blocks(css=css, title="Video Subtitle Translator") as app:
220
+ gr.HTML("""
221
+ <div class="title">🎬 Video Subtitle Translator</div>
222
+ <div class="subtitle">Generate English subtitles from any language video using AI</div>
223
+ """)
224
+
225
+ with gr.Row():
226
+ with gr.Column():
227
+ gr.HTML("""
228
+ <div class="feature-box">
229
+ <h3>πŸš€ Features:</h3>
230
+ <ul>
231
+ <li>πŸ“Ή Upload video files or paste YouTube links</li>
232
+ <li>🎯 Automatic speech recognition with Whisper AI</li>
233
+ <li>🌍 Auto-detect source language</li>
234
+ <li>πŸ“ Generate accurate English subtitles</li>
235
+ <li>⏱️ Perfect timing synchronization</li>
236
+ <li>πŸ’Ύ Download ready-to-use SRT files</li>
237
+ </ul>
238
+ </div>
239
+ """)
240
+
241
+ with gr.Row():
242
+ with gr.Column(scale=1):
243
+ video_input = gr.File(
244
+ label="πŸ“ Upload Video File",
245
+ file_types=[".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v"],
246
+ type="filepath"
247
+ )
248
+
249
+ youtube_input = gr.Textbox(
250
+ label="πŸ”— Or paste YouTube URL",
251
+ placeholder="https://www.youtube.com/watch?v=...",
252
+ lines=1
253
+ )
254
+
255
+ process_btn = gr.Button(
256
+ "πŸš€ Generate Subtitles",
257
+ variant="primary",
258
+ size="lg"
259
+ )
260
+
261
+ with gr.Column(scale=1):
262
+ status_output = gr.Textbox(
263
+ label="πŸ“Š Processing Status",
264
+ lines=6,
265
+ interactive=False
266
+ )
267
+
268
+ srt_output = gr.File(
269
+ label="πŸ’Ύ Download SRT File",
270
+ interactive=False
271
+ )
272
+
273
+ gr.HTML("""
274
+ <div style="text-align: center; margin-top: 2rem; color: #64748b;">
275
+ <p>⚑ Powered by Whisper AI & MarianMT | πŸ€— Running on Hugging Face Spaces</p>
276
+ <p>πŸ’‘ Tip: For best results, use videos with clear audio and minimal background noise</p>
277
+ </div>
278
+ """)
279
+
280
+ # Connect the processing function
281
+ process_btn.click(
282
+ fn=process_video_interface,
283
+ inputs=[video_input, youtube_input],
284
+ outputs=[status_output, srt_output],
285
+ show_progress=True
286
+ )
287
+
288
+ if __name__ == "__main__":
289
+ app.launch()