ui_components.py

#!/usr/bin/env python3
"""
Minimal Working UI Components - BackgroundFX Pro
"""
import gradio as gr
import os
import time
import traceback
from typing import Optional

# Apply Gradio schema patch at the module level
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"
            if isinstance(schema, str):
                return "string"
            if isinstance(schema, (int, float)):
                return "number"
            return "string"
        return original_get_type(schema)
    
    gc_utils.get_type = patched_get_type
    print("✅ UI Components: Gradio schema patch applied")
    
except Exception as e:
    print(f"⚠️ UI Components: Gradio patch failed: {e}")

# Import core functions with error handling
try:
    from app import (
        load_models_with_validation, 
        process_video_fixed, 
        get_model_status, 
        get_cache_status, 
        PROCESS_CANCELLED
    )
    CORE_FUNCTIONS_AVAILABLE = True
    print("✅ UI Components: Core functions imported")
except Exception as e:
    print(f"❌ UI Components: Core functions import failed: {e}")
    CORE_FUNCTIONS_AVAILABLE = False

# Import utilities with error handling
try:
    from utilities import PROFESSIONAL_BACKGROUNDS
    UTILITIES_AVAILABLE = True
    print("✅ UI Components: Utilities imported")
except Exception as e:
    print(f"❌ UI Components: Utilities import failed: {e}")
    UTILITIES_AVAILABLE = False
    PROFESSIONAL_BACKGROUNDS = {}

# Import two-stage with error handling
try:
    from two_stage_processor import CHROMA_PRESETS
    TWO_STAGE_AVAILABLE = True
    print("✅ UI Components: Two-stage processor available")
except ImportError:
    TWO_STAGE_AVAILABLE = False
    CHROMA_PRESETS = {'standard': {}}
    print("⚠️ UI Components: Two-stage processor not available")

def create_interface():
    """Create the main Gradio interface"""
    
    # Safe processing function
    def safe_process_video(video_path, bg_method, custom_img, prof_choice, 
                          use_two_stage, chroma_preset, progress: Optional[gr.Progress] = None):
        """Safe wrapper for video processing"""
        if not CORE_FUNCTIONS_AVAILABLE:
            return None, "Core processing functions not available", "Error: Core functions not loaded"
        
        try:
            # Simple progress callback
            def progress_callback(pct, desc):
                if progress:
                    progress(pct)
                print(f"Progress: {pct:.1%} - {desc}")
                return desc  # Match app.py's expectation
            
            # Call the main processing function
            if bg_method == "professional" and prof_choice:
                result = process_video_fixed(
                    video_path, prof_choice, None, 
                    progress_callback, 
                    use_two_stage=bool(use_two_stage),
                    chroma_preset=chroma_preset or "standard",
                    preview_mask=False,
                    preview_greenscreen=False
                )
            elif bg_method == "upload" and custom_img:
                result = process_video_fixed(
                    video_path, "custom", custom_img, 
                    progress_callback, 
                    use_two_stage=bool(use_two_stage),
                    chroma_preset=chroma_preset or "standard",
                    preview_mask=False,
                    preview_greenscreen=False
                )
            else:
                return None, "Please select a valid background method", "Error: Invalid background selection"
            
            return result[0], result[1], f"Processing completed successfully"
            
        except Exception as e:
            error_msg = f"Processing error: {str(e)}\n{traceback.format_exc()}"
            print(f"Error in safe_process_video: {error_msg}")
            return None, error_msg, f"Error: {error_msg}"
    
    # Safe model loading function
    def safe_load_models(progress: Optional[gr.Progress] = None):
        """Safe wrapper for model loading"""
        if not CORE_FUNCTIONS_AVAILABLE:
            return "Error: Core functions not available", "Core functions not loaded"
        
        try:
            def progress_callback(pct, desc):
                if progress:
                    progress(pct)
                print(f"Model loading: {pct:.1%} - {desc}")
                return desc  # Match app.py's expectation
            
            result = load_models_with_validation(progress_callback)
            return result, f"Model loading completed"
            
        except Exception as e:
            error_msg = f"Model loading error: {str(e)}\n{traceback.format_exc()}"
            print(f"Error in safe_load_models: {error_msg}")
            return error_msg, error_msg
    
    # Create the interface
    with gr.Blocks(
        title="BackgroundFX Pro - Video Background Replacement",
        theme=gr.themes.Soft(),
    ) as demo:
        
        gr.Markdown("# 🎬 BackgroundFX Pro - Video Background Replacement")
        gr.Markdown("Professional quality background replacement using AI segmentation")
        
        if TWO_STAGE_AVAILABLE:
            gr.Markdown("✅ **Two-Stage Green Screen Mode Available** - Cinema-grade processing")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📹 Step 1: Upload Your Video")
                video_input = gr.Video(label="Upload your video", height=300)
                
                gr.Markdown("### 🎨 Step 2: Choose Background Method")
                background_method = gr.Radio(
                    choices=["professional", "upload"],
                    value="professional",
                    label="Background Method"
                )
                
                # Professional backgrounds
                with gr.Group() as professional_group:
                    gr.Markdown("**Professional Presets**")
                    if UTILITIES_AVAILABLE and PROFESSIONAL_BACKGROUNDS:
                        choices = list(PROFESSIONAL_BACKGROUNDS.keys())
                        default_choice = choices[0] if choices else "office_modern"
                    else:
                        choices = ["office_modern", "studio_white", "nature_blur"]
                        default_choice = "office_modern"
                    
                    professional_choice = gr.Dropdown(
                        choices=choices,
                        value=default_choice,
                        label="Select Background"
                    )
                
                # Custom upload
                with gr.Group(visible=False) as upload_group:
                    gr.Markdown("**Upload Custom Background**")
                    custom_background = gr.Image(label="Upload background image", type="filepath")
                
                # Show/hide groups based on method
                def update_visibility(method):
                    return (
                        gr.update(visible=(method == "professional")),
                        gr.update(visible=(method == "upload"))
                    )
                
                background_method.change(
                    fn=update_visibility,
                    inputs=background_method,
                    outputs=[professional_group, upload_group]
                )
                
                gr.Markdown("### ⚙️ Processing Options")
                with gr.Accordion("Advanced Settings", open=False):
                    use_two_stage = gr.Checkbox(
                        label="Enable Two-Stage Processing",
                        value=False,
                        info="Better quality but slower"
                    )
                    
                    if TWO_STAGE_AVAILABLE:
                        chroma_preset = gr.Dropdown(
                            choices=list(CHROMA_PRESETS.keys()),
                            value="standard",
                            label="Quality Preset"
                        )
                    else:
                        chroma_preset = gr.Dropdown(
                            choices=["standard"],
                            value="standard",
                            label="Quality Preset"
                        )
                
                gr.Markdown("### 🚀 Process Video")
                with gr.Row():
                    load_models_btn = gr.Button("Load Models", variant="secondary")
                    process_btn = gr.Button("Process Video", variant="primary", scale=2)
                
                status_text = gr.Textbox(
                    label="Status", 
                    value="Ready - Click 'Load Models' first",
                    interactive=False,
                    lines=3
                )
                
                # Add model and cache status buttons
                with gr.Row():
                    model_status_btn = gr.Button("Check Model Status", variant="secondary")
                    cache_status_btn = gr.Button("Check Cache Status", variant="secondary")
                
                model_status = gr.JSON(label="Model Status", interactive=False)
                cache_status = gr.JSON(label="Cache Status", interactive=False)
            
            with gr.Column(scale=1):
                gr.Markdown("### 🎥 Result")
                video_output = gr.Video(label="Processed Video", height=400)
                
                result_text = gr.Textbox(
                    label="Processing Info",
                    interactive=False,
                    lines=10,
                    placeholder="Processing results will appear here..."
                )
                
                debug_text = gr.Textbox(
                    label="Debug Log",
                    interactive=False,
                    lines=8,
                    placeholder="Debug information will appear here..."
                )
        
        # Event handlers
        load_models_btn.click(
            fn=safe_load_models,
            outputs=[status_text, debug_text],
            show_progress=True
        )
        
        process_btn.click(
            fn=safe_process_video,
            inputs=[
                video_input,
                background_method,
                custom_background,
                professional_choice,
                use_two_stage,
                chroma_preset
            ],
            outputs=[video_output, result_text, debug_text],
            show_progress=True
        )
        
        model_status_btn.click(
            fn=get_model_status,
            outputs=[model_status],
            show_progress=False
        )
        
        cache_status_btn.click(
            fn=get_cache_status,
            outputs=[cache_status],
            show_progress=False
        )
        
        # Info section
        with gr.Accordion("ℹ️ Information", open=False):
            gr.Markdown(f"""
            ### System Status:
            - **Core Functions**: {"✅ Available" if CORE_FUNCTIONS_AVAILABLE else "❌ Not Available"}
            - **Utilities**: {"✅ Available" if UTILITIES_AVAILABLE else "❌ Not Available"}
            - **Two-Stage Mode**: {"✅ Available" if TWO_STAGE_AVAILABLE else "❌ Not Available"}
            
            ### Performance Tips:
            - Use shorter videos (10-30 seconds) for testing
            - Two-stage mode provides better quality but takes longer
            - Professional backgrounds are optimized for best results
            """)
    
    return demo