ui/ui_components.py

#!/usr/bin/env python3
"""
UI Components for BackgroundFX Pro
----------------------------------
* Pure layout: no heavy logic
* All callbacks live in ui/callbacks.py
* Removed first-frame video preview (redundant)
* Improved background section with:
  a) Upload file option
  b) Pre-loaded professional images (dropdown)
  c) Pre-loaded gradients & full colors (dropdown)
  d) AI-generated background option (inline)
"""

from __future__ import annotations
import gradio as gr

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,
)

CSS = """
:root { --radius: 16px; }
.gradio-container { max-width: 1080px !important; margin: auto !important; }
#hero .prose { font-size: 15px; }
.card { border-radius: var(--radius); border: 1px solid rgba(0,0,0,.08); padding: 16px;
        background: linear-gradient(180deg, rgba(255,255,255,.9), rgba(248,250,252,.9));
        box-shadow: 0 10px 30px rgba(0,0,0,.06); }
.footer-note { opacity: 0.7; font-size: 12px; }
.sm { font-size: 13px; opacity: 0.85; }
#statusbox { min-height: 120px; }
.preview-img { border-radius: var(--radius); border: 1px solid rgba(0,0,0,.08); }
"""

# Keep in sync with utils/cv_processing.PROFESSIONAL_BACKGROUNDS
# (We split existing keys into two UI groups without inventing new keys)
_BG_CHOICES = [
    "minimalist",
    "office_modern",
    "studio_blue",
    "studio_green",
    "warm_gradient",
    "tech_dark",
]
PRO_IMAGE_CHOICES = ["minimalist", "office_modern", "studio_blue", "studio_green"]
GRADIENT_COLOR_CHOICES = ["warm_gradient", "tech_dark"]


def create_interface() -> gr.Blocks:
    with gr.Blocks(
        title="🎬 BackgroundFX Pro",
        css=CSS,
        analytics_enabled=False,
        theme=gr.themes.Soft()
    ) as demo:

        # ------------------------------------------------------------------
        # HERO
        # ------------------------------------------------------------------
        with gr.Row(elem_id="hero"):
            gr.Markdown(
                "## 🎬 BackgroundFX Pro (CSP-Safe)\n"
                "Replace your video background with cinema-quality AI matting. "
                "Built for Hugging Face Spaces CSP.\n\n"
                "_Tip: press **Load Models** once after the Space spins up._"
            )

        # ------------------------------------------------------------------
        # TAB – Quick Start
        # ------------------------------------------------------------------
        with gr.Tab("🏁 Quick Start"):

            with gr.Row():
                # ── Left column ────────────────────────────────────────────
                with gr.Column(scale=1):
                    video = gr.Video(label="Upload Video", interactive=True)

                    # Hidden: effective preset key (still used by callbacks / defaults)
                    bg_style = gr.Dropdown(
                        label="Background Style (hidden)",
                        choices=_BG_CHOICES,
                        value="minimalist",
                        visible=False,
                    )

                    # =======================
                    # Background Source block
                    # =======================
                    gr.Markdown("### 🖼️ Background Source")

                    bg_method = gr.Radio(
                        label="Choose method",
                        choices=[
                            "Upload file",
                            "Pre-loaded professional images",
                            "Pre-loaded Gradients / Colors",
                            "AI generated background",
                        ],
                        value="Upload file",
                    )

                    # a) Upload file option
                    with gr.Group(visible=True) as grp_upload:
                        custom_bg = gr.Image(
                            label="Upload Background Image",
                            interactive=True,
                            type="filepath",  # returns file path
                            elem_classes=["preview-img"]
                        )

                    # b) Pre-loaded professional images
                    with gr.Group(visible=False) as grp_pro_images:
                        pro_image_dd = gr.Dropdown(
                            label="Professional Images",
                            choices=PRO_IMAGE_CHOICES,
                            value=PRO_IMAGE_CHOICES[0],
                            info="Pre-defined photo-like backgrounds",
                        )
                        gr.Markdown(
                            "<span class='sm'>Selecting a preset updates the preview below.</span>"
                        )

                    # c) Pre-loaded gradients & full colors
                    with gr.Group(visible=False) as grp_gradients:
                        gradient_dd = gr.Dropdown(
                            label="Gradients & Full Colors",
                            choices=GRADIENT_COLOR_CHOICES,
                            value=GRADIENT_COLOR_CHOICES[0],
                            info="Clean gradients or solid color styles",
                        )
                        gr.Markdown(
                            "<span class='sm'>Selecting a preset updates the preview below.</span>"
                        )

                    # d) AI-generated background (inline, lightweight)
                    with gr.Group(visible=False) as grp_ai:
                        prompt = gr.Textbox(
                            label="Describe vibe",
                            value="modern office",
                            info="e.g. 'soft sunset studio', 'cool tech dark', 'forest ambience'"
                        )
                        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="Bokeh Blur")
                            vignette = gr.Slider(0, 0.6, 0.15, step=0.01, label="Vignette")
                            contrast = gr.Slider(0.8, 1.4, 1.05, step=0.01, label="Contrast")
                        with gr.Row():
                            btn_gen_bg_inline = gr.Button("✨ Generate Background", variant="primary")
                            use_gen_as_custom_inline = gr.Button("📌 Use as Custom Background", variant="secondary")
                        gen_preview = gr.Image(
                            label="Generated Background",
                            interactive=False,
                            elem_classes=["preview-img"]
                        )
                        gen_path = gr.Textbox(label="Saved Path", interactive=False)

                    # ── Advanced options accordion ───────────────────────
                    with gr.Accordion("Advanced", open=False):
                        use_two_stage = gr.Checkbox(
                            label="Use Two-Stage Pipeline",
                            value=False
                        )
                        chroma_preset = gr.Dropdown(
                            label="Chroma Preset",
                            choices=["standard"],    # can add 'studio', 'outdoor' later
                            value="standard"
                        )
                        key_color_mode = gr.Dropdown(
                            label="Key-Colour Mode",
                            choices=["auto", "green", "blue", "cyan", "magenta"],
                            value="auto",
                            info="Auto picks a colour far from your clothes; override if needed."
                        )
                        preview_mask = gr.Checkbox(
                            label="Preview Mask only (mute audio)",
                            value=False
                        )
                        preview_greenscreen = gr.Checkbox(
                            label="Preview Green-screen only (mute audio)",
                            value=False
                        )

                    with gr.Row():
                        btn_load   = gr.Button("🔄 Load Models",    variant="secondary")
                        btn_run    = gr.Button("🎬 Process Video",  variant="primary")
                        btn_cancel = gr.Button("⏹️ Cancel",         variant="secondary")

                # ── Right column ──────────────────────────────────────────
                with gr.Column(scale=1):
                    out_video = gr.Video(label="Processed Output", interactive=False)
                    statusbox = gr.Textbox(label="Status", lines=8, elem_id="statusbox")
                    with gr.Row():
                        btn_refresh = gr.Button("🔍 Refresh Status", variant="secondary")
                        btn_clear   = gr.Button("🧹 Clear",           variant="secondary")

        # ------------------------------------------------------------------
        # TAB – Status & settings
        # ------------------------------------------------------------------
        with gr.Tab("📈 Status & Settings"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes=["card"]):
                    model_status = gr.JSON(label="Model Status")
                with gr.Column(scale=1, elem_classes=["card"]):
                    cache_status = gr.JSON(label="Cache / System Status")

            gr.Markdown(
                "<div class='footer-note'>If models fail to load, fallbacks keep the UI responsive. "
                "Check the runtime log for details.</div>"
            )

        # ------------------------------------------------------------------
        # Callback wiring
        # ------------------------------------------------------------------

        # Toggle which background sub-section is visible
        def _toggle_bg_sections(choice: str):
            return (
                gr.update(visible=(choice == "Upload file")),
                gr.update(visible=(choice == "Pre-loaded professional images")),
                gr.update(visible=(choice == "Pre-loaded Gradients / Colors")),
                gr.update(visible=(choice == "AI generated background")),
            )

        bg_method.change(
            _toggle_bg_sections,
            inputs=[bg_method],
            outputs=[grp_upload, grp_pro_images, grp_gradients, grp_ai],
        )

        # Load models
        btn_load.click(cb_load_models, outputs=statusbox)

        # Process video
        btn_run.click(
            cb_process_video,
            inputs=[
                video,
                bg_style,        # kept for compatibility; ignored if custom_bg is set
                custom_bg,       # becomes the effective background in all methods
                use_two_stage,
                chroma_preset,
                key_color_mode,
                preview_mask,
                preview_greenscreen,
            ],
            outputs=[out_video, statusbox],
        )

        # Cancel / Status / Clear
        btn_cancel.click(cb_cancel, outputs=statusbox)
        btn_refresh.click(cb_status, outputs=[model_status, cache_status])

        btn_clear.click(
            cb_clear,
            outputs=[out_video, statusbox, gen_preview, gen_path, custom_bg]
        )

        # Preloaded presets → update preview (write into custom_bg)
        pro_image_dd.change(
            cb_preset_bg_preview,
            inputs=[pro_image_dd],
            outputs=[custom_bg],
        )
        gradient_dd.change(
            cb_preset_bg_preview,
            inputs=[gradient_dd],
            outputs=[custom_bg],
        )

        # AI background generation (inline)
        btn_gen_bg_inline.click(
            cb_generate_bg,
            inputs=[prompt, gen_width, gen_height, bokeh, vignette, contrast],
            outputs=[gen_preview, gen_path],
        )
        use_gen_as_custom_inline.click(
            cb_use_gen_bg,
            inputs=[gen_path],
            outputs=[custom_bg],
        )

        # Initialize with a default preset preview on load
        demo.load(
            cb_preset_bg_preview,
            inputs=[bg_style],
            outputs=[custom_bg]
        )

    return demo