Update ui/callbacks.py

ui/callbacks.py

@@ -1,8 +1,14 @@
 #!/usr/bin/env python3
 """
 Callbacks for BackgroundFX Pro UI
-- Pure functions wired to UI events
-- Safe imports & graceful fallbacks so the UI keeps working even if optional utils are missing
+---------------------------------
+All functions here are *thin* wrappers wired to the Gradio interface.
+
+Changes in this revision
+------------------------
+•   Added **key_color_mode** positional arg to `cb_process_video`
+    (matches the new dropdown in ui_components.py)
+•   Passes that value straight through to `process_video_fixed(...)`
 """
 
 from __future__ import annotations
@@ -25,127 +31,154 @@
 except Exception:
     pass
 
-def _generate_ai_background(prompt_text: str, width: int, height: int, bokeh: float, vignette: float, contrast: float):
+
+# ------------------------------------------------------------------
+# LIGHTWEIGHT BG GENERATOR (inline fallback)
+# ------------------------------------------------------------------
+def _generate_ai_background(
+    prompt_text: str,
+    width: int,
+    height: int,
+    bokeh: float,
+    vignette: float,
+    contrast: float,
+):
     """
-    Lightweight inline fallback if utils.bg_generator is not present yet.
-    Generates a simple procedural background.
+    If utils.bg_generator.generate_ai_background exists, use it.
+    Otherwise fall back to a tiny procedural background made with PIL & NumPy.
     """
     if _try_bg_gen is not None:
-        return _try_bg_gen(prompt_text, width=width, height=height, bokeh=bokeh, vignette=vignette, contrast=contrast)
-
-    # Fallback (minimal dependency, PIL-only)
+        return _try_bg_gen(
+            prompt_text,
+            width=width,
+            height=height,
+            bokeh=bokeh,
+            vignette=vignette,
+            contrast=contrast,
+        )
+
+    # -------- Tiny fallback (PIL only) --------
     from pathlib import Path
-    import time
-    import random
-    import numpy as np
+    import time, random, numpy as np
     from PIL import Image, ImageFilter, ImageOps
 
     TMP_DIR = Path("/tmp/bgfx")
     TMP_DIR.mkdir(parents=True, exist_ok=True)
 
     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)],
+        "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)],
     }
     p = (prompt_text or "").lower()
-    for k, pal in palettes.items():
-        if k in p:
-            palette = pal
-            break
-    else:
-        random.seed(hash(p) % (2**32 - 1))
+    palette = next((pal for k, pal in palettes.items() if k in p), None)
+    if palette is None:
+        random.seed(hash(p) & 0xFFFFFFFF)
         palette = [tuple(random.randint(90, 200) for _ in range(3)) for _ in range(3)]
 
-    # Fake Perlin-ish noise
     def _noise(h, w, octaves=4):
-        acc = np.zeros((h, w), dtype=np.float32)
+        acc = np.zeros((h, w), np.float32)
         for o in range(octaves):
-            scale = 2 ** o
-            small = np.random.rand(h // scale + 1, w // scale + 1).astype(np.float32)
-            img = Image.fromarray((small * 255).astype(np.uint8)).resize((w, h), Image.BILINEAR)
-            acc += np.array(img, dtype=np.float32) / 255.0 / (o + 1)
-        acc = acc / max(1e-6, acc.max())
+            s = 2**o
+            small = np.random.rand(h // s + 1, w // s + 1).astype(np.float32)
+            acc += cv2.resize(small, (w, h), interpolation=cv2.INTER_LINEAR) / (o + 1)
+        acc /= max(1e-6, acc.max())
         return acc
 
-    def _blend(noise, pal):
-        h, w = noise.shape
-        img = np.zeros((h, w, 3), dtype=np.float32)
+    def _blend(n, pal):
+        h, w = n.shape
         thr = [0.33, 0.66]
-        c0, c1, c2 = [np.array(c, dtype=np.float32) for c in pal]
-        img[noise < thr[0]] = c0
-        mid = (noise >= thr[0]) & (noise < thr[1])
+        img = np.zeros((h, w, 3), np.float32)
+        c0, c1, c2 = [np.array(c, np.float32) for c in pal]
+        img[n < thr[0]] = c0
+        mid = (n >= thr[0]) & (n < thr[1])
         img[mid] = c1
-        img[noise >= thr[1]] = c2
+        img[n >= thr[1]] = c2
         return Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
 
-    n = _noise(height, width, octaves=4)
+    n   = _noise(height, width, 4)
     out = _blend(n, palette)
 
     if bokeh > 0:
-        from PIL import ImageFilter
         out = out.filter(ImageFilter.GaussianBlur(radius=min(50, max(0, bokeh))))
-
     if vignette > 0:
-        import numpy as np
         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).astype(np.float32)
+        mask = (mask**2).astype(np.float32)
         base = np.array(out).astype(np.float32) / 255.0
-        out_arr = base * (mask[..., None] * (1 - vignette) + vignette)
-        out = Image.fromarray(np.clip(out_arr * 255, 0, 255).astype(np.uint8))
-
+        out  = Image.fromarray(np.clip(base * (mask[..., None] * (1 - vignette) + vignette) * 255, 0, 255).astype(np.uint8))
     if contrast != 1.0:
-        from PIL import ImageOps
         out = ImageOps.autocontrast(out, cutoff=1)
-        arr = (np.array(out).astype(np.float32))
+        arr = np.array(out).astype(np.float32)
         mean = arr.mean(axis=(0, 1), keepdims=True)
         arr = (arr - mean) * float(contrast) + mean
         out = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
 
-    ts = int(time.time() * 1000)
+    ts   = int(time.time() * 1000)
     path = str((TMP_DIR / f"ai_bg_{ts}.png").resolve())
     out.save(path)
     return out, path
 
 
-# -------- MODEL MANAGEMENT --------
+# ------------------------------------------------------------------
+# MODEL MANAGEMENT
+# ------------------------------------------------------------------
 def cb_load_models() -> str:
+    """Load SAM2 + MatAnyOne and return human-readable status."""
     return load_models_with_validation()
 
+
+# ------------------------------------------------------------------
+# MAIN video-processing callback
+# ------------------------------------------------------------------
 def cb_process_video(
     vid: str,
     style: str,
     custom_file: dict | None,
     use_two: bool,
     chroma: str,
+    key_color_mode: str,     # <-- NEW ARG from dropdown
     prev_mask: bool,
     prev_green: bool,
 ):
+    """
+    Runs the two-stage (or single-stage) pipeline and returns:
+        (processed_video_path | None, status_message:str)
+    """
+    # Reset any prior cancel flag when user clicks Run
     if PROCESS_CANCELLED.is_set():
         PROCESS_CANCELLED.clear()
+
+    # Resolve custom background path (if provided)
     custom_path = None
     if isinstance(custom_file, dict) and custom_file.get("name"):
         custom_path = custom_file["name"]
+
+    # Fire the core function
     return process_video_fixed(
         video_path=vid,
         background_choice=style,
         custom_background_path=custom_path,
-        progress_callback=None,
+        progress_callback=None,          # UI-level progress handled inside
         use_two_stage=use_two,
         chroma_preset=chroma,
+        key_color_mode=key_color_mode,   # <-- pass straight through
         preview_mask=prev_mask,
         preview_greenscreen=prev_green,
     )
 
+
+# ------------------------------------------------------------------
+# CANCEL / STATUS / CLEAR
+# ------------------------------------------------------------------
 def cb_cancel() -> str:
     try:
         PROCESS_CANCELLED.set()
@@ -160,24 +193,28 @@ def cb_status() -> Tuple[Dict[str, Any], Dict[str, Any]]:
         return {"error": str(e)}, {"error": str(e)}
 
 def cb_clear():
-    # out_video, status, gen_preview, gen_path
+    # Return blanks for (out_video, status, gen_preview, gen_path)
     return None, "", None, ""
 
 
-# -------- AI BACKGROUND --------
+# ------------------------------------------------------------------
+# AI BACKGROUND
+# ------------------------------------------------------------------
 def cb_generate_bg(prompt_text: str, w: int, h: int, b: float, v: float, c: float):
     img, path = _generate_ai_background(prompt_text, int(w), int(h), b, v, c)
     return img, path
 
 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
+    if path_text and os.path.exists(path_text):
+        return {"name": path_text, "size": os.path.getsize(path_text)}
+    return None
 
 
-# -------- PREVIEWS (no-op safe defaults) --------
+# ------------------------------------------------------------------
+# PREVIEWS (safe no-ops — extend later)
+# ------------------------------------------------------------------
 def cb_video_changed(vid_path: str):
-    # Reserved hook for future preview image extraction
     return None
 
 def cb_custom_bg_preview(file_obj: dict | None):
-    # Reserved hook to show uploaded custom background
     return None