Create ui_components.py

ui_components.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+"""
+UI Components for Video Background Replacement
+Contains all Gradio interface components separated from core processing
+"""
+
+import gradio as gr
+import tempfile
+import cv2
+import os
+import time
+import logging
+from typing import Optional
+
+# Import core processing functions
+from app import load_models_fast, process_video_core, get_cache_status
+from utilities import PROFESSIONAL_BACKGROUNDS, create_procedural_background
+
+logger = logging.getLogger(__name__)
+
+# ============================================================================ #
+# GRADIO MONKEY PATCH (BUG FIX for gradio>=4.44.0)
+# ============================================================================ #
+try:
+    import gradio_client.utils as gc_utils
+    original_get_type = gc_utils.get_type
+    def patched_get_type(schema):
+        if not isinstance(schema, dict):
+            if isinstance(schema, bool):
+                return "boolean"
+            return "string"
+        return original_get_type(schema)
+    gc_utils.get_type = patched_get_type
+    logger.info("Applied Gradio schema validation monkey patch.")
+except (ImportError, AttributeError) as e:
+    logger.warning(f"Could not apply Gradio monkey patch: {e}")
+
+# ============================================================================ #
+# UI HELPER FUNCTIONS
+# ============================================================================ #
+def generate_ai_background(prompt, style):
+    """Generate AI background from prompt"""
+    if not prompt or not prompt.strip():
+        return None, "Please enter a prompt"
+    try:
+        bg_image = create_procedural_background(prompt, style, 1920, 1080)
+        if bg_image is not None:
+            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
+                cv2.imwrite(tmp.name, bg_image)
+                return tmp.name, f"Background generated: {prompt[:50]}..."
+        return None, "Generation failed, try different prompt"
+    except Exception as e:
+        logger.error(f"AI generation error: {e}")
+        return None, f"Generation error: {str(e)}"
+
+def switch_background_method(method):
+    """Switch between background method UI groups"""
+    return (
+        gr.update(visible=(method == "upload")),
+        gr.update(visible=(method == "professional")),
+        gr.update(visible=(method == "colors")),
+        gr.update(visible=(method == "ai"))
+    )
+
+def update_cache_status():
+    """Update cache status display"""
+    status = get_cache_status()
+    sam2_status = "✅ Cached" if status["sam2_cached"] else "❌ Not cached"
+    matanyone_status = "✅ Cached" if status["matanyone_cached"] else "❌ Not cached"
+    return f"SAM2: {sam2_status}\nMatAnyone: {matanyone_status}"
+
+# ============================================================================ #
+# PROGRESS WRAPPER
+# ============================================================================ #
+def gradio_progress_wrapper(progress_obj, pct, desc):
+    """Wrapper to safely call Gradio progress"""
+    try:
+        if progress_obj is not None:
+            progress_obj(pct, desc=desc)
+    except Exception:
+        pass
+
+# ============================================================================ #
+# MAIN PROCESSING FUNCTION FOR UI
+# ============================================================================ #
+def process_video_enhanced(
+    video_path, bg_method, custom_img, prof_choice, grad_type,
+    color1, color2, color3, use_third, ai_prompt, ai_style, ai_img,
+    progress: Optional[gr.Progress] = None
+):
+    """Enhanced video processing function that handles all UI inputs"""
+    if not video_path:
+        return None, "No video file provided."
+    
+    def progress_callback(pct, desc):
+        gradio_progress_wrapper(progress, pct, desc)
+    
+    try:
+        if bg_method == "upload":
+            if custom_img and os.path.exists(custom_img):
+                return process_video_core(video_path, "custom", custom_img, progress_callback)
+            return None, "No image uploaded. Please upload a background image."
+            
+        elif bg_method == "professional":
+            if prof_choice and prof_choice in PROFESSIONAL_BACKGROUNDS:
+                return process_video_core(video_path, prof_choice, None, progress_callback)
+            return None, f"Invalid professional background: {prof_choice}"
+            
+        elif bg_method == "colors":
+            try:
+                colors = [color1 or "#3498db", color2 or "#2ecc71"]
+                if use_third and color3:
+                    colors.append(color3)
+                    
+                from utilities import create_professional_background
+                bg_config = {
+                    "type": "gradient" if grad_type != "solid" else "color",
+                    "colors": colors if grad_type != "solid" else [colors[0]],
+                    "direction": grad_type if grad_type != "solid" else "vertical"
+                }
+                gradient_bg = create_professional_background(bg_config, 1920, 1080)
+                temp_path = f"/tmp/gradient_{int(time.time())}.png"
+                cv2.imwrite(temp_path, gradient_bg)
+                return process_video_core(video_path, "custom", temp_path, progress_callback)
+            except Exception as e:
+                return None, f"Error creating gradient: {str(e)}"
+                
+        elif bg_method == "ai":
+            if ai_img and os.path.exists(ai_img):
+                return process_video_core(video_path, "custom", ai_img, progress_callback)
+            return None, "No AI background generated. Click 'Generate Background' first."
+            
+        else:
+            return None, f"Unknown background method: {bg_method}"
+            
+    except Exception as e:
+        logger.error(f"Enhanced processing error: {e}")
+        return None, f"Processing error: {str(e)}"
+
+def load_models_with_progress(progress: Optional[gr.Progress] = None):
+    """Load models with Gradio progress updates"""
+    def progress_callback(pct, desc):
+        gradio_progress_wrapper(progress, pct, desc)
+    
+    return load_models_fast(progress_callback)
+
+# ============================================================================ #
+# MAIN INTERFACE CREATION
+# ============================================================================ #
+def create_interface():
+    """Create the complete Gradio interface"""
+    
+    with gr.Blocks(
+        title="ENHANCED High-Quality Video Background Replacement - Fast Restart",
+        theme=gr.themes.Soft(),
+        css="""
+        .gradio-container { max-width: 1200px !important; }
+        .progress-bar { background: linear-gradient(90deg, #3498db, #2ecc71) !important; }
+        """
+    ) as demo:
+        gr.Markdown("# Cinema-Quality Video Background Replacement")
+        gr.Markdown("**Upload a video → Choose a background → Get professional results with AI**")
+        gr.Markdown("*Powered by SAM2 + MatAnyone with fast restart caching (2-4 min restarts instead of 10+)*")
+        gr.Markdown("---")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### Step 1: Upload Your Video")
+                gr.Markdown("*Supports MP4, MOV, AVI, and other common formats*")
+                video_input = gr.Video(label="Drop your video here", height=300)
+
+                gr.Markdown("### Step 2: Choose Background Method")
+                gr.Markdown("*Select your preferred background creation method*")
+                background_method = gr.Radio(
+                    choices=["upload", "professional", "colors", "ai"],
+                    value="professional",
+                    label="Background Method"
+                )
+                gr.Markdown(
+                    "- **upload** = Upload Image  \n"
+                    "- **professional** = Professional Presets  \n"
+                    "- **colors** = Colors/Gradients  \n"
+                    "- **ai** = AI Generated"
+                )
+
+                with gr.Group(visible=False) as upload_group:
+                    gr.Markdown("**Upload Your Background Image**")
+                    custom_background = gr.Image(label="Drop your background image here", type="filepath")
+
+                with gr.Group(visible=True) as professional_group:
+                    gr.Markdown("**Professional Background Presets**")
+                    professional_choice = gr.Dropdown(
+                        choices=list(PROFESSIONAL_BACKGROUNDS.keys()),
+                        value="office_modern",
+                        label="Select Professional Background"
+                    )
+
+                with gr.Group(visible=False) as colors_group:
+                    gr.Markdown("**Custom Colors & Gradients**")
+                    gradient_type = gr.Dropdown(
+                        choices=["solid", "vertical", "horizontal", "diagonal", "radial", "soft_radial"],
+                        value="vertical",
+                        label="Gradient Type"
+                    )
+                    with gr.Row():
+                        color1 = gr.ColorPicker(label="Color 1", value="#3498db")
+                        color2 = gr.ColorPicker(label="Color 2", value="#2ecc71")
+                    with gr.Row():
+                        color3 = gr.ColorPicker(label="Color 3", value="#e74c3c")
+                        use_third_color = gr.Checkbox(label="Use 3rd color", value=False)
+
+                with gr.Group(visible=False) as ai_group:
+                    gr.Markdown("**AI Generated Background**")
+                    ai_prompt = gr.Textbox(
+                        label="Describe your background",
+                        placeholder="e.g., 'modern office with plants', 'sunset over mountains', 'abstract tech pattern'",
+                        lines=2
+                    )
+                    ai_style = gr.Dropdown(
+                        choices=["photorealistic", "artistic", "abstract", "minimalist", "corporate", "nature"],
+                        value="photorealistic",
+                        label="Style"
+                    )
+                    with gr.Row():
+                        generate_ai_btn = gr.Button("Generate Background", variant="secondary")
+                        ai_generated_image = gr.Image(label="Generated Background", type="filepath", visible=False)
+
+                # Wire up background method switching
+                background_method.change(
+                    fn=switch_background_method,
+                    inputs=background_method,
+                    outputs=[upload_group, professional_group, colors_group, ai_group]
+                )
+
+                gr.Markdown("### Processing Controls")
+                gr.Markdown("*First load the AI models, then process your video*")
+                with gr.Row():
+                    load_models_btn = gr.Button("Step 1: Load AI Models (Fast)", variant="secondary")
+                    process_btn = gr.Button("Step 2: Process Video", variant="primary")
+
+                status_text = gr.Textbox(
+                    label="System Status", 
+                    value="Ready - click 'Load AI Models' to start",
+                    interactive=False, 
+                    lines=3
+                )
+
+                # Cache status indicator
+                cache_status = gr.Textbox(
+                    label="Fast Restart Cache Status",
+                    value=update_cache_status(),
+                    interactive=False,
+                    lines=2
+                )
+
+            with gr.Column(scale=1):
+                gr.Markdown("### Your Results")
+                gr.Markdown("*Processed video will appear here after Step 2*")
+                video_output = gr.Video(label="Your Processed Video", height=400)
+                result_text = gr.Textbox(
+                    label="Processing Results",
+                    interactive=False,
+                    lines=6,
+                    placeholder="Processing status and results will appear here..."
+                )
+
+                gr.Markdown("### Professional Backgrounds Available")
+                bg_preview_html = """
+                <div style='display: grid; grid-template-columns: repeat(3, 1fr); gap: 8px; padding: 10px; max-height: 400px; overflow-y: auto; border: 1px solid #ddd; border-radius: 8px;'>
+                """
+                for key, config in PROFESSIONAL_BACKGROUNDS.items():
+                    colors = config["colors"]
+                    gradient = f"linear-gradient(45deg, {colors[0]}, {colors[-1]})" if len(colors) >= 2 else colors[0]
+                    bg_preview_html += f"""
+                    <div style='padding: 12px 8px; border: 1px solid #ddd; border-radius: 6px; text-align: center; background: {gradient};
+                                min-height: 60px; display: flex; align-items: center; justify-content: center;'>
+                        <div>
+                            <strong style='color: white; text-shadow: 1px 1px 2px rgba(0,0,0,0.8); font-size: 12px; display: block;'>{config["name"]}</strong>
+                            <small style='color: rgba(255,255,255,0.9); text-shadow: 1px 1px 1px rgba(0,0,0,0.6); font-size: 10px;'>{config.get("description", "")[:30]}...</small>
+                        </div>
+                    </div>
+                    """
+                bg_preview_html += "</div>"
+                gr.HTML(bg_preview_html)
+
+        # Event handlers
+        load_models_btn.click(
+            fn=load_models_with_progress, 
+            outputs=status_text
+        )
+        
+        generate_ai_btn.click(
+            fn=generate_ai_background, 
+            inputs=[ai_prompt, ai_style], 
+            outputs=[ai_generated_image, status_text]
+        )
+        
+        process_btn.click(
+            fn=process_video_enhanced,
+            inputs=[video_input, background_method, custom_background, professional_choice,
+                    gradient_type, color1, color2, color3, use_third_color,
+                    ai_prompt, ai_style, ai_generated_image],
+            outputs=[video_output, result_text]
+        )
+
+        with gr.Accordion("ENHANCED Quality & Features", open=False):
+            gr.Markdown("""
+            ### SINGLE-STAGE Cinema-Quality Features:
+            **Processing**: Original → Final Background (SAM2 + MatAnyone)
+            **Quality**: Edge feathering, gamma correction, mask cleanup, H.264 CRF 18, AAC 192kbps.
+            **Fast Restart**: Persistent model caching reduces restart time from 10+ minutes to 2-4 minutes.
+            **Architecture**: Separated UI components for faster startup and better maintainability.
+            """)
+
+        gr.Markdown("---")
+        gr.Markdown("*Cinema-Quality Video Background Replacement — SINGLE-STAGE pipeline with Fast Restart*")
+
+    return demo