Update utils/cv_processing.py

utils/cv_processing.py

@@ -1,20 +1,20 @@
 #!/usr/bin/env python3
 """
-cv_processing.py · slim orchestrator layer (self-contained)
+cv_processing.py · slim orchestrator layer (self-contained, backward-compatible)
 ──────────────────────────────────────────────────────────────────────────────
 Public API (unchanged):
-  - segment_person_hq(frame, predictor=None, fallback_enabled=True) -> mask (H,W) float32 [0..1]
-  - segment_person_hq_original(...) -> alias of segment_person_hq (back-compat)
-  - refine_mask_hq(frame, mask, matanyone=None, fallback_enabled=True) -> mask (H,W) float32 [0..1]
-  - replace_background_hq(frame, mask, background, fallback_enabled=True) -> frame uint8 (H,W,3)
-  - create_professional_background(key_or_cfg, width, height) -> RGB uint8 (H,W,3)
+  - segment_person_hq(frame, predictor=None, fallback_enabled=True, **compat)
+  - segment_person_hq_original(...)
+  - refine_mask_hq(frame, mask, matanyone=None, fallback_enabled=True, **compat)
+  - replace_background_hq(frame, mask, background, fallback_enabled=True)
+  - create_professional_background(key_or_cfg, width, height)
   - validate_video_file(video_path) -> (bool, reason)
 
 Design:
   * NO imports from other utils.* modules → avoids circular imports.
-  * Torch & diffusers imported lazily inside functions.
-  * All masks are single-channel float32 in [0..1] at boundaries between stages.
-  * MatAnyOne step() is fed (N,C,H,W); no 5D tensors.
+  * Torch is imported lazily inside functions.
+  * All masks are single-channel float32 in [0..1] at stage boundaries.
+  * MatAnyOne gets (N,C,H,W) — no 5D tensors.
 """
 
 from __future__ import annotations
@@ -29,7 +29,7 @@
 logger = logging.getLogger(__name__)
 
 # ----------------------------------------------------------------------------
-# Background presets (minimal set; callers can keep their own catalog if needed)
+# Background presets (local copy; safe defaults)
 # ----------------------------------------------------------------------------
 PROFESSIONAL_BACKGROUNDS_LOCAL: Dict[str, Dict[str, Any]] = {
     "office":   {"color": (240, 248, 255), "gradient": True},
@@ -39,30 +39,22 @@
     "white":    {"color": (255, 255, 255), "gradient": False},
     "black":    {"color": (0, 0, 0),       "gradient": False},
 }
+# Optional alias if callers import by this name
+PROFESSIONAL_BACKGROUNDS = PROFESSIONAL_BACKGROUNDS_LOCAL
 
 # ----------------------------------------------------------------------------
 # Helpers
 # ----------------------------------------------------------------------------
 def _ensure_rgb(img: np.ndarray) -> np.ndarray:
-    """Convert BGR→RGB if looks like BGR; otherwise pass-through."""
+    """Convert BGR→RGB if it looks like BGR (OpenCV convention)."""
     if img is None:
         return img
     if img.ndim == 3 and img.shape[2] == 3:
-        # Heuristic: assume OpenCV BGR
         return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     return img
 
-def _ensure_bgr(img: np.ndarray) -> np.ndarray:
-    """Convert RGB→BGR if looks like RGB; otherwise pass-through."""
-    if img is None:
-        return img
-    if img.ndim == 3 and img.shape[2] == 3:
-        # Heuristic: assume non-OpenCV images are RGB
-        return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    return img
-
 def _to_mask01(m: np.ndarray) -> np.ndarray:
-    """Ensure single-channel float32 [0..1]."""
+    """Return single-channel float32 in [0..1]."""
     if m is None:
         return None
     if m.ndim == 3:
@@ -73,7 +65,7 @@ def _to_mask01(m: np.ndarray) -> np.ndarray:
     return np.clip(m, 0.0, 1.0)
 
 def _feather(mask01: np.ndarray, k: int = 2) -> np.ndarray:
-    """Small Gaussian feather for cleaner edges."""
+    """Tiny Gaussian feather for smoother edges."""
     if mask01.ndim == 3:
         mask01 = mask01[..., 0]
     k = max(1, int(k) * 2 + 1)
@@ -90,13 +82,21 @@ def _vertical_gradient(top: Tuple[int,int,int], bottom: Tuple[int,int,int], widt
         bg[y, :] = (r, g, b)
     return bg
 
+def _looks_like_mask(x: Any) -> bool:
+    return (
+        isinstance(x, np.ndarray)
+        and x.ndim in (2, 3)
+        and (x.ndim == 2 or (x.ndim == 3 and x.shape[2] in (1, 3)))
+        and x.dtype != object
+    )
+
 # ----------------------------------------------------------------------------
-# Background creation (kept here to match public API)
+# Background creation (RGB)
 # ----------------------------------------------------------------------------
 def create_professional_background(key_or_cfg: Any, width: int, height: int) -> np.ndarray:
     """
     Accepts:
-      - key: str in local preset dict
+      - key: str in preset dict
       - cfg: {"color": (r,g,b), "gradient": bool}
     Returns RGB uint8 image (H,W,3).
     """
@@ -113,19 +113,14 @@ def create_professional_background(key_or_cfg: Any, width: int, height: int) ->
     if not use_grad:
         return np.full((height, width, 3), color, dtype=np.uint8)
 
-    # Simple vertical gradient dark->base color
     dark = (int(color[0]*0.7), int(color[1]*0.7), int(color[2]*0.7))
-    bg = _vertical_gradient(dark, color, width, height)
-    return bg  # already RGB by convention
+    return _vertical_gradient(dark, color, width, height)
 
 # ----------------------------------------------------------------------------
 # Segmentation
 # ----------------------------------------------------------------------------
 def _simple_person_segmentation(frame_bgr: np.ndarray) -> np.ndarray:
-    """
-    Very simple fallback segmentation by suppressing green/white backgrounds.
-    Returns mask01 (H,W) float32.
-    """
+    """Very simple fallback segmentation by suppressing green/white backgrounds."""
     hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
 
     lower_green = np.array([40, 40, 40], dtype=np.uint8)
@@ -145,13 +140,22 @@ def _simple_person_segmentation(frame_bgr: np.ndarray) -> np.ndarray:
 
     return (person_mask.astype(np.float32) / 255.0)
 
-def segment_person_hq(frame: np.ndarray, predictor: Optional[Any] = None, fallback_enabled: bool = True) -> np.ndarray:
+def segment_person_hq(
+    frame: np.ndarray,
+    predictor: Optional[Any] = None,
+    fallback_enabled: bool = True,
+    # backward-compat shim:
+    use_sam2: Optional[bool] = None,
+    **_compat_kwargs,
+) -> np.ndarray:
     """
     Try SAM2 predictor if available; return single-channel float32 mask in [0..1].
-    - predictor.set_image expects RGB
-    - predictor.predict returns masks with shapes (N,H,W) or (H,W)
+    Backward-compat: accepts use_sam2 (if False → force fallback).
     """
     try:
+        if use_sam2 is False:
+            return _simple_person_segmentation(frame)
+
         if predictor is not None and hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
             rgb = _ensure_rgb(frame)
             predictor.set_image(rgb)
@@ -164,18 +168,11 @@ def segment_person_hq(frame: np.ndarray, predictor: Optional[Any] = None, fallba
                 multimask_output=True
             )
 
-            # Normalize and pick best
-            if isinstance(masks, np.ndarray):
-                m = masks
-            else:
-                m = np.array(masks)
-
-            if m.ndim == 3:  # N,H,W
+            m = np.array(masks)
+            if m.ndim == 3:                 # (N,H,W)
                 idx = int(np.argmax(scores)) if scores is not None else 0
                 m = m[idx]
-            elif m.ndim == 2:  # H,W
-                pass
-            else:
+            elif m.ndim != 2:               # not (H,W)
                 raise RuntimeError(f"Unexpected SAM2 mask shape: {m.shape}")
 
             return _to_mask01(m)
@@ -185,7 +182,7 @@ def segment_person_hq(frame: np.ndarray, predictor: Optional[Any] = None, fallba
 
     return _simple_person_segmentation(frame) if fallback_enabled else np.ones(frame.shape[:2], dtype=np.float32)
 
-# Back-compat alias (some code may import this)
+# Back-compat alias
 segment_person_hq_original = segment_person_hq
 
 # ----------------------------------------------------------------------------
@@ -194,13 +191,11 @@ def segment_person_hq(frame: np.ndarray, predictor: Optional[Any] = None, fallba
 def _to_tensor_chw(img_uint8_bgr: np.ndarray) -> "torch.Tensor":
     import torch
     rgb = cv2.cvtColor(img_uint8_bgr, cv2.COLOR_BGR2RGB)
-    t = torch.from_numpy(rgb).permute(2, 0, 1).contiguous().float() / 255.0  # (3,H,W)
-    return t
+    return torch.from_numpy(rgb).permute(2, 0, 1).contiguous().float() / 255.0  # (3,H,W)
 
 def _mask_to_tensor01(mask01: np.ndarray) -> "torch.Tensor":
     import torch
-    m = torch.from_numpy(mask01.astype(np.float32)).unsqueeze(0).unsqueeze(0)  # (1,1,H,W)
-    return m
+    return torch.from_numpy(mask01.astype(np.float32)).unsqueeze(0).unsqueeze(0)  # (1,1,H,W)
 
 def _tensor_to_mask01(t: "torch.Tensor") -> np.ndarray:
     import torch
@@ -216,23 +211,36 @@ def _simple_mask_refinement(mask01: np.ndarray) -> np.ndarray:
     m = cv2.bilateralFilter(m, 9, 75, 75)
     return (m.astype(np.float32) / 255.0)
 
-def refine_mask_hq(frame: np.ndarray, mask: np.ndarray, matanyone: Optional[Any] = None, fallback_enabled: bool = True) -> np.ndarray:
+def refine_mask_hq(
+    frame: np.ndarray,
+    mask: np.ndarray,
+    matanyone: Optional[Any] = None,
+    fallback_enabled: bool = True,
+    # backward-compat shims:
+    use_matanyone: Optional[bool] = None,
+    **_compat_kwargs,
+) -> np.ndarray:
     """
-    If MatAnyOne processor is available, refine the mask (single-channel).
-    - Converts inputs to tensors with shapes:
-        image:  (1,3,H,W)
-        mask:   (1,1,H,W)
-    - No 5D tensors; avoids conv2d errors like [1,1,3,720,1280].
+    Refine single-channel mask with MatAnyOne if available.
+    Backward-compat:
+      - accepts use_matanyone (False → skip model)
+      - tolerates legacy arg order refine_mask_hq(mask, frame, ...)
     """
-    H, W = frame.shape[:2]
+    # tolerate legacy order: refine_mask_hq(mask, frame, ...)
+    if _looks_like_mask(frame) and isinstance(mask, np.ndarray) and mask.ndim == 3 and mask.shape[2] == 3:
+        frame, mask = mask, frame
+
     mask01 = _to_mask01(mask)
 
     try:
+        if use_matanyone is False:
+            return _simple_mask_refinement(mask01)
+
         if matanyone is not None:
             import torch
 
-            img_t  = _to_tensor_chw(frame).unsqueeze(0)   # (1,3,H,W)
-            mask_t = _mask_to_tensor01(mask01)            # (1,1,H,W)
+            img_t  = _to_tensor_chw(frame).unsqueeze(0)  # (1,3,H,W)
+            mask_t = _mask_to_tensor01(mask01)           # (1,1,H,W)
 
             device = "cuda" if torch.cuda.is_available() else "cpu"
             img_t  = img_t.to(device)
@@ -246,19 +254,15 @@ def refine_mask_hq(frame: np.ndarray, mask: np.ndarray, matanyone: Optional[Any]
                         objects=None,
                         first_frame_pred=True
                     )
-                # out should be (1,1,H,W)
                 if hasattr(matanyone, "output_prob_to_mask"):
                     out = matanyone.output_prob_to_mask(out)
                 return _tensor_to_mask01(out)
 
-            elif hasattr(matanyone, "process"):
-                # Generic .process(image, mask) path; accepts numpy/PIL
+            if hasattr(matanyone, "process"):
                 refined = matanyone.process(frame, mask01)
-                refined = np.asarray(refined).astype(np.float32)
-                return _to_mask01(refined)
+                return _to_mask01(np.asarray(refined))
 
-            else:
-                logger.warning("MatAnyOne provided but no 'step' or 'process' method found.")
+            logger.warning("MatAnyOne provided but neither 'step' nor 'process' found.")
 
     except Exception as e:
         logger.warning("MatAnyOne refinement failed: %s", e)
@@ -268,23 +272,28 @@ def refine_mask_hq(frame: np.ndarray, mask: np.ndarray, matanyone: Optional[Any]
 # ----------------------------------------------------------------------------
 # Compositing
 # ----------------------------------------------------------------------------
-def replace_background_hq(frame: np.ndarray, mask01: np.ndarray, background: np.ndarray, fallback_enabled: bool = True) -> np.ndarray:
+def replace_background_hq(
+    frame: np.ndarray,
+    mask01: np.ndarray,
+    background: np.ndarray,
+    fallback_enabled: bool = True,
+    **_compat,
+) -> np.ndarray:
     """
     Composite frame over background using feathered mask.
     Inputs:
-      - frame:      (H,W,3) uint8 (BGR or RGB, doesn't matter for linear blend)
+      - frame:      (H,W,3) uint8 (BGR or RGB, linear blend anyway)
       - mask01:     (H,W) or (H,W,1) float32 in [0..1]
       - background: (H,W,3) uint8
     Returns:
-      - composited frame (H,W,3) uint8 (same channel order as inputs)
+      - composited frame (H,W,3) uint8
     """
     try:
         H, W = frame.shape[:2]
         if background.shape[:2] != (H, W):
             background = cv2.resize(background, (W, H), interpolation=cv2.INTER_LANCZOS4)
 
-        m = _to_mask01(mask01)
-        m = _feather(m, k=2)
+        m = _feather(_to_mask01(mask01), k=2)
         m3 = np.repeat(m[:, :, None], 3, axis=2)
 
         comp = frame.astype(np.float32) * m3 + background.astype(np.float32) * (1.0 - m3)
@@ -296,7 +305,7 @@ def replace_background_hq(frame: np.ndarray, mask01: np.ndarray, background: np.
         raise
 
 # ----------------------------------------------------------------------------
-# Video validation (detailed)
+# Video validation
 # ----------------------------------------------------------------------------
 def validate_video_file(video_path: str) -> Tuple[bool, str]:
     """
@@ -350,4 +359,5 @@ def validate_video_file(video_path: str) -> Tuple[bool, str]:
     "replace_background_hq",
     "create_professional_background",
     "validate_video_file",
+    "PROFESSIONAL_BACKGROUNDS",
 ]