Update app.py

app.py

@@ -6,11 +6,10 @@ import os
 import math
 from datetime import timedelta
 import subprocess
+import re
 
 # --- Configuration ---
 TRANSLATION_MODEL = "facebook/nllb-200-distilled-1.3B" 
-# We use OpenAI's original small model for better segmentation on CPU
-# It is often better at splitting sentences than Distil-Large for subtitles
 WHISPER_MODEL = "openai/whisper-small"
 
 print("Loading Models...")
@@ -32,8 +31,19 @@ whisper_pipe = pipeline(
 print("Models Loaded Successfully!")
 
 # ---------------------------------------------------------
-# Helper: Extract Audio
+# Helper: Extract Audio & Duration
 # ---------------------------------------------------------
+def get_media_duration(filename):
+    try:
+        result = subprocess.run(
+            ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", filename],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT
+        )
+        return float(result.stdout)
+    except:
+        return 30.0 
+
 def extract_audio(video_path):
     output_audio_path = "temp_audio.mp3"
     if os.path.exists(output_audio_path):
@@ -48,10 +58,30 @@ def extract_audio(video_path):
     return output_audio_path
 
 # ---------------------------------------------------------
-# Helper: Smart SRT Splitter (The Fix!)
+# Helper: VTT Converter (For Browser Preview)
+# ---------------------------------------------------------
+def srt_to_vtt(srt_path):
+    """Converts SRT to VTT format for the HTML5 video player."""
+    vtt_path = srt_path.replace(".srt", ".vtt")
+    with open(srt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    # VTT Header
+    vtt_content = "WEBVTT\n\n"
+    
+    # Replace comma timestamps (00:00:01,000) with dot (00:00:01.000)
+    # And copy the rest
+    vtt_content += re.sub(r'(\d{2}:\d{2}:\d{2}),(\d{3})', r'\1.\2', content)
+    
+    with open(vtt_path, 'w', encoding='utf-8') as f:
+        f.write(vtt_content)
+    
+    return vtt_path
+
+# ---------------------------------------------------------
+# Helper: Smart Splitter logic
 # ---------------------------------------------------------
 def split_text_into_lines(text, max_chars=80):
-    """Breaks long text into smaller lines based on character limit."""
     words = text.split()
     lines = []
     current_line = []
@@ -70,11 +100,7 @@ def split_text_into_lines(text, max_chars=80):
         lines.append(" ".join(current_line))
     return lines
 
-def create_srt_segments(chunks):
-    """
-    Takes raw Whisper chunks and breaks them down into clean SRT subtitles.
-    Distributes time proportionally if a chunk is split into multiple lines.
-    """
+def create_srt_segments(chunks, total_video_duration):
     srt_subtitles = []
     index_counter = 1
 
@@ -82,23 +108,147 @@ def create_srt_segments(chunks):
         text = chunk['text'].strip()
         timestamp = chunk['timestamp']
         
-        # Safe unpacking of timestamps
         if isinstance(timestamp, (list, tuple)):
             start_time, end_time = timestamp
         else:
-            continue # Skip bad chunks
-            
-        if end_time is None: end_time = start_time + 5.0
-
-        # Smart Split: If text is too long (>80 chars), split it
+            start_time, end_time = 0.0, None
+        
+        if end_time is None: end_time = total_video_duration
+        
         lines = split_text_into_lines(text, max_chars=80)
+        duration = end_time - start_time
+        if duration <= 0: duration = 5.0 
         
-        # Calculate duration per line (Proportional split)
-        total_duration = end_time - start_time
-        duration_per_line = total_duration / len(lines) if lines else 0
-
+        step = duration / len(lines) if lines else 0
         current_start = start_time
         
+        for line in lines:
+            current_end = current_start + step
+            srt_subtitles.append(
+                srt.Subtitle(index=index_counter, start=timedelta(seconds=current_start), end=timedelta(seconds=current_end), content=line)
+            )
+            index_counter += 1
+            current_start = current_end
+
+    return srt_subtitles
+
+# ---------------------------------------------------------
+# Logic 1: Translation (NLLB)
+# ---------------------------------------------------------
+def batch_translate(texts, src_lang, tgt_lang, batch_size=8, progress=gr.Progress()):
+    results = []
+    tokenizer_nllb.src_lang = src_lang
+    
+    for i, start_idx in enumerate(range(0, len(texts), batch_size)):
+        batch = texts[start_idx : start_idx + batch_size]
+        inputs = tokenizer_nllb(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        forced_bos_token_id = tokenizer_nllb.convert_tokens_to_ids(tgt_lang)
+        with torch.no_grad():
+            generated_tokens = model_nllb.generate(**inputs, forced_bos_token_id=forced_bos_token_id, max_length=512)
+        results.extend(tokenizer_nllb.batch_decode(generated_tokens, skip_special_tokens=True))
+    return results
+
+def process_translation(filepath, src_lang_code, tgt_lang_code):
+    if filepath is None: return None, None
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            subtitles = list(srt.parse(f.read()))
+    except Exception as e:
+        return f"Error: {str(e)}", None
+
+    texts = [sub.content for sub in subtitles]
+    translated = batch_translate(texts, src_lang_code, tgt_lang_code)
+    
+    for sub, trans in zip(subtitles, translated):
+        sub.content = trans
+
+    out_path = "translated_subtitles.srt"
+    with open(out_path, 'w', encoding='utf-8') as f:
+        f.write(srt.compose(subtitles))
+    
+    return out_path
+
+# ---------------------------------------------------------
+# Logic 2: Video to SRT + Preview
+# ---------------------------------------------------------
+def video_to_srt(video_path, progress=gr.Progress()):
+    if video_path is None: return None, None
+    
+    # 1. Audio & Duration
+    progress(0.1, desc="Extracting Audio...")
+    try:
+        audio_path = extract_audio(video_path)
+        duration = get_media_duration(audio_path)
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+    
+    # 2. Transcribe
+    progress(0.3, desc="Transcribing...")
+    outputs = whisper_pipe(audio_path, return_timestamps=True, generate_kwargs={"language": "english"})
+    chunks = outputs.get("chunks", [])
+    if not chunks: chunks = [{"text": outputs.get("text", ""), "timestamp": (0.0, None)}]
+
+    # 3. Format SRT
+    progress(0.8, desc="Formatting...")
+    srt_subtitles = create_srt_segments(chunks, duration)
+    
+    srt_path = "generated_captions.srt"
+    with open(srt_path, 'w', encoding='utf-8') as f:
+        f.write(srt.compose(srt_subtitles))
+    
+    # 4. Create Preview (HTML + VTT)
+    vtt_path = srt_to_vtt(srt_path)
+    
+    # We use Gradio's /file/ route to serve local files
+    html_preview = f"""
+    <h3>Video Preview</h3>
+    <video controls width="100%" height="400px" style="background:black">
+        <source src="/file={video_path}" type="video/mp4">
+        <track kind="captions" src="/file={vtt_path}" srclang="en" label="English" default>
+        Your browser does not support the video tag.
+    </video>
+    <p style="margin-top:10px; color: #666;">Note: Subtitles are overlaid for preview only. They are not burned into the video.</p>
+    """
+        
+    return srt_path, html_preview
+
+# ---------------------------------------------------------
+# Gradio Interface
+# ---------------------------------------------------------
+with gr.Blocks(title="SRT Master Tool") as demo:
+    gr.Markdown("# 🎬 Auto Subtitle & Translator")
+    
+    with gr.Tabs():
+        # --- TAB 1 ---
+        with gr.TabItem("Step 1: Video to SRT"):
+            gr.Markdown("### 1. Upload Video -> 2. Check Preview -> 3. Download SRT")
+            with gr.Row():
+                video_input = gr.Video(label="Upload Video", sources=["upload"])
+                
+                with gr.Column():
+                    # The new Preview Player
+                    preview_output = gr.HTML(label="Preview Player")
+                    # The download button
+                    srt_output_gen = gr.File(label="Download Generated SRT")
+            
+            btn1 = gr.Button("Generate SRT & Preview", variant="primary")
+            btn1.click(video_to_srt, inputs=video_input, outputs=[srt_output_gen, preview_output])
+
+        # --- TAB 2 ---
+        with gr.TabItem("Step 2: Translate SRT"):
+            gr.Markdown("### Translate Subtitles to Arabic")
+            with gr.Row():
+                srt_input = gr.File(label="Upload SRT")
+                with gr.Column():
+                    src_l = gr.Dropdown(["eng_Latn", "fra_Latn"], label="From", value="eng_Latn")
+                    tgt_l = gr.Dropdown(["arb_Arab", "arz_Arab"], label="To", value="arb_Arab")
+                srt_output_trans = gr.File(label="Translated SRT")
+            
+            btn2 = gr.Button("Translate", variant="primary")
+            btn2.click(process_translation, inputs=[srt_input, src_l, tgt_l], outputs=srt_output_trans)
+
+if __name__ == "__main__":
+    demo.launch()
         for line in lines:
             current_end = current_start + duration_per_line