Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -1,346 +1,677 @@
 #!/usr/bin/env python3
 """
-UI Components for BackgroundFX Pro (forced Two-Stage)
------------------------------------------------------
-* Pure layout (tiny wrapper to set env for quality)
-* All heavy logic stays in ui/callbacks.py
-* Two-stage mode is always active (checkbox removed)
+Two-Stage Green-Screen Processing System  ✅ 2025-08-29
+Stage 1: Original → keyed background (auto-selected colour)
+Stage 2: Keyed video → final composite  (hybrid chroma + segmentation rescue)
+
+Aligned with current project layout:
+  * uses helpers from utils.cv_processing (segment_person_hq, refine_mask_hq)
+  * safe local create_video_writer (no core.app dependency)
+  * cancel support via stop_event
+  * progress_callback(pct, desc)
+  * fully self-contained – just drop in and import TwoStageProcessor
+
+Additional safety:
+  * Ensures MatAnyone receives a valid first-frame mask (bootstraps the session
+    with the first SAM2 mask). This prevents "First frame arrived without a mask"
+    warnings and shape mismatches inside the stateful refiner.
+
+Quality profiles (set via env BFX_QUALITY = speed | balanced | max):
+  * refine cadence, spill suppression, edge softness
+  * hybrid matte mix (segmentation vs chroma), small dilate/blur on mask
+  * optional tiny background blur to hide seams on very flat backgrounds
 """
-
 from __future__ import annotations
-import os
-import gradio as gr
-
-from ui.callbacks import (
-    cb_load_models,
-    cb_process_video,
-    cb_cancel,
-    cb_status,
-    cb_clear,
-    cb_generate_bg,
-    cb_use_gen_bg,
-    cb_preset_bg_preview,
-)
-
-# Typography & UI polish: sharper text + cleaner cards
-CSS = """
-:root {
-  --radius: 16px;
-  --font-sans: 'Inter', system-ui, -apple-system, 'Segoe UI', Roboto,
-               'Helvetica Neue', Arial, sans-serif;
-}
-
-/* Global crisp text */
-html, body, .gradio-container, .gradio-container * {
-  font-family: var(--font-sans) !important;
-  -webkit-font-smoothing: antialiased !important;
-  -moz-osx-font-smoothing: grayscale !important;
-  text-rendering: optimizeLegibility !important;
-  font-synthesis-weight: none;
-}
 
-/* Headings tighter & bolder */
-.gradio-container h1, .gradio-container h2, .gradio-container h3 {
-  letter-spacing: -0.01em;
-  font-weight: 700;
+import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
+from pathlib import Path
+from typing import Optional, Dict, Any, Callable, Tuple, List
+
+from utils.cv_processing import segment_person_hq, refine_mask_hq
+
+# Project logger if available
+try:
+    from utils.logger import get_logger
+    logger = get_logger(__name__)
+except Exception:
+    logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Local video-writer helper
+# ---------------------------------------------------------------------------
+def create_video_writer(output_path: str, fps: float, width: int, height: int, prefer_mp4: bool = True):
+    try:
+        ext = ".mp4" if prefer_mp4 else ".avi"
+        if not output_path:
+            output_path = tempfile.mktemp(suffix=ext)
+        else:
+            base, curr_ext = os.path.splitext(output_path)
+            if curr_ext.lower() not in [".mp4", ".avi", ".mov", ".mkv"]:
+                output_path = base + ext
+
+        fourcc = cv2.VideoWriter_fourcc(*("mp4v" if prefer_mp4 else "XVID"))
+        writer = cv2.VideoWriter(output_path, fourcc, float(fps), (int(width), int(height)))
+        if writer is None or not writer.isOpened():
+            alt_ext = ".avi" if prefer_mp4 else ".mp4"
+            alt_fourcc = cv2.VideoWriter_fourcc(*("XVID" if prefer_mp4 else "mp4v"))
+            alt_path = os.path.splitext(output_path)[0] + alt_ext
+            writer = cv2.VideoWriter(alt_path, alt_fourcc, float(fps), (int(width), int(height)))
+            if writer is None or not writer.isOpened():
+                return None, output_path
+            return writer, alt_path
+        return writer, output_path
+    except Exception as e:
+        logger.error(f"create_video_writer failed: {e}")
+        return None, output_path
+
+# ---------------------------------------------------------------------------
+# Key-colour helpers (fast, no external deps)
+# ---------------------------------------------------------------------------
+def _bgr_to_hsv_hue_deg(bgr: np.ndarray) -> np.ndarray:
+    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
+    # OpenCV H is 0-180; scale to degrees 0-360
+    return hsv[..., 0].astype(np.float32) * 2.0
+
+
+def _hue_distance(a_deg: float, b_deg: float) -> float:
+    """Circular distance on the hue wheel (degrees)."""
+    d = abs(a_deg - b_deg) % 360.0
+    return min(d, 360.0 - d)
+
+
+def _key_candidates_bgr() -> dict:
+    return {
+        "green":   {"bgr": np.array([  0,255,  0], dtype=np.uint8), "hue": 120.0},
+        "blue":    {"bgr": np.array([255,  0,  0], dtype=np.uint8), "hue": 240.0},
+        "cyan":    {"bgr": np.array([255,255,  0], dtype=np.uint8), "hue": 180.0},
+        "magenta": {"bgr": np.array([255,  0,255], dtype=np.uint8), "hue": 300.0},
+    }
+
+
+def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dict:
+    """Pick the candidate colour farthest from the actor's dominant hues."""
+    try:
+        fg = frame_bgr[mask_uint8 > 127]
+        if fg.size < 1_000:
+            return _key_candidates_bgr()["green"]
+
+        fg_hue = _bgr_to_hsv_hue_deg(fg.reshape(-1, 1, 3)).reshape(-1)
+        hist, edges = np.histogram(fg_hue, bins=36, range=(0.0, 360.0))
+        top_idx = np.argsort(hist)[-3:]
+        top_hues = [(edges[i] + edges[i+1]) * 0.5 for i in top_idx]
+
+        best_name, best_score = None, -1.0
+        for name, info in _key_candidates_bgr().items():
+            cand_hue = info["hue"]
+            score = min(abs((cand_hue - th + 180) % 360 - 180) for th in top_hues)
+            if score > best_score:
+                best_name, best_score = name, score
+        return _key_candidates_bgr().get(best_name, _key_candidates_bgr()["green"])
+    except Exception:
+        return _key_candidates_bgr()["green"]
+
+
+# ---------------------------------------------------------------------------
+# Chroma presets
+# ---------------------------------------------------------------------------
+CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {
+    'standard': {'key_color': [0,255,0], 'tolerance': 38, 'edge_softness': 2, 'spill_suppression': 0.35},
+    'studio':   {'key_color': [0,255,0], 'tolerance': 30, 'edge_softness': 1, 'spill_suppression': 0.45},
+    'outdoor':  {'key_color': [0,255,0], 'tolerance': 50, 'edge_softness': 3, 'spill_suppression': 0.25},
 }
 
-/* Body copy slightly tighter */
-#hero .prose, .gr-markdown, .gr-text {
-  letter-spacing: -0.003em;
+# ---------------------------------------------------------------------------
+# Quality profiles (env: BFX_QUALITY = speed | balanced | max)
+# ---------------------------------------------------------------------------
+QUALITY_PROFILES: Dict[str, Dict[str, Any]] = {
+    "speed":    dict(refine_stride=4, spill=0.30, edge_softness=2, mix=0.60, dilate=0, blur=0, bg_sigma=0.0),
+    "balanced": dict(refine_stride=2, spill=0.40, edge_softness=2, mix=0.75, dilate=1, blur=1, bg_sigma=0.6),
+    "max":      dict(refine_stride=1, spill=0.45, edge_softness=3, mix=0.85, dilate=2, blur=2, bg_sigma=1.0),
 }
 
-/* Card look */
-.card {
-  border-radius: var(--radius);
-  border: 1px solid rgba(0,0,0,.08);
-  padding: 16px;
-  background: linear-gradient(180deg, rgba(255,255,255,.94), rgba(248,250,252,.94));
-  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; }
-.preview-img { border-radius: var(--radius); border: 1px solid rgba(0,0,0,.08); }
-
-/* Buttons get a tiny weight bump for clarity */
-button, .gr-button { font-weight: 600; }
-
-/* Inline Quality select between buttons */
-.inline-quality .wrap-inner { min-width: 170px; }
-"""
-
-# Keep in sync with utils/cv_processing.PROFESSIONAL_BACKGROUNDS
-_BG_CHOICES = [
-    "minimalist",
-    "office_modern",
-    "studio_blue",
-    "studio_green",
-    "warm_gradient",
-    "tech_dark",
-]
-PRO_IMAGE_CHOICES = ["minimalist", "office_modern", "studio_blue", "studio_green"]
-GRADIENT_COLOR_CHOICES = ["warm_gradient", "tech_dark"]
-
-
-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 (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._"
+# ---------------------------------------------------------------------------
+# Two-Stage Processor
+# ---------------------------------------------------------------------------
+class TwoStageProcessor:
+    def __init__(self, sam2_predictor=None, matanyone_model=None):
+        self.sam2  = self._unwrap_sam2(sam2_predictor)
+        self.matanyone = matanyone_model
+        self.mask_cache_dir = Path("/tmp/mask_cache")
+        self.mask_cache_dir.mkdir(parents=True, exist_ok=True)
+
+        # Internal flags/state
+        self._mat_bootstrapped = False
+        self._alpha_prev: Optional[np.ndarray] = None  # temporal smoothing
+
+        # Quality selection at construction
+        qname = os.getenv("BFX_QUALITY", "balanced").strip().lower()
+        if qname not in QUALITY_PROFILES:
+            qname = "balanced"
+        self.quality = qname
+        self.q = QUALITY_PROFILES[qname]
+        logger.info(f"TwoStageProcessor quality='{self.quality}' ⇒ {self.q}")
+
+        logger.info(f"TwoStageProcessor init – SAM2: {self.sam2 is not None} | MatAnyOne: {self.matanyone is not None}")
+
+    # --------------------------- internal utils ---------------------------
+
+    def _unwrap_sam2(self, predictor):
+        """Unwrap the SAM2 predictor if needed."""
+        if predictor is None:
+            return None
+        if hasattr(predictor, 'sam_predictor'):
+            return predictor.sam_predictor
+        return predictor
+
+    def _refresh_quality_from_env(self):
+        """Pick up UI changes to BFX_QUALITY without rebuilding models."""
+        qname = os.getenv("BFX_QUALITY", self.quality).strip().lower()
+        if qname not in QUALITY_PROFILES:
+            qname = "balanced"
+        if qname != getattr(self, "quality", None) or not hasattr(self, "q"):
+            self.quality = qname
+            self.q = QUALITY_PROFILES[qname]
+            logger.info(f"Quality switched to '{self.quality}' ⇒ {self.q}")
+
+    def _get_mask(self, frame: np.ndarray) -> np.ndarray:
+        """Get segmentation mask using SAM2 (delegates to project helper)."""
+        if self.sam2 is None:
+            # Fallback: simple luminance threshold (kept to avoid breaking callers)
+            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+            _, mask = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
+            return mask
+
+        try:
+            mask = segment_person_hq(frame, self.sam2)
+            # segment_person_hq returns either uint8(0..255) or float(0..1) in most builds
+            return mask
+        except Exception as e:
+            logger.warning(f"SAM2 segmentation failed: {e}")
+            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+            _, mask = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
+            return mask
+
+    @staticmethod
+    def _to_binary_mask(mask: np.ndarray) -> Optional[np.ndarray]:
+        """Convert mask to uint8(0..255)."""
+        if mask is None:
+            return None
+        if mask.dtype == bool:
+            return mask.astype(np.uint8) * 255
+        if np.issubdtype(mask.dtype, np.floating):
+            m = np.clip(mask, 0.0, 1.0)
+            return (m * 255.0 + 0.5).astype(np.uint8)
+        return mask
+
+    @staticmethod
+    def _to_float01(mask: np.ndarray, h: int = None, w: int = None) -> Optional[np.ndarray]:
+        """Float [0,1] mask, optionally resized to (h,w)."""
+        if mask is None:
+            return None
+        m = mask.astype(np.float32)
+        if m.max() > 1.0:
+            m = m / 255.0
+        if h is not None and w is not None and (m.shape[0] != h or m.shape[1] != w):
+            m = cv2.resize(m, (w, h), interpolation=cv2.INTER_LINEAR)
+        return np.clip(m, 0.0, 1.0)
+
+    def _apply_greenscreen_hard(self, frame: np.ndarray, mask: np.ndarray, bg: np.ndarray) -> np.ndarray:
+        """Apply hard greenscreen compositing."""
+        mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) if mask.ndim == 2 else mask
+        mask_norm = mask_3ch.astype(np.float32) / 255.0
+        result = frame * mask_norm + bg * (1 - mask_norm)
+        return result.astype(np.uint8)
+
+    # -------- improved spill suppression (preserves luminance & skin) --------
+    def _suppress_green_spill(self, frame: np.ndarray, amount: float = 0.35) -> np.ndarray:
+        """
+        Desaturate green dominance while preserving luminance and red skin hues.
+        amount: 0..1
+        """
+        b, g, r = cv2.split(frame.astype(np.float32))
+        green_dom = (g > r) & (g > b)
+        avg_rb = (r + b) * 0.5
+        g2 = np.where(green_dom, g*(1.0-amount) + avg_rb*amount, g)
+        skin = (r > g + 12)  # protect skin tones
+        g2 = np.where(skin, g, g2)
+        out = cv2.merge([np.clip(b,0,255), np.clip(g2,0,255), np.clip(r,0,255)]).astype(np.uint8)
+        return out
+
+    # -------- edge-aware alpha refinement (guided-like) --------
+    def _refine_alpha_edges(self, frame_bgr: np.ndarray, alpha_u8: np.ndarray, radius: int = 3, iters: int = 1) -> np.ndarray:
+        """
+        Fast, dependency-free, guided-like refinement on the alpha border.
+        Returns: uint8 alpha
+        """
+        a = alpha_u8.astype(np.uint8)
+        if radius <= 0:
+            return a
+
+        band = cv2.Canny(a, 32, 64)
+        if band.max() == 0:
+            return a
+
+        for _ in range(max(1, iters)):
+            a_blur = cv2.GaussianBlur(a, (radius*2+1, radius*2+1), 0)
+            b,g,r = cv2.split(frame_bgr.astype(np.float32))
+            green_dom = (g > r) & (g > b)
+            spill_mask = (green_dom & (a > 96) & (a < 224)).astype(np.uint8)*255
+            u = cv2.bitwise_or(band, spill_mask)
+            a = np.where(u>0, a_blur, a).astype(np.uint8)
+
+        return a
+
+    # -------- soft key based on chosen color (robust to blue/cyan/magenta) --------
+    def _soft_key_mask(self, frame_bgr: np.ndarray, key_bgr: np.ndarray, tol: int = 40) -> np.ndarray:
+        """
+        Soft chroma mask (uint8 0..255, 255=keep subject) using CbCr distance.
+        """
+        if key_bgr is None:
+            return np.full(frame_bgr.shape[:2], 255, np.uint8)
+
+        ycbcr = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2YCrCb).astype(np.float32)
+        kycbcr = cv2.cvtColor(key_bgr.reshape(1,1,3).astype(np.uint8), cv2.COLOR_BGR2YCrCb).astype(np.float32)[0,0]
+        d = np.linalg.norm((ycbcr[...,1:] - kycbcr[1:]), axis=-1)
+        d = cv2.GaussianBlur(d, (5,5), 0)
+        alpha = 255.0 * np.clip((d - tol) / (tol*1.7), 0.0, 1.0)  # far from key = keep (255)
+        return alpha.astype(np.uint8)
+
+    # --------------------- NEW: MatAnyone bootstrap ----------------------
+    def _bootstrap_matanyone_if_needed(self, frame_bgr: np.ndarray, coarse_mask: np.ndarray):
+        """
+        Call the MatAnyone session ONCE with the first coarse mask to initialize
+        its memory. This guarantees downstream calls never hit "first frame without a mask".
+        """
+        if self.matanyone is None or self._mat_bootstrapped:
+            return
+        try:
+            h, w = frame_bgr.shape[:2]
+            mask_f = self._to_float01(coarse_mask, h, w)
+            rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+            _ = self.matanyone(rgb, mask_f)  # boot only; ignore returned alpha
+            self._mat_bootstrapped = True
+            logger.info("MatAnyone session bootstrapped with first-frame mask.")
+        except Exception as e:
+            logger.warning(f"MatAnyone bootstrap failed (continuing without): {e}")
+
+    # ---------------------------------------------------------------------
+    # Stage 1 – Original → keyed (green/blue/…)  -- chooses colour on 1st frame
+    # ---------------------------------------------------------------------
+    def stage1_extract_to_greenscreen(
+        self,
+        video_path: str,
+        output_path: str,
+        *,
+        key_color_mode: str = "auto",           # "auto" | "green" | "blue" | "cyan" | "magenta"
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+    ) -> Tuple[Optional[dict], str]:
+
+        def _prog(p, d):
+            if progress_callback:
+                try:
+                    progress_callback(float(p), str(d))
+                except Exception:
+                    pass
+
+        try:
+            # pick up any new quality selection
+            self._refresh_quality_from_env()
+
+            _prog(0.0, "Stage 1: opening video…")
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                return None, "Could not open input video"
+
+            fps   = cap.get(cv2.CAP_PROP_FPS) or 25.0
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
+            w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+            writer, out_path = create_video_writer(output_path, fps, w, h)
+            if writer is None:
+                cap.release()
+                return None, "Could not create output writer"
+
+            key_info: dict | None = None
+            chosen_bgr = np.array([0, 255, 0], np.uint8)  # default
+            probe_done = False
+            masks: List[np.ndarray] = []
+            frame_idx = 0
+
+            solid_bg = np.zeros((h, w, 3), np.uint8)  # overwritten per-frame
+
+            while True:
+                if stop_event and stop_event.is_set():
+                    _prog(1.0, "Stage 1: cancelled")
+                    break
+
+                ok, frame = cap.read()
+                if not ok:
+                    break
+
+                # --- SAM2 segmentation ---
+                mask = self._get_mask(frame)
+
+                # --- MatAnyone bootstrap exactly once (first frame) ---
+                if frame_idx == 0 and self.matanyone is not None:
+                    try:
+                        self._bootstrap_matanyone_if_needed(frame, mask)
+                    except Exception as e:
+                        logger.warning(f"Bootstrap error (non-fatal): {e}")
+
+                # --- Decide key colour once ---
+                if not probe_done:
+                    if key_color_mode.lower() == "auto":
+                        key_info = _choose_best_key_color(frame, self._to_binary_mask(mask))
+                        chosen_bgr = key_info["bgr"]
+                    else:
+                        cand = _key_candidates_bgr().get(key_color_mode.lower())
+                        if cand is not None:
+                            chosen_bgr = cand["bgr"]
+                    probe_done = True
+                    logger.info(f"[TwoStage] Using key colour: {key_color_mode} → {chosen_bgr.tolist()}")
+
+                # --- Optional refinement via MatAnyone (profile cadence) ---
+                stride = int(self.q.get("refine_stride", 3))
+                if self.matanyone and (frame_idx % max(1, stride) == 0):
+                    try:
+                        mask = refine_mask_hq(frame, mask, self.matanyone, fallback_enabled=True)
+                    except Exception as e:
+                        logger.warning(f"MatAnyOne refine fail f={frame_idx}: {e}")
+
+                # --- Composite onto solid key colour ---
+                solid_bg[:] = chosen_bgr
+                mask_u8 = self._to_binary_mask(mask)
+                gs = self._apply_greenscreen_hard(frame, mask_u8, solid_bg)
+                writer.write(gs)
+                masks.append(mask_u8)
+
+                frame_idx += 1
+                pct = 0.05 + 0.9 * (frame_idx / total) if total else min(0.95, 0.05 + frame_idx * 0.002)
+                _prog(pct, f"Stage 1: {frame_idx}/{total or '?'}")
+
+            cap.release()
+            writer.release()
+
+            # save mask cache
+            try:
+                cache_file = self.mask_cache_dir / (Path(out_path).stem + "_masks.pkl")
+                with open(cache_file, "wb") as f:
+                    pickle.dump(masks, f)
+            except Exception as e:
+                logger.warning(f"mask cache save fail: {e}")
+
+            _prog(1.0, "Stage 1: complete")
+            return (
+                {"path": out_path, "frames": frame_idx, "key_bgr": chosen_bgr.tolist()},
+                f"Green-screen video created ({frame_idx} frames)"
             )
 
-        # ------------------------------------------------------------------
-        # TAB – Quick Start
-        # ------------------------------------------------------------------
-        with gr.Tab("🏁 Quick Start"):
-
-            with gr.Row():
-                # ── Left column ────────────────────────────────────────────
-                with gr.Column(scale=1):
-                    video = gr.Video(label="Upload Video", interactive=True)
-
-                    # Hidden: effective preset key (still used by callbacks / defaults)
-                    bg_style = gr.Dropdown(
-                        label="Background Style (hidden)",
-                        choices=_BG_CHOICES,
-                        value="minimalist",
-                        visible=False,
+        except Exception as e:
+            logger.error(f"Stage 1 error: {e}\n{traceback.format_exc()}")
+            return None, f"Stage 1 failed: {e}"
+
+    # ---------------------------------------------------------------------
+    # Stage 2 – keyed video → final composite  (hybrid matte)
+    # ---------------------------------------------------------------------
+    def stage2_greenscreen_to_final(
+        self,
+        gs_path: str,
+        background: np.ndarray | str,
+        output_path: str,
+        *,
+        chroma_settings: Optional[Dict[str, Any]] = None,
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+        key_bgr: Optional[np.ndarray] = None,   # pass chosen key color
+    ) -> Tuple[Optional[str], str]:
+
+        def _prog(p, d):
+            if progress_callback:
+                try:
+                    progress_callback(float(p), str(d))
+                except Exception:
+                    pass
+
+        try:
+            # pick up any new quality selection
+            self._refresh_quality_from_env()
+
+            _prog(0.0, "Stage 2: opening keyed video…")
+            cap = cv2.VideoCapture(gs_path)
+            if not cap.isOpened():
+                return None, "Could not open keyed video"
+
+            fps   = cap.get(cv2.CAP_PROP_FPS) or 25.0
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
+            w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+            # Load or prepare background
+            if isinstance(background, str):
+                bg = cv2.imread(background)
+                if bg is None:
+                    cap.release()
+                    return None, "Could not load background image"
+                bg = cv2.resize(bg, (w, h))
+            else:
+                bg = cv2.resize(background, (w, h))
+
+            # Optional tiny BG blur per profile to hide seams on flat BGs
+            sigma = float(self.q.get("bg_sigma", 0.0))
+            if sigma > 0:
+                bg = cv2.GaussianBlur(bg, (0, 0), sigmaX=sigma, sigmaY=sigma)
+
+            writer, out_path = create_video_writer(output_path, fps, w, h)
+            if writer is None:
+                cap.release()
+                return None, "Could not create output writer"
+
+            # Load cached masks if available
+            masks = None
+            try:
+                cache_file = self.mask_cache_dir / (Path(gs_path).stem + "_masks.pkl")
+                if cache_file.exists():
+                    with open(cache_file, "rb") as f:
+                        masks = pickle.load(f)
+                    logger.info(f"Loaded {len(masks)} cached masks")
+            except Exception as e:
+                logger.warning(f"Could not load mask cache: {e}")
+
+            # Get chroma settings and override with profile
+            settings = chroma_settings or CHROMA_PRESETS.get('standard', {})
+            tolerance = int(settings.get('tolerance', 38))  # keep user tolerance
+            edge_softness = int(self.q.get('edge_softness', settings.get('edge_softness', 2)))
+            spill_suppression = float(self.q.get('spill', settings.get('spill_suppression', 0.35)))
+
+            # If caller didn't pass key_bgr, try preset or default green
+            if key_bgr is None:
+                key_bgr = np.array(settings.get('key_color', [0,255,0]), dtype=np.uint8)
+
+            self._alpha_prev = None  # reset temporal smoothing per render
+
+            frame_idx = 0
+            while True:
+                if stop_event and stop_event.is_set():
+                    _prog(1.0, "Stage 2: cancelled")
+                    break
+
+                ok, frame = cap.read()
+                if not ok:
+                    break
+
+                # Apply chroma keying with optional mask assistance
+                if masks and frame_idx < len(masks):
+                    mask = masks[frame_idx]
+                    final_frame = self._hybrid_composite(
+                        frame, bg, mask,
+                        tolerance=tolerance,
+                        edge_softness=edge_softness,
+                        spill_suppression=spill_suppression,
+                        key_bgr=key_bgr
                     )
-
-                    # =======================
-                    # Background Source block
-                    # =======================
-                    gr.Markdown("### 🖼️ Background Source")
-
-                    bg_method = gr.Radio(
-                        label="Choose method",
-                        choices=[
-                            "Upload file",
-                            "Pre-loaded professional images",
-                            "Pre-loaded Gradients / Colors",
-                            "AI generated background",
-                        ],
-                        value="Upload file",
+                else:
+                    # Pure chroma key
+                    final_frame = self._chroma_key_composite(
+                        frame, bg,
+                        tolerance=tolerance,
+                        edge_softness=edge_softness,
+                        spill_suppression=spill_suppression,
+                        key_bgr=key_bgr
                     )
 
-                    # a) Upload file option
-                    with gr.Group(visible=True) as grp_upload:
-                        custom_bg = gr.Image(
-                            label="Upload Background Image",
-                            interactive=True,
-                            type="filepath",  # returns file path
-                            elem_classes=["preview-img"]
-                        )
-
-                    # b) Pre-loaded professional images
-                    with gr.Group(visible=False) as grp_pro_images:
-                        pro_image_dd = gr.Dropdown(
-                            label="Professional Images",
-                            choices=PRO_IMAGE_CHOICES,
-                            value=PRO_IMAGE_CHOICES[0],
-                            info="Pre-defined photo-like backgrounds",
-                        )
-                        gr.Markdown(
-                            "<span class='sm'>Selecting a preset updates the preview below.</span>"
-                        )
-
-                    # c) Pre-loaded gradients & full colors
-                    with gr.Group(visible=False) as grp_gradients:
-                        gradient_dd = gr.Dropdown(
-                            label="Gradients & Full Colors",
-                            choices=GRADIENT_COLOR_CHOICES,
-                            value=GRADIENT_COLOR_CHOICES[0],
-                            info="Clean gradients or solid color styles",
-                        )
-                        gr.Markdown(
-                            "<span class='sm'>Selecting a preset updates the preview below.</span>"
-                        )
-
-                    # d) AI-generated background (inline, lightweight)
-                    with gr.Group(visible=False) as grp_ai:
-                        prompt = gr.Textbox(
-                            label="Describe vibe",
-                            value="modern office",
-                            info="e.g. 'soft sunset studio', 'cool tech dark', 'forest ambience'"
-                        )
-                        with gr.Row():
-                            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, 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")
-                        with gr.Row():
-                            btn_gen_bg_inline = gr.Button("✨ Generate Background", variant="primary")
-                            use_gen_as_custom_inline = gr.Button("📌 Use as Custom Background", variant="secondary")
-                        gen_preview = gr.Image(
-                            label="Generated Background",
-                            interactive=False,
-                            elem_classes=["preview-img"]
-                        )
-                        gen_path = gr.Textbox(label="Saved Path", interactive=False)
-
-                    # ── Advanced options accordion ───────────────────────
-                    with gr.Accordion("Advanced", open=False):
-                        chroma_preset = gr.Dropdown(
-                            label="Chroma Preset",
-                            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 only (mute audio)",
-                            value=False
-                        )
-                        preview_greenscreen = gr.Checkbox(
-                            label="Preview Green-screen only (mute audio)",
-                            value=False
-                        )
-
-                    # ── Controls row: Load → Quality → Process / Cancel ──
-                    with gr.Row():
-                        btn_load = gr.Button("🔄 Load Models", variant="secondary")
-
-                        quality = gr.Dropdown(
-                            label="Quality",
-                            choices=["speed", "balanced", "max"],
-                            value=os.getenv("BFX_QUALITY", "balanced"),
-                            info="Speed = fastest; Max = best edges & spill control.",
-                            elem_classes=["inline-quality"],
-                        )
-
-                        btn_run = gr.Button("🎬 Process Video", variant="primary")
-                        btn_cancel = gr.Button("⏹️ Cancel", variant="secondary")
-
-                # ── 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")
-
-        # ------------------------------------------------------------------
-        # 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 the runtime log for details.</div>"
-            )
-
-        # ------------------------------------------------------------------
-        # Callback wiring
-        # ------------------------------------------------------------------
-
-        # Toggle which background sub-section is visible
-        def _toggle_bg_sections(choice: str):
-            return (
-                gr.update(visible=(choice == "Upload file")),
-                gr.update(visible=(choice == "Pre-loaded professional images")),
-                gr.update(visible=(choice == "Pre-loaded Gradients / Colors")),
-                gr.update(visible=(choice == "AI generated background")),
-            )
-
-        bg_method.change(
-            _toggle_bg_sections,
-            inputs=[bg_method],
-            outputs=[grp_upload, grp_pro_images, grp_gradients, grp_ai],
+                writer.write(final_frame)
+                frame_idx += 1
+                pct = 0.05 + 0.9 * (frame_idx / total) if total else min(0.95, 0.05 + frame_idx * 0.002)
+                _prog(pct, f"Stage 2: {frame_idx}/{total or '?'}")
+
+            cap.release()
+            writer.release()
+
+            _prog(1.0, "Stage 2: complete")
+            return out_path, f"Final composite created ({frame_idx} frames)"
+
+        except Exception as e:
+            logger.error(f"Stage 2 error: {e}\n{traceback.format_exc()}")
+            return None, f"Stage 2 failed: {e}"
+
+    # ---------------- chroma + hybrid compositors (polished) ----------------
+    def _chroma_key_composite(self, frame, bg, *, tolerance=38, edge_softness=2, spill_suppression=0.35, key_bgr: Optional[np.ndarray] = None):
+        """Apply chroma key compositing with soft color distance + edge refinement."""
+        # 1) spill first
+        if spill_suppression > 0:
+            frame = self._suppress_green_spill(frame, spill_suppression)
+
+        # 2) build alpha
+        if key_bgr is not None:
+            alpha = self._soft_key_mask(frame, key_bgr, tol=int(tolerance))
+        else:
+            # Fallback: HSV green range
+            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+            lower_green = np.array([40, 40, 40])
+            upper_green = np.array([80, 255, 255])
+            alpha = cv2.bitwise_not(cv2.inRange(hsv, lower_green, upper_green))
+
+        # 3) soft edges + refinement
+        if edge_softness > 0:
+            k = edge_softness * 2 + 1
+            alpha = cv2.GaussianBlur(alpha, (k, k), 0)
+        alpha = self._refine_alpha_edges(frame, alpha, radius=max(1, edge_softness), iters=1)
+
+        # 4) temporal smoothing
+        if self._alpha_prev is not None and self._alpha_prev.shape == alpha.shape:
+            alpha = cv2.addWeighted(alpha, 0.75, self._alpha_prev, 0.25, 0)
+        self._alpha_prev = alpha
+
+        # 5) composite
+        mask_3ch = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR).astype(np.float32) / 255.0
+        out = frame.astype(np.float32) * mask_3ch + bg.astype(np.float32) * (1.0 - mask_3ch)
+        return np.clip(out, 0, 255).astype(np.uint8)
+
+    def _hybrid_composite(self, frame, bg, mask, *, tolerance=38, edge_softness=2, spill_suppression=0.35, key_bgr: Optional[np.ndarray] = None):
+        """Apply hybrid compositing using both chroma key and cached mask, with profile controls."""
+        chroma_result = self._chroma_key_composite(
+            frame, bg,
+            tolerance=tolerance,
+            edge_softness=edge_softness,
+            spill_suppression=spill_suppression,
+            key_bgr=key_bgr
         )
-
-        # Load models
-        btn_load.click(cb_load_models, outputs=statusbox)
-
-        # Tiny wrapper to set env for quality before calling the existing callback
-        def _run_with_quality(video_pth, bg_style_val, custom_bg_pth,
-                              use_two_stage_state, chroma_p, key_mode,
-                              prev_mask, prev_gs, quality_val):
-            os.environ["BFX_QUALITY"] = (quality_val or "balanced")
-            return cb_process_video(
-                video_pth, bg_style_val, custom_bg_pth,
-                use_two_stage_state, chroma_p, key_mode, prev_mask, prev_gs
+        if mask is None:
+            return chroma_result
+
+        # profile-driven dilate/feather on cached mask to close pinholes + soften edges
+        m = mask
+        d = int(self.q.get("dilate", 0))
+        if d > 0:
+            k = 2*d + 1
+            se = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
+            m = cv2.dilate(m, se, iterations=1)
+        b = int(self.q.get("blur", 0))
+        if b > 0:
+            m = cv2.GaussianBlur(m, (2*b+1, 2*b+1), 0)
+
+        m3 = cv2.cvtColor(m, cv2.COLOR_GRAY2BGR) if m.ndim == 2 else m
+        m3f = (m3.astype(np.float32) / 255.0)
+
+        seg_comp = frame.astype(np.float32) * m3f + bg.astype(np.float32) * (1.0 - m3f)
+
+        mix = float(self.q.get("mix", 0.7))  # weight towards segmentation on "max"
+        out = chroma_result.astype(np.float32) * (1.0 - mix) + seg_comp * mix
+        return np.clip(out, 0, 255).astype(np.uint8)
+
+    # ---------------------------------------------------------------------
+    # Combined pipeline
+    # ---------------------------------------------------------------------
+    def process_full_pipeline(
+        self,
+        video_path: str,
+        background: np.ndarray | str,
+        output_path: str,
+        *,
+        key_color_mode: str = "auto",
+        chroma_settings: Optional[Dict[str, Any]] = None,
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+    ) -> Tuple[Optional[str], str]:
+        """Run both stages in sequence."""
+
+        def _combined_progress(pct, desc):
+            # Scale progress: Stage 1 is 0-50%, Stage 2 is 50-100%
+            if "Stage 1" in desc:
+                actual_pct = pct * 0.5
+            else:  # Stage 2
+                actual_pct = 0.5 + pct * 0.5
+
+            if progress_callback:
+                try:
+                    progress_callback(actual_pct, desc)
+                except Exception:
+                    pass
+
+        try:
+            # pick up any new quality selection once per run
+            self._refresh_quality_from_env()
+
+            # Reset per-video state
+            self._mat_bootstrapped = False
+            self._alpha_prev = None
+            if self.matanyone is not None and hasattr(self.matanyone, "reset"):
+                try:
+                    self.matanyone.reset()
+                except Exception:
+                    pass
+
+            # Stage 1
+            temp_gs_path = tempfile.mktemp(suffix="_greenscreen.mp4")
+            stage1_result, stage1_msg = self.stage1_extract_to_greenscreen(
+                video_path, temp_gs_path,
+                key_color_mode=key_color_mode,
+                progress_callback=_combined_progress,
+                stop_event=stop_event
+            )
+            if stage1_result is None:
+                return None, stage1_msg
+
+            # Stage 2 (pass through chosen key color)
+            key_bgr = np.array(stage1_result.get("key_bgr", [0,255,0]), dtype=np.uint8)
+            final_path, stage2_msg = self.stage2_greenscreen_to_final(
+                stage1_result["path"], background, output_path,
+                chroma_settings=chroma_settings,
+                progress_callback=_combined_progress,
+                stop_event=stop_event,
+                key_bgr=key_bgr,
             )
 
-        # Always two-stage: pass use_two_stage=True to callback via State
-        btn_run.click(
-            _run_with_quality,
-            inputs=[
-                video,
-                bg_style,
-                custom_bg,
-                gr.State(value=True),  # Always two-stage
-                chroma_preset,
-                key_color_mode,
-                preview_mask,
-                preview_greenscreen,
-                quality,               # <-- the new control
-            ],
-            outputs=[out_video, statusbox],
-        )
-
-        # Cancel / Status / Clear
-        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]
-        )
-
-        # Preloaded presets → update preview (write into custom_bg)
-        pro_image_dd.change(
-            cb_preset_bg_preview,
-            inputs=[pro_image_dd],
-            outputs=[custom_bg],
-        )
-        gradient_dd.change(
-            cb_preset_bg_preview,
-            inputs=[gradient_dd],
-            outputs=[custom_bg],
-        )
+            # Clean up temp file
+            try:
+                os.remove(temp_gs_path)
+            except Exception:
+                pass
 
-        # AI background generation (inline)
-        btn_gen_bg_inline.click(
-            cb_generate_bg,
-            inputs=[prompt, gen_width, gen_height, bokeh, vignette, contrast],
-            outputs=[gen_preview, gen_path],
-        )
-        use_gen_as_custom_inline.click(
-            cb_use_gen_bg,
-            inputs=[gen_path],
-            outputs=[custom_bg],
-        )
-
-        # Initialize with a default preset preview on load
-        demo.load(
-            cb_preset_bg_preview,
-            inputs=[bg_style],
-            outputs=[custom_bg]
-        )
+            return final_path, stage2_msg
 
-    return demo
+        except Exception as e:
+            logger.error(f"Full pipeline error: {e}\n{traceback.format_exc()}")
+            return None, f"Pipeline failed: {e}"