Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -14,13 +14,14 @@
 
 from __future__ import annotations
 
-import cv2, numpy as np, os, io, gc, pickle, logging, tempfile, traceback, math, threading
+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
 
-try:                       # project logger if available
+# Project logger if available
+try:
     from utils.logger import get_logger
     logger = get_logger(__name__)
 except Exception:
@@ -28,7 +29,7 @@
 
 
 # ---------------------------------------------------------------------------
-# ――― Local video-writer helper (unchanged from your previous file) ―――
+# Local video-writer helper
 # ---------------------------------------------------------------------------
 def create_video_writer(output_path: str, fps: float, width: int, height: int, prefer_mp4: bool = True):
     try:
@@ -57,7 +58,7 @@ def create_video_writer(output_path: str, fps: float, width: int, height: int, p
 
 
 # ---------------------------------------------------------------------------
-# ――― NEW: key-colour helpers (fast, no external deps) ―――
+# 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)
@@ -95,7 +96,7 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
         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)
+            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"])
@@ -104,7 +105,7 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
 
 
 # ---------------------------------------------------------------------------
-# ――― Chroma presets (same keys, but tolerance now gets overwritten) ―――
+# Chroma presets
 # ---------------------------------------------------------------------------
 CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {
     'standard': {'key_color': [0,255,0], 'tolerance': 38, 'edge_softness': 2, 'spill_suppression': 0.35},
@@ -114,13 +115,14 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
 
 
 # ---------------------------------------------------------------------------
-# ――― Two-Stage Processor ―――
+# 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)
+        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}")
 
     # ---------------------------------------------------------------------
@@ -132,59 +134,70 @@ def stage1_extract_to_greenscreen(
         output_path: str,
         *,
         key_color_mode: str = "auto",           # "auto" | "green" | "blue" | "cyan" | "magenta"
-        progress_callback: Optional[Callable[[float,str],None]] = None,
+        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
+
+        def _prog(p, d):
+            if progress_callback:
+                try:
+                    progress_callback(float(p), str(d))
+                except Exception:
+                    pass
 
         try:
             _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
             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"
+            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] = []
+            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
+            green_bg_template = 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
+                    _prog(1.0, "Stage 1: cancelled")
+                    break
 
-                ok,frame = cap.read()
-                if not ok: break
+                ok, frame = cap.read()
+                if not ok:
+                    break
 
                 mask = self._get_mask(frame)
 
-                # -------- decide key colour once --------
+                # 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"]
+                        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
+                        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 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}")
+                    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
@@ -193,18 +206,21 @@ def _prog(p,d):
                 masks.append(self._to_binary_mask(mask))
 
                 frame_idx += 1
-                pct = 0.05 + 0.9 * (frame_idx/total) if total else min(0.95, 0.05+frame_idx*0.002)
+                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 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}")
+                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")
+            _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)"
@@ -223,186 +239,20 @@ def stage2_greenscreen_to_final(
         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,
     ) -> 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: 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))
-
-            writer,out_path = create_video_writer(output_path, fps, w, h)
-            if writer is None: cap.release(); return None,"Could not create output writer"
-
-            # 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"
-            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)
+        def _prog(p, d):
+            if progress_callback:
+                try:
+                    progress_callback(float(p), str(d))
+                except Exception:
+                    pass
 
-            # 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 and stop_event.is_set(): _prog(1.0,"Stage 2: cancelled"); break
-                ok,frame = cap.read()
-                if not ok: break
-
-                seg_mask = None
-                if cached_masks and frame_idx < len(cached_masks):
-                    seg_mask = cached_masks[frame_idx]
-                else:
-                    seg_mask = self._segmentation_mask_on_stage2(frame)
-
-                composite = self._chroma_key_advanced(frame, bg, settings, seg_mask)
-
-                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 – now passes chosen key into Stage 2
-    # ---------------------------------------------------------------------
-    def process_full_pipeline(
-        self,
-        video_path: str,
-        background: np.ndarray | str,
-        final_output: 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]:
-        gs_tmp = tempfile.mktemp(suffix="_gs.mp4")
         try:
-            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 gs_info is None: return None,msg1
-
-            # 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:
-            try: os.remove(gs_tmp)
-            except Exception: pass
-            gc.collect()
-
-    # ---------------------------------------------------------------------
-    # Internal helpers  (mostly unchanged + new hybrid / seg)
-    # ---------------------------------------------------------------------
-    def _unwrap_sam2(self,obj):
-        try:
-            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:
-        try: return segment_person_hq(frame,self.sam2,fallback_enabled=True)
-        except Exception as e:
-            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 self.sam2 is None: return None
-            return self._get_mask(frame_bgr)
-        except Exception: return None
-
-    # ---------- 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))
-
-            f = frame_bgr.astype(np.float32)
-            b = bg_bgr.astype(np.float32)
-
-            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 key error: {e}")
-            return frame_bgr
+            _prog(0.0, "Stage 2: opening keyed video…")
+            cap = cv2.VideoCapture(gs_path)
+            if not cap.isOpened():
+                return None, "