Update ui/ui_components.py

ui/ui_components.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 """
 UI Components for BackgroundFX Pro
-- Layout only (no heavy logic here)
-- Delegates callbacks to ui/callbacks.py
-- Adds previews for uploaded video (first frame) and custom background image
+----------------------------------
+* Pure layout: no heavy logic
+* All callbacks live in ui/callbacks.py
+* Adds live previews for the first frame of the uploaded video
+  and for the custom background image
 """
 
 from __future__ import annotations
@@ -17,8 +19,8 @@
     cb_clear,
     cb_generate_bg,
     cb_use_gen_bg,
-    cb_video_changed,          # should return a preview image (or None)
-    cb_custom_bg_preview,      # should return a preview image (or None)
+    cb_video_changed,
+    cb_custom_bg_preview,
 )
 
 CSS = """
@@ -34,7 +36,7 @@
 .preview-img { border-radius: var(--radius); border: 1px solid rgba(0,0,0,.08); }
 """
 
-# NOTE: Keep these keys in sync with utils/cv_processing.PROFESSIONAL_BACKGROUNDS
+# Keep in sync with utils/cv_processing.PROFESSIONAL_BACKGROUNDS
 _BG_CHOICES = [
     "minimalist",
     "office_modern",
@@ -44,6 +46,7 @@
     "tech_dark",
 ]
 
+
 def create_interface() -> gr.Blocks:
     with gr.Blocks(
         title="🎬 BackgroundFX Pro",
@@ -52,21 +55,28 @@ def create_interface() -> gr.Blocks:
         theme=gr.themes.Soft()
     ) as demo:
 
-        # ---------- HERO ----------
+        # ------------------------------------------------------------------
+        # HERO
+        # ------------------------------------------------------------------
         with gr.Row(elem_id="hero"):
             gr.Markdown(
-                "## 🎬 BackgroundFX Pro\n"
-                "Polished matting & background replacement for video. Runs on Hugging Face Spaces.\n"
-                "Tip: **Load models** before processing for best results."
+                "## 🎬 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._"
             )
 
-        # ---------- QUICK START ----------
+        # ------------------------------------------------------------------
+        # TAB – Quick Start
+        # ------------------------------------------------------------------
         with gr.Tab("🏁 Quick Start"):
+
             with gr.Row():
-                # Left column: inputs & controls
+                # ── Left column ────────────────────────────────────────────
                 with gr.Column(scale=1):
                     video = gr.Video(label="Upload Video", interactive=True)
-                    # Small live preview (first frame)
+
+                    # First-frame live preview
                     video_preview = gr.Image(
                         label="Video Preview (first frame)",
                         interactive=False,
@@ -77,93 +87,133 @@ def create_interface() -> gr.Blocks:
                         label="Background Style",
                         choices=_BG_CHOICES,
                         value="minimalist",
-                        info="Matches built-in PROFESSIONAL_BACKGROUNDS keys."
+                        info="Matches keys in PROFESSIONAL_BACKGROUNDS"
                     )
+
                     custom_bg = gr.File(
-                        label="Custom Background (Optional)",
+                        label="Custom Background (optional)",
                         file_types=["image"],
                         interactive=True
                     )
-                    # Custom background preview
+
                     custom_bg_preview_img = gr.Image(
                         label="Custom Background Preview",
                         interactive=False,
                         elem_classes=["preview-img"]
                     )
 
+                    # ── Advanced options accordion ───────────────────────
                     with gr.Accordion("Advanced", open=False):
-                        use_two_stage = gr.Checkbox(label="Use Two-Stage Pipeline", value=False)
+                        use_two_stage = gr.Checkbox(
+                            label="Use Two-Stage Pipeline",
+                            value=False
+                        )
+
                         chroma_preset = gr.Dropdown(
                             label="Chroma Preset",
-                            choices=["standard"],
+                            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 (no audio remix)",
+                            label="Preview Mask only (mute audio)",
                             value=False
                         )
+
                         preview_greenscreen = gr.Checkbox(
-                            label="Preview Greenscreen (no audio remix)",
+                            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")
+                        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: outputs & status
+                # ── 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")
+                        btn_clear   = gr.Button("🧹 Clear",           variant="secondary")
 
-        # ---------- AI BACKGROUND ----------
+        # ------------------------------------------------------------------
+        # TAB – AI background (optional lightweight generator)
+        # ------------------------------------------------------------------
         with gr.Tab("🧠 AI Background (Lightweight)"):
+
             with gr.Row():
                 with gr.Column(scale=1):
                     prompt = gr.Textbox(
                         label="Describe vibe",
                         value="modern office",
-                        info="Examples: 'modern office', 'soft sunset studio', 'cool tech dark', 'forest ambience'"
+                        info="e.g. 'soft sunset studio', 'cool tech dark', 'forest ambience'"
                     )
+
                     with gr.Row():
-                        gen_width = gr.Slider(640, 1920, value=1280, step=10, label="Width")
-                        gen_height = gr.Slider(360, 1080, value=720, step=10, label="Height")
+                        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, value=8, step=1, label="Bokeh Blur")
-                        vignette = gr.Slider(0.0, 0.6, value=0.15, step=0.01, label="Vignette")
-                        contrast = gr.Slider(0.8, 1.4, value=1.05, step=0.01, label="Contrast")
+                        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")
 
                     btn_gen_bg = gr.Button("✨ Generate Background", variant="primary")
 
                 with gr.Column(scale=1):
-                    gen_preview = gr.Image(label="Generated Background", interactive=False, elem_classes=["preview-img"])
+                    gen_preview = gr.Image(
+                        label="Generated Background",
+                        interactive=False,
+                        elem_classes=["preview-img"]
+                    )
                     gen_path = gr.Textbox(label="Saved Path", interactive=False)
-                    use_gen_as_custom = gr.Button("📌 Use As Custom Background", variant="secondary")
+                    use_gen_as_custom = gr.Button("📌 Use as Custom Background", variant="secondary")
 
-        # ---------- STATUS ----------
+        # ------------------------------------------------------------------
+        # 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 logs for details.</div>"
+                "Check the runtime log for details.</div>"
             )
 
-        # ---------- CALLBACK WIRING ----------
-        # Core actions
+        # ------------------------------------------------------------------
+        # Callback wiring
+        # ------------------------------------------------------------------
         btn_load.click(cb_load_models, outputs=statusbox)
+
         btn_run.click(
             cb_process_video,
-            inputs=[video, bg_style, custom_bg, use_two_stage, chroma_preset, preview_mask, preview_greenscreen],
+            inputs=[
+                video,
+                bg_style,
+                custom_bg,
+                use_two_stage,
+                chroma_preset,
+                key_color_mode,          # NEW
+                preview_mask,
+                preview_greenscreen,
+            ],
             outputs=[out_video, statusbox],
         )
+
         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])
@@ -177,9 +227,6 @@ def create_interface() -> gr.Blocks:
         use_gen_as_custom.click(cb_use_gen_bg, inputs=[gen_path], outputs=[custom_bg])
 
         # Live previews
-        # Expectation:
-        #   - cb_video_changed(video_path) -> PIL.Image | numpy array | None
-        #   - cb_custom_bg_preview(file_obj) -> PIL.Image | numpy array | None
         video.change(cb_video_changed, inputs=[video], outputs=[video_preview])
         custom_bg.change(cb_custom_bg_preview, inputs=[custom_bg], outputs=[custom_bg_preview_img])