ui_components.py

#!/usr/bin/env python3
"""
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)
- NEW:
    * Preview of uploaded custom background
    * Preview of the video's first frame when a video is uploaded
    * Background style keys aligned with utils.cv_processing.PROFESSIONAL_BACKGROUNDS
"""

from __future__ import annotations

import os
import time
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 cv2

# Import core wrappers (core/app.py only imports UI from inside main(), no 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:
    ts = int(time.time() * 1000)
    p = TMP_DIR / f"{stem}_{ts}.{ext}"
    img.save(p)
    return str(p)


def _pil_from_path(path: str) -> Optional[Image.Image]:
    try:
        return Image.open(path).convert("RGB")
    except Exception:
        return None


def _first_frame(path: str, max_side: int = 960) -> Optional[Image.Image]:
    """Extract the first frame of a video for preview."""
    try:
        cap = cv2.VideoCapture(path)
        ok, frame = cap.read()
        cap.release()
        if not ok or frame is None:
            return None
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w = frame.shape[:2]
        scale = min(1.0, max_side / max(h, w))
        if scale < 1.0:
            frame = cv2.resize(frame, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
        return Image.fromarray(frame)
    except Exception:
        return None


# --------------------------
# Lightweight "AI" background generator
# --------------------------

_PALETTES = {
    "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
    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:
    acc = np.zeros((h, w), dtype=np.float32)
    for o in range(octaves):
        scale = 2 ** o
        small = np.random.rand(h // scale + 1, w // scale + 1).astype(np.float32)
        small = Image.fromarray((small * 255).astype(np.uint8)).resize((w, h), Image.BILINEAR)
        arr = np.array(small).astype(np.float32) / 255.0
        acc += arr / (o + 1)
    acc = acc / max(1e-6, acc.max())
    return acc


def _blend_palette(noise: np.ndarray, palette: List[tuple]) -> Image.Image:
    h, w = noise.shape
    img = np.zeros((h, w, 3), dtype=np.float32)
    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]:
    palette = _palette_from_prompt(prompt)
    noise = _perlin_like_noise(height, width, octaves=4)
    img = _blend_palette(noise, palette)

    if bokeh > 0:
        img = img.filter(ImageFilter.GaussianBlur(radius=max(0, min(50, bokeh))))

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

    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) * float(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."
            )

        with gr.Tab("🏁 Quick Start"):
            with gr.Row():
                with gr.Column(scale=1):
                    # Inputs
                    video = gr.Video(label="Upload Video")
                    video_preview = gr.Image(label="Video First Frame (Preview)", interactive=False)

                    # Align keys with utils.cv_processing.PROFESSIONAL_BACKGROUNDS
                    bg_style = gr.Dropdown(
                        label="Background Style",
                        choices=[
                            "minimalist",
                            "office_modern",
                            "studio_blue",
                            "studio_green",
                            "warm_gradient",
                            "tech_dark",
                        ],
                        value="minimalist",
                    )
                    custom_bg = gr.File(label="Custom Background (Optional)", file_types=["image"])
                    custom_bg_preview = gr.Image(label="Custom Background Preview", interactive=False)

                    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():
                PROCESS_CANCELLED.clear()
            custom_path = None
            if isinstance(custom_file, dict) and custom_file.get("name"):
                # Gradio passes {"name": "/tmp/...", "size": int, ...}
                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, "", 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):
            return (
                {"name": path_text, "size": os.path.getsize(path_text)}
                if path_text and os.path.exists(path_text) else None
            )

        # Video change -> extract first frame
        def _cb_video_changed(vid_path: str):
            if not vid_path:
                return None
            img = _first_frame(vid_path)
            return img

        # Custom background change -> preview image
        def _cb_custom_bg_preview(file_obj: dict | None):
            try:
                if isinstance(file_obj, dict) and file_obj.get("name") and os.path.exists(file_obj["name"]):
                    pil = _pil_from_path(file_obj["name"])
                    return pil
            except Exception:
                pass
            return 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, custom_bg_preview])

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

        # Live previews
        video.change(_cb_video_changed, inputs=[video], outputs=[video_preview])
        custom_bg.change(_cb_custom_bg_preview, inputs=[custom_bg], outputs=[custom_bg_preview])

    return demo