ui_components.py

#!/usr/bin/env python3
"""
UI Components for BackgroundFX Pro - with Start/Stop, frame counter, and live FPS
"""

import gradio as gr
import time
import threading

# --- Import backend logic ---
try:
    from core.app import (
        processor,
        load_models_with_validation,
        process_video_fixed,
        get_model_status,
        get_cache_status,
    )
    CORE_FUNCTIONS_AVAILABLE = True
except Exception as e:
    print(f"[UI] Core functions import failed: {e}")
    CORE_FUNCTIONS_AVAILABLE = False

try:
    from utils.backgrounds import PROFESSIONAL_BACKGROUNDS
    UTILITIES_AVAILABLE = True
except Exception as e:
    PROFESSIONAL_BACKGROUNDS = {"office_modern": {"name": "Modern Office", "description": "Default office background"}}
    UTILITIES_AVAILABLE = False

# --- UI state management ---
stop_event = threading.Event()

def create_interface():
    is_processing = gr.State(False)
    frame_progress = gr.State({"current": 0, "total": 0, "fps": 0.0})

    def enhanced_process_video(
        video_path, bg_method, custom_img, prof_choice, 
        use_two_stage, chroma_preset, quality_preset,
        is_processing_state, frame_progress_state, progress=gr.Progress()
    ):
        """
        Handles video processing (start/stop) and live frame/fps reporting.
        """
        # If already processing: user wants to STOP
        if is_processing_state:
            stop_event.set()
            return None, "Processing stopped by user.", "Processing stopped.", False, {"current": 0, "total": 0, "fps": 0.0}
        
        # Otherwise: START processing
        stop_event.clear()
        is_processing_state = True
        frame_progress_state = {"current": 0, "total": 0, "fps": 0.0}
        last_update = time.time()
        last_frame = 0

        if not CORE_FUNCTIONS_AVAILABLE:
            return None, "Core backend not available.", "System error.", False, frame_progress_state
        if not processor.models_loaded:
            return None, "Models not loaded.", "Load models first.", False, frame_progress_state
        if not video_path:
            return None, "No video uploaded.", "Upload a video.", False, frame_progress_state

        # Choose background
        if bg_method == "professional" and not prof_choice:
            return None, "No professional background selected.", "Choose a background.", False, frame_progress_state
        if bg_method == "upload" and not custom_img:
            return None, "No custom background image.", "Upload a background image.", False, frame_progress_state

        try:
            def progress_callback(pct, desc, current_frame=None, total_frames=None):
                # User cancellation check
                if stop_event.is_set():
                    raise Exception("Processing stopped by user.")

                # Standard Gradio progress
                if progress:
                    progress(pct, desc)

                # Update frame count/fps
                nonlocal last_update, last_frame, frame_progress_state
                now = time.time()
                fps = 0.0
                if current_frame is not None and total_frames is not None:
                    dt = now - last_update if last_update else 0.01
                    frames_done = current_frame - last_frame if last_frame is not None else 0
                    fps = frames_done / dt if dt > 0 else 0.0
                    last_update = now
                    last_frame = current_frame
                    frame_progress_state = {
                        "current": int(current_frame),
                        "total": int(total_frames),
                        "fps": round(fps, 2)
                    }
                return desc

            result_path, result_message = process_video_fixed(
                video_path=video_path,
                background_choice=prof_choice if bg_method == "professional" else "custom",
                custom_background_path=custom_img if bg_method == "upload" else None,
                progress_callback=lambda pct, desc: progress_callback(pct, desc),  # You may want to pass current_frame/total_frames
                use_two_stage=bool(use_two_stage),
                chroma_preset=chroma_preset or "standard",
                preview_mask=False,
                preview_greenscreen=False,
                stop_event=stop_event,   # <-- Backend must support this argument!
            )
            is_processing_state = False
            return result_path, result_message, "Processing completed.", False, frame_progress_state

        except Exception as e:
            is_processing_state = False
            return None, f"Error: {str(e)}", "Error during processing.", False, frame_progress_state

    def handle_model_loading(progress=gr.Progress()):
        """
        Wrapper function for model loading with proper progress handling
        """
        try:
            # Call the model loading function with progress callback
            result = load_models_with_validation(progress)
            return result
        except Exception as e:
            return f"Model loading failed: {str(e)}"

    # --- UI setup ---
    with gr.Blocks(
        title="BackgroundFX Pro",
        css="""
        .main-header { text-align: center; margin-bottom: 20px; }
        .status-box { background: #f8f9fa; padding: 15px; border-radius: 8px; margin: 10px 0; }
        .live-progress { font-size: 1.1em; font-weight: bold; color: #004085; }
        """
    ) as demo:
        with gr.Row():
            gr.Markdown("# BackgroundFX Pro - Video Background Replacement", elem_classes=["main-header"])

        # --- Main UI fields ---
        with gr.Row():
            with gr.Column():
                video_input = gr.Video(label="Upload Video", height=280)
                background_method = gr.Radio(
                    ["professional", "upload"], value="professional", label="Background Method"
                )
                professional_choice = gr.Dropdown(
                    choices=[(bg['name'], key) for key, bg in PROFESSIONAL_BACKGROUNDS.items()],
                    value=list(PROFESSIONAL_BACKGROUNDS.keys())[0],
                    label="Professional Background"
                )
                custom_background = gr.Image(label="Upload Background", type="filepath", visible=False)
                quality_preset = gr.Dropdown(
                    choices=[
                        ("Fast", "fast"),
                        ("Balanced", "balanced"),
                        ("High Quality", "high")
                    ],
                    value="balanced",
                    label="Quality Preset"
                )
                use_two_stage = gr.Checkbox("Enable Two-Stage Mode", value=False)
                chroma_preset = gr.Dropdown(
                    choices=[("Standard", "standard"), ("Studio", "studio"), ("Outdoor", "outdoor")],
                    value="standard",
                    label="Chroma Key Preset"
                )
                # Dynamic show/hide
                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_choice, custom_background]
                )
                with gr.Row():
                    load_models_btn = gr.Button("Load Models", variant="secondary")
                    process_btn = gr.Button("Process Video", variant="primary")

                status_text = gr.Textbox(label="Status", value="Ready", interactive=False)
                frame_info = gr.Textbox(label="Frame Progress", value="", interactive=False)
                fps_info = gr.Textbox(label="Frames/sec", value="", interactive=False)

            with gr.Column():
                video_output = gr.Video(label="Processed Video", height=360)
                result_info = gr.Textbox(label="Processing Info", lines=8, interactive=False)
                debug_info = gr.Textbox(label="Debug Log", lines=6, interactive=False, visible=False)

        # --- Button actions ---
        def update_process_btn(is_processing):
            return gr.update(value="Stop Processing" if is_processing else "Process Video")
        is_processing.change(update_process_btn, inputs=is_processing, outputs=process_btn)

        def update_frame_fields(frame_progress):
            total = frame_progress.get("total", 0)
            current = frame_progress.get("current", 0)
            fps = frame_progress.get("fps", 0.0)
            txt = f"{current}/{total}" if total else ""
            return txt, f"{fps:.2f} FPS" if fps else ""
        frame_progress.change(update_frame_fields, inputs=frame_progress, outputs=[frame_info, fps_info])

        process_btn.click(
            fn=enhanced_process_video,
            inputs=[
                video_input,
                background_method,
                custom_background,
                professional_choice,
                use_two_stage,
                chroma_preset,
                quality_preset,
                is_processing,
                frame_progress
            ],
            outputs=[video_output, result_info, debug_info, is_processing, frame_progress],
            show_progress=True
        )

        # Fixed model loading button - let Gradio handle progress automatically
        load_models_btn.click(
            fn=handle_model_loading,
            outputs=[status_text],
            show_progress=True
        )

    return demo

def create_ui():
    return create_interface()