Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -15,6 +15,11 @@
   * 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
 
@@ -115,6 +120,15 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
     'outdoor':  {'key_color': [0,255,0], 'tolerance': 50, 'edge_softness': 3, 'spill_suppression': 0.25},
 }
 
+# ---------------------------------------------------------------------------
+# 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),
+}
+
 # ---------------------------------------------------------------------------
 # Two-Stage Processor
 # ---------------------------------------------------------------------------
@@ -129,6 +143,14 @@ def __init__(self, sam2_predictor=None, matanyone_model=None):
         self._mat_bootstrapped = False
         self._alpha_prev: Optional[np.ndarray] = None  # temporal smoothing
 
+        # Quality selection
+        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 ---------------------------
@@ -197,12 +219,10 @@ def _suppress_green_spill(self, frame: np.ndarray, amount: float = 0.35) -> np.n
         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)
+        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
@@ -237,7 +257,6 @@ def _soft_key_mask(self, frame_bgr: np.ndarray, key_bgr: np.ndarray, tol: int =
         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)
 
         ycbcr = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2YCrCb).astype(np.float32)
@@ -340,8 +359,9 @@ 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 ---
-                if self.matanyone and (frame_idx % 3 == 0):
+                # --- 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:
@@ -391,7 +411,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
+        key_bgr: Optional[np.ndarray] = None,   # pass chosen key color
     ) -> Tuple[Optional[str], str]:
 
         def _prog(p, d):
@@ -422,6 +442,11 @@ def _prog(p, d):
             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()
@@ -438,11 +463,11 @@ def _prog(p, d):
             except Exception as e:
                 logger.warning(f"Could not load mask cache: {e}")
 
-            # Get chroma settings
+            # Get chroma settings and override with profile
             settings = chroma_settings or CHROMA_PRESETS.get('standard', {})
-            tolerance = int(settings.get('tolerance', 38))
-            edge_softness = int(settings.get('edge_softness', 2))
-            spill_suppression = float(settings.get('spill_suppression', 0.35))
+            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:
@@ -529,7 +554,7 @@ def _chroma_key_composite(self, frame, bg, *, tolerance=38, edge_softness=2, spi
         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."""
+        """Apply hybrid compositing using both chroma key and cached mask, with profile controls."""
         chroma_result = self._chroma_key_composite(
             frame, bg,
             tolerance=tolerance,
@@ -537,13 +562,28 @@ def _hybrid_composite(self, frame, bg, mask, *, tolerance=38, edge_softness=2, s
             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
-            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
+        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