Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Two-Stage Green-Screen Processing System  ✅ 2025-08-26
+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)
 
@@ -11,7 +11,7 @@
   * progress_callback(pct, desc)
   * fully self-contained – just drop in and import TwoStageProcessor
 
-Additional safety (2025-08-29):
+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.
@@ -31,6 +31,9 @@
 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"
@@ -56,6 +59,9 @@ def create_video_writer(output_path: str, fps: float, width: int, height: int, p
         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
@@ -100,12 +106,18 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
         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},
 }
 
+# ---------------------------------------------------------------------------
+# Two-Stage Processor
+# ---------------------------------------------------------------------------
 class TwoStageProcessor:
     def __init__(self, sam2_predictor=None, matanyone_model=None):
         self.sam2  = self._unwrap_sam2(sam2_predictor)
@@ -113,8 +125,9 @@ def __init__(self, sam2_predictor=None, matanyone_model=None):
         self.mask_cache_dir = Path("/tmp/mask_cache")
         self.mask_cache_dir.mkdir(parents=True, exist_ok=True)
 
-        # Internal flag: we "boot" MatAnyone exactly once per video
+        # Internal flags/state
         self._mat_bootstrapped = False
+        self._alpha_prev: Optional[np.ndarray] = None  # temporal smoothing
 
         logger.info(f"TwoStageProcessor init – SAM2: {self.sam2 is not None} | MatAnyOne: {self.matanyone is not None}")
 
@@ -147,7 +160,7 @@ def _get_mask(self, frame: np.ndarray) -> np.ndarray:
             return mask
 
     @staticmethod
-    def _to_binary_mask(mask: np.ndarray) -> np.ndarray:
+    def _to_binary_mask(mask: np.ndarray) -> Optional[np.ndarray]:
         """Convert mask to uint8(0..255)."""
         if mask is None:
             return None
@@ -159,7 +172,7 @@ def _to_binary_mask(mask: np.ndarray) -> np.ndarray:
         return mask
 
     @staticmethod
-    def _to_float01(mask: np.ndarray, h: int = None, w: int = None) -> np.ndarray:
+    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
@@ -177,37 +190,79 @@ def _apply_greenscreen_hard(self, frame: np.ndarray, mask: np.ndarray, bg: np.nd
         result = frame * mask_norm + bg * (1 - mask_norm)
         return result.astype(np.uint8)
 
-    def _suppress_green_spill(self, frame, amount=0.35):
-        """Suppress green spill in the frame."""
-        b, g, r = cv2.split(frame)
-        green_dominant = (g > b) & (g > r)
-        avg_br = (b.astype(np.float32) + r.astype(np.float32)) / 2
-        g_float = g.astype(np.float32)
-        g_float[green_dominant] = g_float[green_dominant] * (1 - amount) + avg_br[green_dominant] * amount
-        g = np.clip(g_float, 0, 255).astype(np.uint8)
-        return cv2.merge([b, g, r])
+    # -------- 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))
+        y = 0.299*r + 0.587*g + 0.114*b  # luminance (unused directly but good for future tuning)
+        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)
+        # protect skin tones (red significantly above green)
+        skin = (r > g + 12)
+        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:
+            # fallback to keep-all (no key)
+            return np.full(frame_bgr.shape[:2], 255, np.uint8)
 
-    # --------------------- NEW: MatAnyone bootstrap ----------------------
+        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 (e.g., refine_mask_hq) never
-        hit "first frame without a mask".
+        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)
-            # MatAnyone loaders expect RGB inputs; convert BGR→RGB
             rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
-            # Call the stateful session directly; it returns alpha [H,W] but we ignore it here.
-            _ = self.matanyone(rgb, mask_f)
+            _ = 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:
+        except Exception as e:
             logger.warning(f"MatAnyone bootstrap failed (continuing without): {e}")
 
     # ---------------------------------------------------------------------
@@ -252,7 +307,7 @@ def _prog(p, d):
             masks: List[np.ndarray] = []
             frame_idx = 0
 
-            green_bg_template = np.zeros((h, w, 3), np.uint8)  # overwritten per-frame
+            solid_bg = np.zeros((h, w, 3), np.uint8)  # overwritten per-frame
 
             while True:
                 if stop_event and stop_event.is_set():
@@ -263,7 +318,7 @@ def _prog(p, d):
                 if not ok:
                     break
 
-                # --- SAM2 segmentation (project helper) ---
+                # --- SAM2 segmentation ---
                 mask = self._get_mask(frame)
 
                 # --- MatAnyone bootstrap exactly once (first frame) ---
@@ -285,18 +340,17 @@ def _prog(p, d):
                     probe_done = True
                     logger.info(f"[TwoStage] Using key colour: {key_color_mode} → {chosen_bgr.tolist()}")
 
-                # --- Optional refinement via MatAnyone every few frames (keep your cadence) ---
+                # --- Optional refinement via MatAnyone every few frames ---
                 if self.matanyone and (frame_idx % 3 == 0):
                     try:
-                        # refine_mask_hq(frame_bgr, mask_uint8_or_float01, mat_session, fallback_enabled=True)
                         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 ---
-                green_bg_template[:] = chosen_bgr
+                solid_bg[:] = chosen_bgr
                 mask_u8 = self._to_binary_mask(mask)
-                gs = self._apply_greenscreen_hard(frame, mask_u8, green_bg_template)
+                gs = self._apply_greenscreen_hard(frame, mask_u8, solid_bg)
                 writer.write(gs)
                 masks.append(mask_u8)
 
@@ -307,7 +361,7 @@ def _prog(p, d):
             cap.release()
             writer.release()
 
-            # save mask cache (unchanged)
+            # save mask cache
             try:
                 cache_file = self.mask_cache_dir / (Path(out_path).stem + "_masks.pkl")
                 with open(cache_file, "wb") as f:
@@ -337,6 +391,7 @@ def stage2_greenscreen_to_final(
         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,   # <-- NEW: pass chosen key color
     ) -> Tuple[Optional[str], str]:
 
         def _prog(p, d):
@@ -385,9 +440,15 @@ def _prog(p, d):
 
             # Get chroma settings
             settings = chroma_settings or CHROMA_PRESETS.get('standard', {})
-            tolerance = settings.get('tolerance', 38)
-            edge_softness = settings.get('edge_softness', 2)
-            spill_suppression = settings.get('spill_suppression', 0.35)
+            tolerance = int(settings.get('tolerance', 38))
+            edge_softness = int(settings.get('edge_softness', 2))
+            spill_suppression = float(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:
@@ -406,7 +467,8 @@ def _prog(p, d):
                         frame, bg, mask,
                         tolerance=tolerance,
                         edge_softness=edge_softness,
-                        spill_suppression=spill_suppression
+                        spill_suppression=spill_suppression,
+                        key_bgr=key_bgr
                     )
                 else:
                     # Pure chroma key
@@ -414,7 +476,8 @@ def _prog(p, d):
                         frame, bg,
                         tolerance=tolerance,
                         edge_softness=edge_softness,
-                        spill_suppression=spill_suppression
+                        spill_suppression=spill_suppression,
+                        key_bgr=key_bgr
                     )
 
                 writer.write(final_frame)
@@ -432,39 +495,54 @@ def _prog(p, d):
             logger.error(f"Stage 2 error: {e}\n{traceback.format_exc()}")
             return None, f"Stage 2 failed: {e}"
 
-    def _chroma_key_composite(self, frame, bg, tolerance=38, edge_softness=2, spill_suppression=0.35):
-        """Apply chroma key compositing."""
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-
-        # Basic green range (kept as-is)
-        lower_green = np.array([40, 40, 40])
-        upper_green = np.array([80, 255, 255])
+    # ---------------- 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)
 
-        mask = cv2.inRange(hsv, lower_green, upper_green)
-        mask = cv2.bitwise_not(mask)
+        # 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
-            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
-            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
-            mask = cv2.GaussianBlur(mask, (k, k), 0)
+            alpha = cv2.GaussianBlur(alpha, (k, k), 0)
+        alpha = self._refine_alpha_edges(frame, alpha, radius=max(1, edge_softness), iters=1)
 
-        if spill_suppression > 0:
-            frame = self._suppress_green_spill(frame, spill_suppression)
+        # 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
 
-        mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
-        mask_norm = mask_3ch.astype(np.float32) / 255.0
-        result = frame * mask_norm + bg * (1 - mask_norm)
-        return result.astype(np.uint8)
+        # 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):
+    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."""
-        chroma_result = self._chroma_key_composite(frame, bg, tolerance, edge_softness, spill_suppression)
+        chroma_result = self._chroma_key_composite(
+            frame, bg,
+            tolerance=tolerance,
+            edge_softness=edge_softness,
+            spill_suppression=spill_suppression,
+            key_bgr=key_bgr
+        )
         if mask is not None:
             mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) if mask.ndim == 2 else mask
             mask_norm = mask_3ch.astype(np.float32) / 255.0
-            result = chroma_result * 0.7 + (frame * mask_norm + bg * (1 - mask_norm)) * 0.3
-            return result.astype(np.uint8)
+            guided = frame.astype(np.float32) * mask_norm + bg.astype(np.float32) * (1.0 - mask_norm)
+            result = chroma_result.astype(np.float32) * 0.7 + guided * 0.3
+            return np.clip(result, 0, 255).astype(np.uint8)
         return chroma_result
 
     # ---------------------------------------------------------------------
@@ -497,8 +575,9 @@ def _combined_progress(pct, desc):
                     pass
 
         try:
-            # Reset MatAnyone bootstrap flag between videos
+            # 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()
@@ -513,16 +592,17 @@ def _combined_progress(pct, desc):
                 progress_callback=_combined_progress,
                 stop_event=stop_event
             )
-
             if stage1_result is None:
                 return None, stage1_msg
 
-            # Stage 2
+            # 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
+                stop_event=stop_event,
+                key_bgr=key_bgr,
             )
 
             # Clean up temp file