Update ui_components.py

ui_components.py

@@ -1,231 +1,330 @@
 #!/usr/bin/env python3
 """
-UI Components for BackgroundFX Pro - with Start/Stop, frame counter, and live FPS
+UI Components for BackgroundFX Pro (Hugging Face Spaces, CSP-safe)
+
+- Clean, modern layout with tabs
+- Keeps existing functionality:
+    * Load models
+    * Process video (single-stage / two-stage switch, previews, etc.)
+    * Status panel
+- Adds lightweight "AI Background" generator (procedural, no heavy deps)
 """
 
-import gradio as gr
+from __future__ import annotations
+
+import os
+import io
 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!
+import math
+import random
+from pathlib import Path
+from typing import Optional, Tuple, Dict, Any, List
+
+import gradio as gr
+from PIL import Image, ImageFilter, ImageOps
+import numpy as np
+
+# Import core wrappers (these are defined in core/app.py)
+# NOTE: core/app.py imports ui_components only *inside* main(), so this won’t create a circular import.
+from core.app import (
+    load_models_with_validation,
+    process_video_fixed,
+    get_model_status,
+    get_cache_status,
+    PROCESS_CANCELLED,
+)
+
+# --------------------------
+# Helpers: file paths, io
+# --------------------------
+
+TMP_DIR = Path("/tmp/bgfx")
+TMP_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _save_pil(img: Image.Image, stem: str = "gen_bg", ext: str = "png") -> str:
+    """Save a PIL image into /tmp and return its path."""
+    ts = int(time.time() * 1000)
+    p = TMP_DIR / f"{stem}_{ts}.{ext}"
+    img.save(p)
+    return str(p)
+
+
+def _image_to_np(img: Image.Image) -> np.ndarray:
+    return np.array(img)
+
+
+# --------------------------
+# Lightweight "AI" background generator
+# (procedural: palette from prompt + Perlin-ish noise + bokeh blur)
+# --------------------------
+
+_PALETTES = {
+    # very light keyword mapping; expand anytime
+    "office": [(240, 245, 250), (210, 220, 230), (180, 190, 200)],
+    "studio": [(18, 18, 20), (32, 32, 36), (58, 60, 64)],
+    "sunset": [(255, 183, 77), (255, 138, 101), (244, 143, 177)],
+    "forest": [(46, 125, 50), (102, 187, 106), (165, 214, 167)],
+    "ocean": [(33, 150, 243), (3, 169, 244), (0, 188, 212)],
+    "minimal": [(245, 246, 248), (230, 232, 236), (214, 218, 224)],
+    "warm": [(255, 224, 178), (255, 204, 128), (255, 171, 145)],
+    "cool": [(197, 202, 233), (179, 229, 252), (178, 235, 242)],
+    "royal": [(63, 81, 181), (121, 134, 203), (159, 168, 218)],
+}
+
+def _palette_from_prompt(prompt: str) -> List[tuple]:
+    p = (prompt or "").lower()
+    for key, pal in _PALETTES.items():
+        if key in p:
+            return pal
+    # fallback: hash to palette
+    random.seed(hash(p) % (2**32 - 1))
+    return [tuple(random.randint(90, 200) for _ in range(3)) for _ in range(3)]
+
+
+def _perlin_like_noise(h: int, w: int, octaves: int = 4) -> np.ndarray:
+    """Fast fake-perlin using summed blurred noise."""
+    acc = np.zeros((h, w), dtype=np.float32)
+    for o in range(octaves):
+        scale = 2 ** o
+        noise = np.random.rand(h // scale + 1, w // scale + 1).astype(np.float32)
+        noise = Image.fromarray((noise * 255).astype(np.uint8)).resize((w, h), Image.BILINEAR)
+        noise = np.array(noise).astype(np.float32) / 255.0
+        acc += noise / (o + 1)
+    acc = acc / acc.max()
+    return acc
+
+
+def _blend_palette(noise: np.ndarray, palette: List[tuple]) -> Image.Image:
+    """Map grayscale noise to a 3-color gradient."""
+    h, w = noise.shape
+    img = np.zeros((h, w, 3), dtype=np.float32)
+    # Tri-color mapping
+    thresholds = [0.33, 0.66]
+    c0, c1, c2 = [np.array(c, dtype=np.float32) for c in palette]
+    mask0 = noise < thresholds[0]
+    mask1 = (noise >= thresholds[0]) & (noise < thresholds[1])
+    mask2 = noise >= thresholds[1]
+    img[mask0] = c0
+    img[mask1] = c1
+    img[mask2] = c2
+    img = np.clip(img, 0, 255).astype(np.uint8)
+    return Image.fromarray(img)
+
+
+def generate_ai_background(
+    prompt: str,
+    width: int = 1280,
+    height: int = 720,
+    bokeh: float = 0.0,
+    vignette: float = 0.15,
+    contrast: float = 1.05,
+) -> Tuple[Image.Image, str]:
+    """Procedural 'AI-ish' background for fast, dependency-free generation."""
+    palette = _palette_from_prompt(prompt)
+    noise = _perlin_like_noise(height, width, octaves=4)
+    img = _blend_palette(noise, palette)
+
+    # Subtle blur / bokeh
+    if bokeh > 0:
+        radius = max(0, min(50, bokeh))
+        img = img.filter(ImageFilter.GaussianBlur(radius=radius))
+
+    # Vignette
+    if vignette > 0:
+        y, x = np.ogrid[:height, :width]
+        cx, cy = width / 2, height / 2
+        r = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
+        mask = 1 - np.clip(r / (max(width, height) / 1.2), 0, 1)
+        mask = mask ** 2
+        mask = (mask * (1 - vignette) + (1 - (1 - vignette))).astype(np.float32)
+        base = np.array(img).astype(np.float32) / 255.0
+        out = np.empty_like(base)
+        for c in range(3):
+            out[..., c] = base[..., c] * mask
+        img = Image.fromarray(np.clip(out * 255, 0, 255).astype(np.uint8))
+
+    # Simple contrast
+    if contrast != 1.0:
+        img = ImageOps.autocontrast(img, cutoff=1)
+        arr = np.array(img).astype(np.float32)
+        mean = arr.mean(axis=(0, 1), keepdims=True)
+        arr = (arr - mean) * contrast + mean
+        img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
+
+    path = _save_pil(img, stem="ai_bg", ext="png")
+    return img, path
+
+
+# --------------------------
+# Gradio UI
+# --------------------------
+
+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; }
+"""
+
+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\n"
+                "Polished matting & background replacement for video. Runs on Hugging Face Spaces.\n"
+                "Tip: **Load models** before processing for best results."
             )
-            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"))
+
+        with gr.Tab("🏁 Quick Start"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    video = gr.Video(label="Upload Video")
+                    bg_style = gr.Dropdown(
+                        label="Background Style",
+                        choices=[
+                            "minimalist", "office", "studio", "ocean", "forest", "sunset",
+                            "royal", "warm", "cool"
+                        ],
+                        value="minimalist",
                     )
-                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
+                    custom_bg = gr.File(label="Custom Background (Optional)", file_types=["image"])
+
+                    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"], value="standard"
+                        )
+                        preview_mask = gr.Checkbox(label="Preview Mask (no audio remix)", value=False)
+                        preview_greenscreen = gr.Checkbox(label="Preview Greenscreen (no audio remix)", 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")
+
+                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")
+
+        # ---------- AI BACKGROUND ----------
+        with gr.Tab("🧠 AI Background (Lightweight)"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    prompt = gr.Textbox(label="Describe the vibe (e.g., 'modern office', 'soft sunset studio')", value="modern office")
+                    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")
+                    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")
+
+                    btn_gen_bg = gr.Button("✨ Generate Background", variant="primary")
+
+                with gr.Column(scale=1):
+                    gen_preview = gr.Image(label="Generated Background", interactive=False)
+                    gen_path = gr.Textbox(label="Saved Path", interactive=False)
+                    use_gen_as_custom = gr.Button("📌 Use As Custom Background", variant="secondary")
+
+        # ---------- STATUS ----------
+        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>")
+
+        # ---------- CALLBACKS ----------
+        # Load Models
+        def _cb_load_models() -> str:
+            return load_models_with_validation()
+
+        # Process
+        def _cb_process(
+            vid: str,
+            style: str,
+            custom_file: dict | None,
+            use_two: bool,
+            chroma: str,
+            prev_mask: bool,
+            prev_green: bool,
+        ):
+            if PROCESS_CANCELLED.is_set():
+                # if user cancelled previously, reset it so a new run can proceed
+                PROCESS_CANCELLED.clear()
+            custom_path = None
+            if isinstance(custom_file, dict) and custom_file.get("name"):
+                custom_path = custom_file["name"]
+            return process_video_fixed(
+                video_path=vid,
+                background_choice=style,
+                custom_background_path=custom_path,
+                progress_callback=None,
+                use_two_stage=use_two,
+                chroma_preset=chroma,
+                preview_mask=prev_mask,
+                preview_greenscreen=prev_green,
+            )
+
+        # Cancel processing
+        def _cb_cancel() -> str:
+            try:
+                PROCESS_CANCELLED.set()
+                return "Cancellation requested."
+            except Exception as e:
+                return f"Cancel failed: {e}"
+
+        # Refresh status
+        def _cb_status() -> Tuple[Dict[str, Any], Dict[str, Any]]:
+            try:
+                return get_model_status(), get_cache_status()
+            except Exception as e:
+                return {"error": str(e)}, {"error": str(e)}
+
+        # Clear
+        def _cb_clear():
+            return None, "", None, ""
+
+        # AI background generation
+        def _cb_generate_bg(prompt_text: str, w: int, h: int, b: float, v: float, c: float):
+            img, path = generate_ai_background(prompt_text, width=int(w), height=int(h), bokeh=b, vignette=v, contrast=c)
+            return img, path
+
+        # Use AI gen as custom
+        def _cb_use_gen_bg(path_text: str):
+            # The Quick Start tab expects a "file" object. We can simply echo the path
+            # and let the "Process" callback read it if provided.
+            return {"name": path_text, "size": os.path.getsize(path_text)} if path_text and os.path.exists(path_text) else None
+
+        # Wire events
+        btn_load.click(_cb_load_models, outputs=statusbox)
+        btn_run.click(
+            _cb_process,
+            inputs=[video, bg_style, custom_bg, use_two_stage, chroma_preset, 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])
 
-        # Fixed model loading button - let Gradio handle progress automatically
-        load_models_btn.click(
-            fn=handle_model_loading,
-            outputs=[status_text],
-            show_progress=True
+        btn_gen_bg.click(
+            _cb_generate_bg,
+            inputs=[prompt, gen_width, gen_height, bokeh, vignette, contrast],
+            outputs=[gen_preview, gen_path],
+        )
+        use_gen_as_custom.click(
+            _cb_use_gen_bg,
+            inputs=[gen_path],
+            outputs=[custom_bg],
         )
 
     return demo
-
-def create_ui():
-    return create_interface()