ui_components.py

#!/usr/bin/env python3
"""
FIXED UI Components for Video Background Replacement
Updated to use the fixed core processing functions with working SAM2 + MatAnyone
"""

import gradio as gr
import tempfile
import cv2
import os
import time
import logging
from typing import Optional

# Import FIXED core processing functions
from app import load_models_with_validation, process_video_fixed, get_model_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 with actual model state - FIXED"""
    try:
        status = get_model_status()
        
        sam2_status = f"✅ Loaded & Validated" if status['sam2'] == 'Ready' else "❌ Not loaded"
        matanyone_status = f"✅ Loaded & Validated" if status['matanyone'] == 'Ready' else "❌ Not loaded"
        models_status = "✅ Models Validated" if status['validated'] else "❌ Not validated"
        
        return f"SAM2: {sam2_status}\nMatAnyone: {matanyone_status}\nStatus: {models_status}"
    except Exception as e:
        return f"Status check error: {str(e)}"

# ============================================================================ #
# 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 (FIXED)
# ============================================================================ #
def process_video_enhanced_fixed(
    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
):
    """FIXED enhanced video processing function that uses validated SAM2 + MatAnyone"""
    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_fixed(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_fixed(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_fixed(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_fixed(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_fixed(progress: Optional[gr.Progress] = None):
    """FIXED model loading with Gradio progress updates and validation"""
    def progress_callback(pct, desc):
        gradio_progress_wrapper(progress, pct, desc)
    
    return load_models_with_validation(progress_callback)

# ============================================================================ #
# MAIN INTERFACE CREATION (FIXED)
# ============================================================================ #
def create_interface():
    """Create the complete Gradio interface with FIXED processing"""
    
    with gr.Blocks(
        title="FIXED High-Quality Video Background Replacement",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container { max-width: 1200px !important; }
        .progress-bar { background: linear-gradient(90deg, #3498db, #2ecc71) !important; }
        .status-box { background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 6px; }
        """
    ) as demo:
        gr.Markdown("# FIXED Cinema-Quality Video Background Replacement")
        gr.Markdown("**Upload a video → Choose a background → Get professional results with VALIDATED AI**")
        gr.Markdown("*Powered by PROPERLY INTEGRATED SAM2 + MatAnyone with comprehensive validation*")
        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 and validate the AI models, then process your video*")
                with gr.Row():
                    load_models_btn = gr.Button("Step 1: Load & Validate AI Models", variant="secondary")
                    process_btn = gr.Button("Step 2: Process Video (FIXED)", variant="primary")

                status_text = gr.Textbox(
                    label="System Status", 
                    value="Ready - click 'Load & Validate AI Models' to start",
                    interactive=False, 
                    lines=4,
                    elem_classes=["status-box"]
                )

                # Enhanced cache/validation status indicator
                validation_status = gr.Textbox(
                    label="AI Models Validation Status",
                    value=update_cache_status(),
                    interactive=False,
                    lines=3,
                    elem_classes=["status-box"]
                )

            with gr.Column(scale=1):
                gr.Markdown("### Your Results")
                gr.Markdown("*FIXED processed video with validated SAM2 + MatAnyone will appear here*")
                video_output = gr.Video(label="Your FIXED Processed Video", height=400)
                result_text = gr.Textbox(
                    label="Processing Results",
                    interactive=False,
                    lines=8,
                    placeholder="FIXED processing status and detailed results will appear here...",
                    elem_classes=["status-box"]
                )

                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 with FIXED functions
        load_models_btn.click(
            fn=load_models_with_progress_fixed, 
            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_fixed,
            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]
        )

        # Auto-update validation status when models are loaded
        load_models_btn.click(
            fn=update_cache_status,
            outputs=[validation_status]
        )

        with gr.Accordion("FIXED Quality & Features", open=False):
            gr.Markdown("""
            ### FIXED Single-Stage Cinema-Quality Features:
            **Processing**: Original → Final Background (VALIDATED SAM2 + MatAnyone)
            **Quality**: Comprehensive validation, enhanced fallbacks, edge feathering, gamma correction
            **Models**: SAM2 with functionality testing + MatAnyone with proper error handling
            **Segmentation**: Multi-point SAM2 strategy with advanced OpenCV fallback
            **Refinement**: MatAnyone refinement with enhanced OpenCV backup
            **Architecture**: Separated UI components with proper error handling and validation
            **Output**: H.264 CRF 18, AAC 192kbps with preserved audio
            """)

        gr.Markdown("---")
        gr.Markdown("*FIXED Cinema-Quality Video Background Replacement — Validated SAM2 + MatAnyone pipeline*")

    return demo