Update ui_components.py

ui_components.py

@@ -8,14 +8,16 @@
     * 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 io
 import time
-import math
 import random
 from pathlib import Path
 from typing import Optional, Tuple, Dict, Any, List
@@ -23,9 +25,9 @@
 import gradio as gr
 from PIL import Image, ImageFilter, ImageOps
 import numpy as np
+import cv2
 
-# 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.
+# 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,
@@ -43,24 +45,42 @@
 
 
 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)
+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
-# (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)],
@@ -77,29 +97,25 @@ def _palette_from_prompt(prompt: str) -> List[tuple]:
     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()
+        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:
-    """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]
@@ -120,17 +136,13 @@ def generate_ai_background(
     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))
+        img = img.filter(ImageFilter.GaussianBlur(radius=max(0, min(50, bokeh))))
 
-    # Vignette
     if vignette > 0:
         y, x = np.ogrid[:height, :width]
         cx, cy = width / 2, height / 2
@@ -144,12 +156,11 @@ def generate_ai_background(
             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
+        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")
@@ -183,22 +194,29 @@ def create_interface() -> gr.Blocks:
         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", "studio", "ocean", "forest", "sunset",
-                            "royal", "warm", "cool"
+                            "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"
-                        )
+                        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)
 
@@ -218,7 +236,10 @@ def create_interface() -> gr.Blocks:
         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")
+                    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")
@@ -259,10 +280,10 @@ def _cb_process(
             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"):
+                # Gradio passes {"name": "/tmp/...", "size": int, ...}
                 custom_path = custom_file["name"]
             return process_video_fixed(
                 video_path=vid,
@@ -292,7 +313,7 @@ def _cb_status() -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
         # Clear
         def _cb_clear():
-            return None, "", None, ""
+            return None, "", None, "", None
 
         # AI background generation
         def _cb_generate_bg(prompt_text: str, w: int, h: int, b: float, v: float, c: float):
@@ -301,9 +322,27 @@ def _cb_generate_bg(prompt_text: str, w: int, h: int, b: float, v: float, c: flo
 
         # 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
+            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)
@@ -314,17 +353,17 @@ def _cb_use_gen_bg(path_text: str):
         )
         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])
+        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],
-        )
+        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