ui_components.py

#!/usr/bin/env python3
"""
Enhanced UI Components - BackgroundFX Pro
Streamlined interface with better error handling and user experience
"""

import gradio as gr
import os
import json
import time
import traceback
from typing import Optional, Dict, Any, Tuple
from pathlib import Path

# Remove redundant Gradio schema patching - handled in app.py
print("UI Components: Initializing interface...")

# Import core functions with comprehensive error handling
try:
    from app import (
        VideoProcessor,
        processor,
        load_models_with_validation, 
        process_video_fixed, 
        get_model_status, 
        get_cache_status
    )
    CORE_FUNCTIONS_AVAILABLE = True
    print("UI Components: Core functions imported successfully")
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 successfully")
except Exception as e:
    print(f"UI Components: Utilities import failed: {e}")
    UTILITIES_AVAILABLE = False
    PROFESSIONAL_BACKGROUNDS = {
        "office_modern": {"name": "Modern Office", "description": "Clean office environment"},
        "studio_blue": {"name": "Professional Blue", "description": "Blue studio background"},
        "minimalist": {"name": "Minimalist White", "description": "Clean white background"}
    }

# Import two-stage processor 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': {'name': 'Standard Quality'},
        'balanced': {'name': 'Balanced'},
        'high': {'name': 'High Quality'}
    }
    print("UI Components: Two-stage processor not available")

class UIStateManager:
    """Manage UI state and provide user feedback"""
    
    def __init__(self):
        self.processing_active = False
        self.models_loaded = False
        self.last_processing_time = None
        self.processing_history = []
        
    def update_processing_state(self, active: bool):
        self.processing_active = active
        if not active and self.last_processing_time:
            duration = time.time() - self.last_processing_time
            self.processing_history.append({
                'timestamp': time.time(),
                'duration': duration
            })
        elif active:
            self.last_processing_time = time.time()
    
    def get_average_processing_time(self) -> float:
        if not self.processing_history:
            return 0
        recent_history = self.processing_history[-5:]  # Last 5 processing sessions
        return sum(h['duration'] for h in recent_history) / len(recent_history)

# Global UI state
ui_state = UIStateManager()

def create_interface():
    """Create the enhanced Gradio interface with better UX"""
    
    # Enhanced processing function with better error handling
    def enhanced_process_video(
        video_path, bg_method, custom_img, prof_choice, 
        use_two_stage, chroma_preset, quality_preset,
        progress: Optional[gr.Progress] = None
    ):
        """Enhanced video processing with comprehensive error handling and user feedback"""
        
        if not CORE_FUNCTIONS_AVAILABLE:
            return None, "Error: Core processing functions not available", "System Error: Please check installation"
        
        if not processor.models_loaded:
            return None, "Error: Models not loaded", "Please load models first using the 'Load Models' button"
        
        if not video_path:
            return None, "Error: No video uploaded", "Please upload a video file first"
        
        # Validate inputs
        if bg_method == "professional" and not prof_choice:
            return None, "Error: No background selected", "Please select a professional background"
        
        if bg_method == "upload" and not custom_img:
            return None, "Error: No custom background", "Please upload a custom background image"
        
        try:
            ui_state.update_processing_state(True)
            
            # Set quality preset in processor config
            if quality_preset and hasattr(processor, 'config'):
                processor.config.quality_preset = quality_preset
            
            def progress_callback(pct, desc):
                if progress:
                    progress(pct, desc)
                return desc
            
            # Determine background choice
            if bg_method == "professional":
                background_choice = prof_choice
                custom_background_path = None
            else:
                background_choice = "custom"
                custom_background_path = custom_img
            
            # Process video
            result_path, result_message = process_video_fixed(
                video_path=video_path,
                background_choice=background_choice,
                custom_background_path=custom_background_path,
                progress_callback=progress_callback,
                use_two_stage=bool(use_two_stage),
                chroma_preset=chroma_preset or "standard",
                preview_mask=False,
                preview_greenscreen=False
            )
            
            ui_state.update_processing_state(False)
            
            if result_path:
                # Enhanced success message
                avg_time = ui_state.get_average_processing_time()
                success_info = f"""
Processing Complete!

Results:
{result_message}

Performance:
- Average processing time: {avg_time:.1f}s
- Two-stage mode: {'Enabled' if use_two_stage else 'Disabled'}
- Quality preset: {quality_preset or 'Default'}

Tips:
- Try two-stage mode for better quality
- Use 'fast' preset for quicker processing
- Shorter videos process faster
"""
                return result_path, success_info, "Processing completed successfully!"
            else:
                return None, f"Processing failed: {result_message}", f"Error: {result_message}"
                
        except Exception as e:
            ui_state.update_processing_state(False)
            error_msg = f"Processing error: {str(e)}"
            print(f"UI Error: {error_msg}\n{traceback.format_exc()}")
            return None, error_msg, f"System Error: {error_msg}"
    
    # Enhanced model loading with better feedback
    def enhanced_load_models(progress: Optional[gr.Progress] = None):
        """Enhanced model loading with detailed feedback"""
        
        if not CORE_FUNCTIONS_AVAILABLE:
            return "Error: Core functions not available", "System Error: Installation incomplete"
        
        try:
            def progress_callback(pct, desc):
                if progress:
                    progress(pct, desc)
                return desc
            
            result = load_models_with_validation(progress_callback)
            
            if "SUCCESS" in result or "successful" in result.lower():
                ui_state.models_loaded = True
                enhanced_result = f"""
Models Loaded Successfully!

Status:
{result}

Ready for Processing:
- High-quality segmentation (SAM2)
- Professional mask refinement (MatAnyone)
- {'Two-stage green screen mode available' if TWO_STAGE_AVAILABLE else 'Single-stage processing only'}

Next Steps:
1. Upload your video
2. Choose background method
3. Click 'Process Video'
"""
                return enhanced_result, "Models loaded successfully! Ready to process videos."
            else:
                return result, f"Model loading failed: {result}"
                
        except Exception as e:
            error_msg = f"Model loading error: {str(e)}"
            print(f"UI Model Loading Error: {error_msg}\n{traceback.format_exc()}")
            return error_msg, error_msg
    
    # Enhanced status functions
    def get_enhanced_model_status():
        """Get enhanced model status with user-friendly formatting"""
        try:
            status = get_model_status()
            if isinstance(status, dict):
                formatted_status = {
                    "SAM2 Segmentation": "Ready" if status.get('sam2_available') else "Not Loaded",
                    "MatAnyone Refinement": "Ready" if status.get('matanyone_available') else "Not Loaded",
                    "Two-Stage Mode": "Available" if status.get('two_stage_available') else "Not Available",
                    "Device": status.get('device', 'Unknown'),
                    "Models Validated": "Yes" if status.get('models_loaded') else "No"
                }
                if 'memory_usage' in status and status['memory_usage']:
                    memory = status['memory_usage']
                    if 'gpu_percent' in memory:
                        formatted_status["GPU Memory"] = f"{memory['gpu_percent']:.1f}% used"
                return formatted_status
            else:
                return {"Status": str(status)}
        except Exception as e:
            return {"Error": f"Failed to get status: {e}"}
    
    def get_enhanced_cache_status():
        """Get enhanced cache status with detailed information"""
        try:
            status = get_cache_status()
            if isinstance(status, dict):
                return {
                    "Cache Status": "Active" if status.get('models_loaded') else "Inactive",
                    "Processing Mode": "Two-Stage" if status.get('two_stage_available') else "Single-Stage",
                    "Configuration": status.get('config', {}),
                    "System Device": status.get('device', 'Unknown')
                }
            else:
                return {"Cache": str(status)}
        except Exception as e:
            return {"Error": f"Failed to get cache info: {e}"}
    
    # Create the main interface
    with gr.Blocks(
        title="BackgroundFX Pro - Professional Video Background Replacement",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="gray",
            neutral_hue="slate"
        ),
        css="""
        .main-header { text-align: center; margin-bottom: 20px; }
        .status-box { background: #f8f9fa; padding: 15px; border-radius: 8px; margin: 10px 0; }
        .error-box { background: #fee; border-left: 4px solid #dc3545; padding: 15px; }
        .success-box { background: #efe; border-left: 4px solid #28a745; padding: 15px; }
        .feature-list { columns: 2; column-gap: 20px; }
        """
    ) as demo:
        
        # Header
        with gr.Row():
            gr.Markdown("""
            # BackgroundFX Pro - Video Background Replacement
            
            Professional-quality video background replacement using AI segmentation and advanced compositing techniques.
            """, elem_classes=["main-header"])
        
        # System status indicator
        with gr.Row():
            with gr.Column(scale=1):
                system_status = gr.HTML(f"""
                <div class="status-box">
                    <h4>System Status</h4>
                    <ul>
                        <li>Core Functions: {'Available' if CORE_FUNCTIONS_AVAILABLE else 'Not Available'}</li>
                        <li>Utilities: {'Available' if UTILITIES_AVAILABLE else 'Not Available'}</li>
                        <li>Two-Stage Mode: {'Available' if TWO_STAGE_AVAILABLE else 'Not Available'}</li>
                    </ul>
                </div>
                """)
        
        with gr.Row():
            # Left column - Input and controls
            with gr.Column(scale=1):
                gr.Markdown("### Step 1: Upload Your Video")
                video_input = gr.Video(
                    label="Upload your video (MP4, AVI, MOV supported)",
                    height=300
                )
                
                with gr.Accordion("Video Requirements", open=False):
                    gr.Markdown("""
                    **Supported Formats:** MP4, AVI, MOV, MKV
                    **Max Duration:** 5 minutes (300 seconds)
                    **Max Resolution:** 4096x4096
                    **Max File Size:** 2GB
                    
                    **Recommendations:**
                    - Use 1080p or lower for faster processing
                    - Shorter videos (10-30s) are ideal for testing
                    - Ensure good lighting and clear person visibility
                    """)
                
                gr.Markdown("### Step 2: Choose Background")
                background_method = gr.Radio(
                    choices=["professional", "upload"],
                    value="professional",
                    label="Background Method"
                )
                
                # Professional backgrounds
                with gr.Group(visible=True) as professional_group:
                    gr.Markdown("**Professional Background Presets**")
                    
                    if UTILITIES_AVAILABLE and PROFESSIONAL_BACKGROUNDS:
                        choices = [(f"{bg['name']} - {bg['description']}", key) 
                                 for key, bg in PROFESSIONAL_BACKGROUNDS.items()]
                        default_choice = list(PROFESSIONAL_BACKGROUNDS.keys())[0]
                    else:
                        choices = [("Modern Office - Clean office environment", "office_modern")]
                        default_choice = "office_modern"
                    
                    professional_choice = gr.Dropdown(
                        choices=choices,
                        value=default_choice,
                        label="Select Professional Background"
                    )
                
                # Custom upload - FIXED: Removed 'info' parameter
                with gr.Group(visible=False) as upload_group:
                    gr.Markdown("**Upload Custom Background**")
                    custom_background = gr.Image(
                        label="Upload background image",
                        type="filepath"
                    )
                    gr.Markdown("JPG, PNG supported. Will be resized to match video resolution.")
                
                # Background method visibility control
                def update_background_visibility(method):
                    return (
                        gr.update(visible=(method == "professional")),
                        gr.update(visible=(method == "upload"))
                    )
                
                background_method.change(
                    fn=update_background_visibility,
                    inputs=background_method,
                    outputs=[professional_group, upload_group]
                )
                
                gr.Markdown("### Step 3: Processing Options")
                
                with gr.Row():
                    quality_preset = gr.Dropdown(
                        choices=[
                            ("Fast - Quick processing", "fast"),
                            ("Balanced - Good quality/speed", "balanced"),
                            ("High - Best quality", "high")
                        ],
                        value="balanced",
                        label="Quality Preset"
                    )
                
                with gr.Accordion("Advanced Settings", open=False):
                    use_two_stage = gr.Checkbox(
                        label="Enable Two-Stage Processing",
                        value=False,
                        interactive=TWO_STAGE_AVAILABLE
                    )
                    gr.Markdown("Cinema-quality green screen mode (slower but much better quality)")
                    
                    if TWO_STAGE_AVAILABLE:
                        chroma_preset = gr.Dropdown(
                            choices=[
                                ("Standard - General use", "standard"),
                                ("Studio - Broadcast quality", "studio"),
                                ("Outdoor - Challenging lighting", "outdoor")
                            ],
                            value="standard",
                            label="Chroma Key Preset"
                        )
                    else:
                        chroma_preset = gr.Dropdown(
                            choices=[("Standard", "standard")],
                            value="standard",
                            label="Chroma Key Preset",
                            interactive=False
                        )
                
                gr.Markdown("### Step 4: Process")
                
                with gr.Row():
                    load_models_btn = gr.Button(
                        "Load Models", 
                        variant="secondary",
                        size="lg"
                    )
                    process_btn = gr.Button(
                        "Process Video", 
                        variant="primary", 
                        size="lg",
                        scale=2
                    )
                
                # Status and feedback
                status_text = gr.Textbox(
                    label="Status Updates", 
                    value="Ready - Click 'Load Models' to begin",
                    interactive=False,
                    lines=6,
                    max_lines=10
                )
                
                # System monitoring
                with gr.Accordion("System Monitoring", open=False):
                    with gr.Row():
                        model_status_btn = gr.Button("Check Models", variant="secondary")
                        cache_status_btn = gr.Button("Check Cache", variant="secondary")
                    
                    model_status_display = gr.JSON(label="Model Status", visible=False)
                    cache_status_display = gr.JSON(label="Cache Status", visible=False)
            
            # Right column - Output and results
            with gr.Column(scale=1):
                gr.Markdown("### Results")
                
                video_output = gr.Video(
                    label="Processed Video",
                    height=400
                )
                
                result_info = gr.Textbox(
                    label="Processing Information",
                    interactive=False,
                    lines=12,
                    max_lines=15,
                    placeholder="Processing results and statistics will appear here..."
                )
                
                debug_info = gr.Textbox(
                    label="Debug Information",
                    interactive=False,
                    lines=6,
                    max_lines=10,
                    placeholder="Debug and system information will appear here...",
                    visible=False
                )
                
                # Toggle debug visibility
                show_debug_btn = gr.Button("Show Debug Info", variant="secondary", size="sm")
                
                def toggle_debug_visibility(current_visibility):
                    new_visibility = not current_visibility
                    return (
                        gr.update(visible=new_visibility),
                        "Hide Debug Info" if new_visibility else "Show Debug Info"
                    )
                
                show_debug_btn.click(
                    fn=lambda: toggle_debug_visibility(False),  # Will be managed by state
                    outputs=[debug_info, show_debug_btn]
                )
        
        # Event handlers
        load_models_btn.click(
            fn=enhanced_load_models,
            outputs=[status_text, debug_info],
            show_progress=True
        )
        
        process_btn.click(
            fn=enhanced_process_video,
            inputs=[
                video_input,
                background_method,
                custom_background,
                professional_choice,
                use_two_stage,
                chroma_preset,
                quality_preset
            ],
            outputs=[video_output, result_info, debug_info],
            show_progress=True
        )
        
        model_status_btn.click(
            fn=get_enhanced_model_status,
            outputs=[model_status_display],
            show_progress=False
        ).then(
            fn=lambda: gr.update(visible=True),
            outputs=[model_status_display]
        )
        
        cache_status_btn.click(
            fn=get_enhanced_cache_status,
            outputs=[cache_status_display],
            show_progress=False
        ).then(
            fn=lambda: gr.update(visible=True),
            outputs=[cache_status_display]
        )
        
        # Information and help section
        with gr.Accordion("Help & Information", open=False):
            gr.Markdown(f"""
            ### How to Use
            
            1. **Load Models**: Click 'Load Models' and wait for completion (first-time setup)
            2. **Upload Video**: Choose a video file (MP4 recommended, under 5 minutes)
            3. **Select Background**: Use professional presets or upload your own image
            4. **Configure Quality**: Choose preset based on your speed/quality preference
            5. **Process**: Click 'Process Video' and wait for completion
            
            ### Processing Modes
            
            **Single-Stage (Default)**
            - Direct background replacement
            - Faster processing (2-5x speed)
            - Good quality for most use cases
            - Recommended for: Social media, quick edits, testing
            
            **Two-Stage (Premium)**
            - Green screen intermediate step
            - Cinema-quality edge compositing
            - Advanced chroma key algorithms
            - Recommended for: Professional content, broadcast, film
            
            ### Performance Tips
            
            - **Fast Processing**: Use 'fast' preset, disable two-stage mode
            - **Best Quality**: Use 'high' preset, enable two-stage mode
            - **GPU Memory**: Processing automatically manages memory and provides fallbacks
            - **Video Length**: Shorter videos (10-30s) process much faster
            
            ### Troubleshooting
            
            **Models Won't Load**
            - Check internet connection (models download from Hugging Face)
            - Wait for downloads to complete (may take several minutes first time)
            - Try restarting if stuck
            
            **Processing Fails**
            - Ensure video file is not corrupted
            - Try shorter clips first (under 30 seconds)
            - Check video format (MP4 works best)
            - Verify sufficient disk space
            
            **Poor Quality Results**
            - Use higher quality preset
            - Enable two-stage mode
            - Ensure good lighting in original video
            - Try different professional backgrounds
            
            ### System Information
            
            - **Core Functions**: {'Available' if CORE_FUNCTIONS_AVAILABLE else 'Not Available'}
            - **Background Library**: {'Available' if UTILITIES_AVAILABLE else 'Not Available'}
            - **Two-Stage Processing**: {'Available' if TWO_STAGE_AVAILABLE else 'Not Available'}
            - **Professional Backgrounds**: {len(PROFESSIONAL_BACKGROUNDS)} presets available
            """)
    
    return demo

# Compatibility function for existing imports
def create_ui():
    """Compatibility wrapper for create_interface"""
    return create_interface()