Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -1,56 +1,36 @@
 #!/usr/bin/env python3
 """
-Two-Stage Green Screen Processing System
-Stage 1: Original → Green Screen
-Stage 2: Green Screen → Final Background
-
-This version is aligned with the current project structure:
-- Uses segment/refine helpers from utils.cv_processing
-- Has its own safe create_video_writer (no core.app dependency)
-- Supports cancel via stop_event
-- Robust SAM2 predictor handling
+Two-Stage Green-Screen Processing System  ✅ 2025-08-26
+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
 """
 
 from __future__ import annotations
 
-import cv2
-import numpy as np
-import os
-import io
-import gc
-import pickle
-import logging
-import tempfile
-import traceback
+import cv2, numpy as np, os, io, gc, pickle, logging, tempfile, traceback, math, threading
 from pathlib import Path
-from typing import Optional, Dict, Any, Callable
+from typing import Optional, Dict, Any, Callable, Tuple, List
 
-from utils.cv_processing import (
-    segment_person_hq,
-    refine_mask_hq,
-)
+from utils.cv_processing import segment_person_hq, refine_mask_hq
 
-try:
+try:                       # project logger if available
     from utils.logger import get_logger
     logger = get_logger(__name__)
 except Exception:
     logger = logging.getLogger(__name__)
 
 
-# ---------------------------
-# Small local video I/O helper
-# ---------------------------
-def create_video_writer(
-    output_path: str,
-    fps: float,
-    width: int,
-    height: int,
-    prefer_mp4: bool = True,
-):
-    """
-    Create a cv2.VideoWriter with sane defaults.
-    Returns (writer, actual_output_path) or (None, output_path) on failure.
-    """
+# ---------------------------------------------------------------------------
+# ――― Local video-writer helper (unchanged from your previous file) ―――
+# ---------------------------------------------------------------------------
+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:
@@ -60,392 +40,369 @@ def create_video_writer(
             if curr_ext.lower() not in [".mp4", ".avi", ".mov", ".mkv"]:
                 output_path = base + ext
 
-        # pick codec
-        # mp4v works widely on Spaces; if that fails, try XVID
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v") if prefer_mp4 else cv2.VideoWriter_fourcc(*"XVID")
+        fourcc = cv2.VideoWriter_fourcc(*("mp4v" if prefer_mp4 else "XVID"))
         writer = cv2.VideoWriter(output_path, fourcc, float(fps), (int(width), int(height)))
-        if not writer or not writer.isOpened():
-            # Fallback
+        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 cv2.VideoWriter_fourcc(*"mp4v")
+            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 not writer or not writer.isOpened():
+            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
 
 
-# ---------------------------
-# Chroma key presets
-# ---------------------------
+# ---------------------------------------------------------------------------
+# ――― NEW: 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(_hue_distance(cand_hue, th) 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 (same keys, but tolerance now gets overwritten) ―――
+# ---------------------------------------------------------------------------
 CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {
-    'standard': {
-        'key_color': [0, 255, 0],     # pure green (BGR)
-        'tolerance': 38,              # color distance threshold
-        'edge_softness': 2,           # Gaussian kernel radius
-        'spill_suppression': 0.35,    # 0..1
-    },
-    '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,
-    },
+    '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:
-    """
-    Handle two-stage video processing with a green screen intermediate.
-    - Stage 1: generate clean green screen video (hard edges; great for chroma key)
-    - Stage 2: chroma-key that green to your final background
-    """
-
     def __init__(self, sam2_predictor=None, matanyone_model=None):
-        # We expect `sam2_predictor` to behave like SAM2ImagePredictor:
-        #   .set_image(np.ndarray)
-        #   .predict(point_coords=..., point_labels=..., multimask_output=True)
-        # If you passed a wrapper, we’ll try to unwrap it.
-        self.sam2 = self._unwrap_sam2(sam2_predictor)
+        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)
+        logger.info(f"TwoStageProcessor init – SAM2: {self.sam2 is not None} | MatAnyOne: {self.matanyone is not None}")
 
-        self.mask_cache_dir = Path("/tmp/mask_cache")
-        self.mask_cache_dir.mkdir(exist_ok=True, parents=True)
-
-        logger.info("TwoStageProcessor initialized. "
-                    f"SAM2 available: {self.sam2 is not None} | "
-                    f"MatAnyOne available: {self.matanyone is not None}")
-
-    # ---------------------------
-    # Stage 1: Original → Green
-    # ---------------------------
+    # ---------------------------------------------------------------------
+    # Stage 1 – Original → keyed (green/blue/…)  -- chooses colour on 1st frame
+    # ---------------------------------------------------------------------
     def stage1_extract_to_greenscreen(
         self,
         video_path: str,
         output_path: str,
-        progress_callback: Optional[Callable[[float, str], None]] = None,
+        *,
+        key_color_mode: str = "auto",           # "auto" | "green" | "blue" | "cyan" | "magenta"
+        progress_callback: Optional[Callable[[float,str],None]] = None,
         stop_event: Optional["threading.Event"] = None,
-    ):
-        """
-        Extract foreground to a pure green background.
-        Saves per-frame masks (pickle) next to the output for optional reuse.
-        """
-        def _prog(pct: float, desc: str):
-            if progress_callback:
-                try:
-                    progress_callback(float(pct), str(desc))
-                except Exception:
-                    pass
+    ) -> Tuple[Optional[dict], str]:
+        def _prog(p,d): 
+            if progress_callback: 
+                try: progress_callback(float(p), str(d)); except Exception: pass
 
         try:
-            _prog(0.0, "Stage 1: Preparing…")
+            _prog(0.0, "Stage 1: opening video…")
             cap = cv2.VideoCapture(video_path)
-            if not cap.isOpened():
-                return None, "Could not open input video"
+            if not cap.isOpened():   return None, "Could not open input video"
 
-            fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+            fps   = cap.get(cv2.CAP_PROP_FPS) or 25.0
             total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-            writer, output_path = create_video_writer(output_path, fps, width, height)
-            if writer is None:
-                cap.release()
-                return None, "Could not create output writer"
+            w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-            green_bg = np.zeros((height, width, 3), dtype=np.uint8)
-            green_bg[:, :] = [0, 255, 0]  # BGR Pure Green
+            writer,out_path = create_video_writer(output_path, fps, w, h)
+            if writer is None: 
+                cap.release(); return None, "Could not create output writer"
 
-            masks: list[np.ndarray] = []
+            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
 
+            green_bg_template = np.zeros((h,w,3), np.uint8)    # we’ll overwrite per-frame
+
             while True:
-                if stop_event is not None and stop_event.is_set():
-                    _prog(1.0, "Stage 1: Cancelled")
-                    break
+                if stop_event and stop_event.is_set():
+                    _prog(1.0, "Stage 1: cancelled"); break
 
-                ok, frame = cap.read()
-                if not ok:
-                    break
+                ok,frame = cap.read()
+                if not ok: break
 
-                # 1) get a mask (SAM2 w/ smart points via segment_person_hq)
                 mask = self._get_mask(frame)
 
-                # 2) refine occasionally with MatAnyOne to keep it light
-                if (self.matanyone is not None) and (frame_idx % 3 == 0):
-                    try:
-                        mask = refine_mask_hq(frame, mask, self.matanyone, fallback_enabled=True)
-                    except Exception as e:
-                        logger.warning(f"MatAnyOne refine failed (frame {frame_idx}): {e}")
-
-                masks.append(mask)
-
-                # 3) HARD-edge composite to green (no feather here)
-                green = self._apply_greenscreen_hard(frame, mask, green_bg)
-                writer.write(green)
+                # -------- decide key colour once --------
+                if not probe_done:
+                    if key_color_mode.lower() == "auto":
+                        key_info  = _choose_best_key_color(frame, mask)
+                        chosen_bgr= key_info["bgr"]
+                    else:
+                        cand = _key_candidates_bgr().get(key_color_mode.lower())
+                        chosen_bgr = cand["bgr"] if cand is not None else chosen_bgr
+                    probe_done = True
+                    logger.info(f"[TwoStage] Using key colour: {key_color_mode} → {chosen_bgr.tolist()}")
+
+                # optional refine
+                if self.matanyone and frame_idx % 3 == 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
+                green_bg_template[:] = chosen_bgr
+                gs = self._apply_greenscreen_hard(frame, mask, green_bg_template)
+                writer.write(gs)
+                masks.append(self._to_binary_mask(mask))
 
                 frame_idx += 1
-                if total > 0:
-                    pct = 0.05 + 0.9 * (frame_idx / total)
-                else:
-                    pct = min(0.95, 0.05 + frame_idx * 0.002)
-                _prog(pct, f"Stage 1: {frame_idx}/{total or '?'} frames")
+                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()
+            cap.release(); writer.release()
 
-            # Save masks (best-effort)
+            # save mask cache
             try:
-                mask_file = self.mask_cache_dir / (Path(output_path).stem + "_masks.pkl")
-                with open(mask_file, "wb") as f:
-                    pickle.dump(masks, f)
-                logger.info(f"Stage 1: saved masks → {mask_file}")
-            except Exception as e:
-                logger.warning(f"Stage 1: failed to save masks: {e}")
-
-            _prog(1.0, "Stage 1: Complete")
-            return output_path, f"Green screen video created ({frame_idx} frames)"
+                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)"
+            )
 
         except Exception as e:
             logger.error(f"Stage 1 error: {e}\n{traceback.format_exc()}")
             return None, f"Stage 1 failed: {e}"
 
-    # ---------------------------
-    # Stage 2: Green → Final BG
-    # ---------------------------
+    # ---------------------------------------------------------------------
+    # Stage 2 – keyed video → final composite  (hybrid matte)
+    # ---------------------------------------------------------------------
     def stage2_greenscreen_to_final(
         self,
-        greenscreen_path: str,
+        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,
+        *,
+        chroma_settings: Optional[Dict[str,Any]] = None,
+        progress_callback: Optional[Callable[[float,str],None]] = None,
         stop_event: Optional["threading.Event"] = None,
-    ):
-        """
-        Replace green screen with the given background using chroma keying.
-        `background` may be a path or an already-loaded image (BGR).
-        """
-        def _prog(pct: float, desc: str):
-            if progress_callback:
-                try:
-                    progress_callback(float(pct), str(desc))
-                except Exception:
-                    pass
+    ) -> Tuple[Optional[str], str]:
+        def _prog(p,d): 
+            if progress_callback: 
+                try: progress_callback(float(p),str(d)); except Exception: pass
 
         try:
-            _prog(0.0, "Stage 2: Preparing…")
-            cap = cv2.VideoCapture(greenscreen_path)
-            if not cap.isOpened():
-                return None, "Could not open green screen video"
+            _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
+            fps   = cap.get(cv2.CAP_PROP_FPS) or 25.0
             total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-            writer, output_path = create_video_writer(output_path, fps, width, height)
-            if writer is None:
-                cap.release()
-                return None, "Could not create output writer"
+            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/resize background
-            if isinstance(background, str):
+            # background
+            if isinstance(background,str):
                 bg = cv2.imread(background, cv2.IMREAD_COLOR)
-                if bg is None:
-                    cap.release()
-                    writer.release()
-                    return None, "Could not load background image"
-            else:
-                bg = background
-            bg = cv2.resize(bg, (width, height), interpolation=cv2.INTER_LANCZOS4)
-
-            settings = dict(CHROMA_PRESETS.get('standard', {}))
-            if chroma_settings:
-                settings.update(chroma_settings)
-
-            frame_idx = 0
-
+                if bg is None: cap.release(); writer.release(); return None,"Could not load background"
+            else: bg = background
+            bg = cv2.resize(bg,(w,h),interpolation=cv2.INTER_LANCZOS4).astype(np.uint8)
+
+            # settings
+            settings = dict(CHROMA_PRESETS['standard'])
+            if chroma_settings: settings.update(chroma_settings)
+
+            # load cached masks if any
+            cache_file = self.mask_cache_dir / (Path(gs_path).stem + "_masks.pkl")
+            cached_masks = None
+            if cache_file.exists():
+                try: cached_masks = pickle.load(open(cache_file,'rb'))
+                except Exception as e: logger.warning(f"mask cache load fail: {e}")
+
+            frame_idx=0
             while True:
-                if stop_event is not None and stop_event.is_set():
-                    _prog(1.0, "Stage 2: Cancelled")
-                    break
-
-                ok, frame = cap.read()
-                if not ok:
-                    break
+                if stop_event and stop_event.is_set(): _prog(1.0,"Stage 2: cancelled"); break
+                ok,frame = cap.read()
+                if not ok: break
 
-                out = self._chroma_key_advanced(frame, bg, settings)
-                writer.write(out)
-
-                frame_idx += 1
-                if total > 0:
-                    pct = 0.05 + 0.9 * (frame_idx / total)
+                seg_mask = None
+                if cached_masks and frame_idx < len(cached_masks):
+                    seg_mask = cached_masks[frame_idx]
                 else:
-                    pct = min(0.95, 0.05 + frame_idx * 0.002)
-                _prog(pct, f"Stage 2: {frame_idx}/{total or '?'} frames")
+                    seg_mask = self._segmentation_mask_on_stage2(frame)
 
-            cap.release()
-            writer.release()
-            _prog(1.0, "Stage 2: Complete")
+                composite = self._chroma_key_advanced(frame, bg, settings, seg_mask)
 
-            return output_path, f"Final video created ({frame_idx} frames)"
+                writer.write(composite)
+                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 video 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}"
 
-    # ---------------------------
-    # Full pipeline
-    # ---------------------------
+    # ---------------------------------------------------------------------
+    # Full pipeline – now passes chosen key into Stage 2
+    # ---------------------------------------------------------------------
     def process_full_pipeline(
         self,
         video_path: str,
         background: np.ndarray | str,
         final_output: str,
-        chroma_settings: Optional[Dict[str, Any]] = None,
-        progress_callback: Optional[Callable[[float, str], None]] = None,
+        *,
+        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,
-    ):
-        """
-        Stage 1 (to temp greenscreen) → Stage 2 (final composite).
-        """
-        gs_temp = tempfile.mktemp(suffix="_greenscreen.mp4")
+    ) -> Tuple[Optional[str], str]:
+        gs_tmp = tempfile.mktemp(suffix="_gs.mp4")
         try:
-            gs_path, msg1 = self.stage1_extract_to_greenscreen(
-                video_path, gs_temp, progress_callback=progress_callback, stop_event=stop_event
-            )
-            if gs_path is None:
-                return None, msg1
-
-            result, msg2 = self.stage2_greenscreen_to_final(
-                gs_path, background, final_output,
-                chroma_settings=chroma_settings,
-                progress_callback=progress_callback,
-                stop_event=stop_event
+            gs_info,msg1 = self.stage1_extract_to_greenscreen(
+                video_path, gs_tmp,
+                key_color_mode=key_color_mode,
+                progress_callback=progress_callback, stop_event=stop_event
             )
-            if result is None:
-                return None, msg2
+            if gs_info is None: return None,msg1
 
-            return result, msg2
+            # inject key colour into chroma settings for Stage 2
+            chosen_key = gs_info.get("key_bgr",[0,255,0])
+            cs = dict(chroma_settings or CHROMA_PRESETS['standard'])
+            cs['key_color'] = chosen_key
 
+            result,msg2 = self.stage2_greenscreen_to_final(
+                gs_info["path"], background, final_output,
+                chroma_settings=cs, progress_callback=progress_callback, stop_event=stop_event
+            )
+            return result,msg2
         finally:
-            # best-effort cleanup
-            try:
-                if os.path.exists(gs_temp):
-                    os.remove(gs_temp)
-            except Exception:
-                pass
+            try: os.remove(gs_tmp)
+            except Exception: pass
             gc.collect()
 
-    # ---------------------------
-    # Internals
-    # ---------------------------
-    def _unwrap_sam2(self, obj):
-        """
-        Try to get a callable SAM2-like predictor from whatever was passed.
-        Accepts:
-          - direct predictor (has set_image + predict)
-          - wrapper with .model that has set_image + predict
-          - wrapper with .predictor
-        """
+    # ---------------------------------------------------------------------
+    # Internal helpers  (mostly unchanged + new hybrid / seg)
+    # ---------------------------------------------------------------------
+    def _unwrap_sam2(self,obj):
         try:
-            if obj is None:
-                return None
-            # predictor directly?
-            if hasattr(obj, "set_image") and hasattr(obj, "predict"):
-                return obj
-            # wrapper.model?
-            model = getattr(obj, "model", None)
-            if model is not None and hasattr(model, "set_image") and hasattr(model, "predict"):
-                return model
-            # wrapper.predictor?
-            predictor = getattr(obj, "predictor", None)
-            if predictor is not None and hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
-                return predictor
-        except Exception as e:
-            logger.warning(f"SAM2 unwrap failed: {e}")
+            if obj is None: return None
+            if all(hasattr(obj,attr) for attr in ("set_image","predict")): return obj
+            for attr in ("model","predictor"):
+                inner=getattr(obj,attr,None)
+                if inner and all(hasattr(inner,a) for a in ("set_image","predict")): return inner
+        except Exception as e: logger.warning(f"SAM2 unwrap fail: {e}")
         return None
 
-    def _get_mask(self, frame: np.ndarray) -> np.ndarray:
-        """
-        Use our project’s enhanced segmentation helper so validation/fallbacks are consistent.
-        """
-        predictor = self.sam2
-        try:
-            mask = segment_person_hq(frame, predictor, fallback_enabled=True)
-            return mask
+    def _get_mask(self,frame:np.ndarray)->np.ndarray:
+        try: return segment_person_hq(frame,self.sam2,fallback_enabled=True)
         except Exception as e:
-            logger.warning(f"Segmentation failed, using geometric fallback: {e}")
-            h, w = frame.shape[:2]
-            m = np.zeros((h, w), dtype=np.uint8)
-            m[h//6:5*h//6, w//4:3*w//4] = 255
-            return m
-
-    def _apply_greenscreen_hard(self, frame: np.ndarray, mask: np.ndarray, green_bg: np.ndarray) -> np.ndarray:
-        """
-        Hard-edge composite to pure green for very clean keying later.
-        """
+            logger.warning(f"Segmentation fallback: {e}")
+            h,w=frame.shape[:2]; m=np.zeros((h,w),np.uint8); m[h//6:5*h//6,w//4:3*w//4]=255; return m
+
+    # ---------- stage-1 composite (same as before) ----------
+    def _apply_greenscreen_hard(self,frame,mask,green_bg):
+        mask_u8=self._to_binary_mask(mask)
+        mk=cv2.cvtColor(mask_u8,cv2.COLOR_GRAY2BGR).astype(np.float32)/255.0
+        out=frame.astype(np.float32)*mk+green_bg.astype(np.float32)*(1.0-mk)
+        return np.clip(out,0,255).astype(np.uint8)
+
+    @staticmethod
+    def _to_binary_mask(mask:np.ndarray)->np.ndarray:
+        if mask.ndim==3: mask=cv2.cvtColor(mask,cv2.COLOR_BGR2GRAY)
+        if mask.dtype!=np.uint8:
+            mask=(np.clip(mask,0,1)*255).astype(np.uint8) if mask.max()<=1.0 else np.clip(mask,0,255).astype(np.uint8)
+        _,binm=cv2.threshold(mask,127,255,cv2.THRESH_BINARY); return binm
+
+    # ---------- segmentation rescue for stage-2 ----------
+    def _segmentation_mask_on_stage2(self,frame_bgr:np.ndarray)->Optional[np.ndarray]:
         try:
-            if mask.ndim == 3:
-                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-            if mask.dtype != np.uint8:
-                mask = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
-
-            _, mask_bin = cv2.threshold(mask, 140, 255, cv2.THRESH_BINARY)
-            mask3 = cv2.cvtColor(mask_bin, cv2.COLOR_GRAY2BGR).astype(np.float32) / 255.0
+            if self.sam2 is None: return None
+            return self._get_mask(frame_bgr)
+        except Exception: return None
 
-            out = frame.astype(np.float32) * mask3 + green_bg.astype(np.float32) * (1.0 - mask3)
-            return np.clip(out, 0, 255).astype(np.uint8)
-        except Exception as e:
-            logger.error(f"Greenscreen composite failed: {e}")
-            return frame
-
-    def _chroma_key_advanced(self, frame_bgr: np.ndarray, bg_bgr: np.ndarray, settings: Dict[str, Any]) -> np.ndarray:
-        """
-        Distance-to-key color mask + soft edge + spill suppression (green reduction).
-        """
+    # ---------- hybrid chroma key ----------
+    def _chroma_key_advanced(
+        self,
+        frame_bgr: np.ndarray,
+        bg_bgr:    np.ndarray,
+        settings:  Dict[str,Any],
+        seg_mask: Optional[np.ndarray] = None,
+    )->np.ndarray:
         try:
-            key = np.array(settings.get("key_color", [0, 255, 0]), dtype=np.float32)
-            tol = float(settings.get("tolerance", 40))
-            soft = int(settings.get("edge_softness", 2))
-            spill = float(settings.get("spill_suppression", 0.3))
+            key  = np.array(settings.get("key_color",[0,255,0]),dtype=np.float32)
+            tol  = float(settings.get("tolerance",40))
+            soft = int  (settings.get("edge_softness",2))
+            spill= float(settings.get("spill_suppression",0.3))
 
             f = frame_bgr.astype(np.float32)
             b = bg_bgr.astype(np.float32)
 
-            # distance (BGR space)
-            diff = f - key
-            dist = np.sqrt((diff ** 2).sum(axis=2))
-
-            # inside green → 0, far from green → 1
-            mask = np.clip((dist - tol) / max(tol, 1.0), 0.0, 1.0)
-
-            if soft > 0:
-                ksize = max(1, soft * 2 + 1)
-                mask = cv2.GaussianBlur(mask.astype(np.float32), (ksize, ksize), soft)
-
-            # spill suppression
-            if spill > 0:
-                # where mask < 1.0 (near edges), reduce green channel proportionally
-                spill_zone = 1.0 - mask
-                g = f[:, :, 1]
-                f[:, :, 1] = np.clip(g - g * spill_zone * spill, 0, 255)
-
-            mask3 = np.stack([mask] * 3, axis=2)
-            out = f * mask3 + b * (1.0 - mask3)
-            return np.clip(out, 0, 255).astype(np.uint8)
+            diff = np.linalg.norm(f-key,axis=2)
+            alpha = np.clip((diff - tol*0.6) / max(1e-6,tol*0.4), 0.0, 1.0)
+            if soft>0:
+                k=soft*2+1; alpha=cv2.GaussianBlur(alpha,(k,k),soft)
+
+            # ---------- segmentation rescue ----------
+            if seg_mask is not None:
+                if seg_mask.ndim==3: seg_mask=cv2.cvtColor(seg_mask,cv2.COLOR_BGR2GRAY)
+                seg = seg_mask.astype(np.float32)/255.0
+                seg = cv2.GaussianBlur(seg,(5,5),1.0)
+                alpha=np.clip(np.maximum(alpha,seg*0.85),0.0,1.0)
+
+            # ---------- spill suppression ----------
+            if spill>0:
+                zone = 1.0-alpha
+                g=f[:,:,1]; f[:,:,1]=np.clip(g - g*zone*spill,0,255)
+
+            mask3=np.stack([alpha]*3,axis=2)
+            out = f*mask3 + b*(1.0-mask3)
+            return np.clip(out,0,255).astype(np.uint8)
         except Exception as e:
-            logger.error(f"Chroma keying failed: {e}")
+            logger.error(f"Chroma key error: {e}")
             return frame_bgr