Upload google_colabs/gradio_app.py with huggingface_hub

google_colabs/gradio_app.py

@@ -26,11 +26,13 @@ sys.path.insert(0, os.path.join(PROJECT_ROOT, 'backend'))
 from backend import config
 from backend.config import InpaintMode
 from backend.main import SubtitleRemover
+from backend.subtitle_extractor import SubtitleExtractor
 
 
 def process_video(
     video_file,
     algorithm,
+    extract_subtitles,
     sttn_skip_detection,
     sttn_max_load,
     lama_super_fast,
@@ -40,9 +42,9 @@ def process_video(
     diffueraser_max_load,
     progress=gr.Progress()
 ):
-    """Process video to remove subtitles"""
+    """Process video to remove subtitles and optionally extract them"""
     if video_file is None:
-        return None, "❌ Please upload a video first"
+        return None, None, "❌ Please upload a video first"
     
     try:
         progress(0, desc="Initializing...")
@@ -87,14 +89,32 @@ def process_video(
         remover = SubtitleRemover(video_path, gui_mode=False)
         remover.run()
         
+        srt_file = None
+        
+        # Extract subtitles if requested
+        if extract_subtitles:
+            progress(0.8, desc="Extracting subtitles...")
+            try:
+                extractor = SubtitleExtractor(video_path)
+                srt_file = extractor.extract_to_srt(
+                    progress_callback=lambda p, desc: progress(0.8 + p * 0.2, desc)
+                )
+            except Exception as e:
+                print(f"Warning: Subtitle extraction failed: {e}")
+                srt_file = None
+        
         progress(1.0, desc="Complete!")
         
-        return remover.video_out_name, f"✅ Success! Processed with {algorithm}"
+        status_msg = f"✅ Success! Processed with {algorithm}"
+        if srt_file:
+            status_msg += f"\n📄 Subtitles extracted to SRT file"
+        
+        return remover.video_out_name, srt_file, status_msg
         
     except Exception as e:
         import traceback
         error_msg = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
-        return None, error_msg
+        return None, None, error_msg
 
 
 def create_interface():
@@ -129,6 +149,12 @@ def create_interface():
                     info="STTN: Fastest. DiffuEraser: Best quality (needs 12GB+ VRAM)."
                 )
                 
+                extract_subs = gr.Checkbox(
+                    label="📝 Extract Subtitles to SRT",
+                    value=False,
+                    info="Extract detected subtitles using OCR and save as SRT file"
+                )
+                
                 with gr.Accordion("⚙️ Advanced Settings", open=False):
                     with gr.Tab("STTN"):
                         sttn_skip = gr.Checkbox(
@@ -176,6 +202,7 @@ def create_interface():
             with gr.Column():
                 # Output
                 video_output = gr.Video(label="✨ Output Video")
+                srt_output = gr.File(label="📄 Extracted Subtitles (SRT)", visible=True)
                 status_output = gr.Textbox(label="📊 Status", lines=3)
                 
                 gr.Markdown("""
@@ -196,23 +223,23 @@ def create_interface():
         process_btn.click(
             fn=process_video,
             inputs=[
-                video_input, algorithm,
+                video_input, algorithm, extract_subs,
                 sttn_skip, sttn_max_load,
                 lama_fast,
                 sd_steps, sd_guidance,
                 diffueraser_steps, diffueraser_max_load
             ],
-            outputs=[video_output, status_output]
+            outputs=[video_output, srt_output, status_output]
         )
         
         # Examples
         gr.Examples(
             examples=[
-                ["STTN (Fast)", True, 40, False, 50, 7.5, 50, 40],
-                ["DiffuEraser (Recommended)", True, 80, False, 50, 7.5, 50, 40],
+                ["STTN (Fast)", False, True, 40, False, 50, 7.5, 50, 40],
+                ["DiffuEraser (Recommended)", True, True, 80, False, 50, 7.5, 50, 40],
             ],
             inputs=[
-                algorithm,
+                algorithm, extract_subs,
                 sttn_skip, sttn_max_load,
                 lama_fast,
                 sd_steps, sd_guidance,