ui/ui_components.py

#!/usr/bin/env python3
"""
UI Components for BackgroundFX Pro
Complete interface with all features
"""

import gradio as gr
import os

# Background style choices
BG_STYLES = ["minimalist", "office_modern", "studio_blue", "studio_green", "warm_gradient", "tech_dark"]

def create_interface():
    """Create the full BackgroundFX Pro interface"""
    
    # Import callbacks inside function to avoid circular imports
    from ui.callbacks import (
        cb_load_models,
        cb_process_video,
        cb_cancel,
        cb_status,
        cb_clear,
        cb_generate_bg,
        cb_use_gen_bg,
        cb_preset_bg_preview,
    )
    
    with gr.Blocks(title="🎬 BackgroundFX Pro", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
            # 🎬 BackgroundFX Pro
            ### AI-Powered Video Background Replacement
            Replace your video background with professional quality using SAM2 + MatAnyone
            """
        )
        
        with gr.Tabs():
            # Main Processing Tab
            with gr.Tab("🎥 Process Video"):
                with gr.Row():
                    # Left Column - Inputs
                    with gr.Column(scale=1):
                        video_input = gr.Video(
                            label="Upload Video",
                            interactive=True
                        )
                        
                        gr.Markdown("### Background Settings")
                        
                        bg_method = gr.Radio(
                            label="Background Source",
                            choices=[
                                "Preset Styles",
                                "Upload Custom",
                                "Generate AI Background"
                            ],
                            value="Preset Styles"
                        )
                        
                        # Preset styles
                        with gr.Group(visible=True) as preset_group:
                            bg_style = gr.Dropdown(
                                label="Background Style",
                                choices=BG_STYLES,
                                value="minimalist"
                            )
                            preset_preview = gr.Image(
                                label="Preview",
                                interactive=False,
                                height=200
                            )
                        
                        # Custom upload
                        with gr.Group(visible=False) as upload_group:
                            custom_bg = gr.Image(
                                label="Upload Background Image",
                                type="filepath",
                                interactive=True
                            )
                        
                        # AI Generation
                        with gr.Group(visible=False) as generate_group:
                            prompt = gr.Textbox(
                                label="Describe Background",
                                value="modern office space",
                                placeholder="e.g., 'warm sunset studio', 'minimalist white room'"
                            )
                            with gr.Row():
                                gen_width = gr.Slider(640, 1920, 1280, step=10, label="Width")
                                gen_height = gr.Slider(360, 1080, 720, step=10, label="Height")
                            with gr.Row():
                                bokeh = gr.Slider(0, 30, 8, step=1, label="Blur")
                                vignette = gr.Slider(0, 0.6, 0.15, step=0.01, label="Vignette")
                            btn_generate = gr.Button("✨ Generate Background", variant="secondary")
                            gen_preview = gr.Image(label="Generated Background", interactive=False)
                            gen_path = gr.Textbox(visible=False)
                            btn_use_generated = gr.Button("Use This Background", variant="secondary")
                        
                        # Processing Options
                        with gr.Accordion("⚙️ Advanced Options", open=False):
                            use_two_stage = gr.Checkbox(
                                label="Use Two-Stage Processing",
                                value=True,
                                info="Better quality but slower"
                            )
                            chroma_preset = gr.Dropdown(
                                label="Chroma Preset",
                                choices=["standard"],
                                value="standard"
                            )
                            key_color = gr.Dropdown(
                                label="Key Color Mode",
                                choices=["auto", "green", "blue"],
                                value="auto"
                            )
                            quality = gr.Radio(
                                label="Quality",
                                choices=["speed", "balanced", "max"],
                                value="balanced"
                            )
                        
                        # Action Buttons
                        gr.Markdown("### Actions")
                        with gr.Row():
                            btn_load = gr.Button("📦 Load Models", variant="secondary", scale=1)
                            btn_process = gr.Button("🎬 Process Video", variant="primary", scale=2)
                        with gr.Row():
                            btn_cancel = gr.Button("⏹ Cancel", variant="stop")
                            btn_clear = gr.Button("🗑️ Clear All")
                    
                    # Right Column - Outputs
                    with gr.Column(scale=1):
                        video_output = gr.Video(
                            label="Processed Video",
                            interactive=False
                        )
                        status = gr.Textbox(
                            label="Status",
                            lines=8,
                            max_lines=12,
                            interactive=False
                        )
                        with gr.Row():
                            processing_time = gr.Textbox(
                                label="Processing Time",
                                interactive=False
                            )
                            quality_info = gr.Textbox(
                                label="Quality Mode",
                                interactive=False
                            )
            
            # Status Tab
            with gr.Tab("📊 System Status"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Model Status")
                        model_status = gr.JSON(label="Models")
                    with gr.Column():
                        gr.Markdown("### System Info")
                        system_status = gr.JSON(label="System")
                btn_refresh_status = gr.Button("🔄 Refresh Status")
        
        # Wire up all the callbacks
        
        # Background method switching
        def switch_bg_method(method):
            return (
                gr.update(visible=(method == "Preset Styles")),
                gr.update(visible=(method == "Upload Custom")),
                gr.update(visible=(method == "Generate AI Background"))
            )
        
        bg_method.change(
            switch_bg_method,
            inputs=[bg_method],
            outputs=[preset_group, upload_group, generate_group]
        )
        
        # Preset preview
        bg_style.change(
            cb_preset_bg_preview,
            inputs=[bg_style],
            outputs=[preset_preview]
        )
        
        # AI Generation
        btn_generate.click(
            cb_generate_bg,
            inputs=[prompt, gen_width, gen_height, bokeh, vignette, gr.State(1.0)],
            outputs=[gen_preview, gen_path]
        )
        
        btn_use_generated.click(
            cb_use_gen_bg,
            inputs=[gen_path],
            outputs=[custom_bg]
        )
        
        # Main actions
        btn_load.click(
            cb_load_models,
            outputs=[status]
        )
        
        # Process with proper background selection
        def process_with_bg(video, method, style, custom, gen_path, two_stage, chroma, key, quality):
            # Set quality environment variable
            os.environ["BFX_QUALITY"] = quality.lower()
            
            # Determine which background to use
            if method == "Upload Custom":
                bg_choice = "custom"
                bg_path = custom
            elif method == "Generate AI Background":
                bg_choice = "custom"
                bg_path = gen_path
            else:  # Preset Styles
                bg_choice = style
                bg_path = None
            
            return cb_process_video(
                video, bg_choice, bg_path, two_stage, 
                chroma, key, False, False
            )
        
        btn_process.click(
            process_with_bg,
            inputs=[
                video_input, bg_method, bg_style, custom_bg, 
                gen_path, use_two_stage, chroma_preset, key_color, quality
            ],
            outputs=[video_output, status]
        )
        
        btn_cancel.click(
            cb_cancel,
            outputs=[status]
        )
        
        btn_clear.click(
            cb_clear,
            outputs=[video_output, status, gen_preview, gen_path, custom_bg]
        )
        
        btn_refresh_status.click(
            cb_status,
            outputs=[model_status, system_status]
        )
        
        # Update quality info when changed
        quality.change(
            lambda q: f"Mode: {q} ({'fastest' if q=='speed' else 'best quality' if q=='max' else 'balanced'})",
            inputs=[quality],
            outputs=[quality_info]
        )
        
        # COMMENTED OUT to prevent cv2 import on startup - fixes WebAssembly error
        # Users can still get preview by changing the dropdown selection
        # demo.load(
        #     cb_preset_bg_preview,
        #     inputs=[bg_style],
        #     outputs=[preset_preview]
        # )
    
    return demo