Update utils/cv_processing.py

utils/cv_processing.py

@@ -1,72 +1,302 @@
 #!/usr/bin/env python3
 """
-cv_processing.py · slim orchestrator layer
+cv_processing.py · slim orchestrator layer (self-contained)
 ──────────────────────────────────────────────────────────────────────────────
-Keeps the public API (segment_person_hq, refine_mask_hq, replace_background_hq,
-create_professional_background, validate_video_file) exactly the same so that
-existing callers do **not** need to change their imports.
-
-All heavy-lifting implementations live in:
-    utils.segmentation
-    utils.refinement
-    utils.compositing
-    utils.background_factory
-    utils.background_presets
+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)
+  - 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.
 """
 
 from __future__ import annotations
 
-# ── std / 3rd-party ────────────────────────────────────────────────────────
-import os, logging, cv2, numpy as np
+import logging
 from pathlib import Path
-from typing   import Tuple, Dict, Any, Optional
-
-# ── project helpers (new modules we split out) ─────────────────────────────
-from utils.segmentation       import (
-    segment_person_hq,
-    segment_person_hq_original,
-    SegmentationError,
-)
-from utils.refinement         import (
-    refine_mask_hq, MaskRefinementError,
-)
-from .compositing import (
-    replace_background_hq, BackgroundReplacementError,
-)
-from utils.background_factory import create_professional_background
-from utils.background_presets import PROFESSIONAL_BACKGROUNDS      # still used in the UI
+from typing import Any, Dict, Optional, Tuple
+
+import cv2
+import numpy as np
 
 logger = logging.getLogger(__name__)
 
 # ----------------------------------------------------------------------------
-# LIGHT CONFIG – only what the UI still needs
+# Background presets (minimal set; callers can keep their own catalog if needed)
+# ----------------------------------------------------------------------------
+PROFESSIONAL_BACKGROUNDS_LOCAL: Dict[str, Dict[str, Any]] = {
+    "office":   {"color": (240, 248, 255), "gradient": True},
+    "studio":   {"color": (32, 32, 32),    "gradient": False},
+    "nature":   {"color": (34, 139, 34),   "gradient": True},
+    "abstract": {"color": (75, 0, 130),    "gradient": True},
+    "white":    {"color": (255, 255, 255), "gradient": False},
+    "black":    {"color": (0, 0, 0),       "gradient": False},
+}
+
+# ----------------------------------------------------------------------------
+# Helpers
 # ----------------------------------------------------------------------------
-USE_AUTO_TEMPORAL_CONSISTENCY = True   # placeholder for future smoothing
+def _ensure_rgb(img: np.ndarray) -> np.ndarray:
+    """Convert BGR→RGB if looks like BGR; otherwise pass-through."""
+    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
 
-# Validator soft-limits (kept here because validate_video_file still lives here)
-MIN_AREA_RATIO = 0.015
-MAX_AREA_RATIO = 0.97
+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]."""
+    if m is None:
+        return None
+    if m.ndim == 3:
+        m = m[..., 0]
+    m = m.astype(np.float32)
+    if m.max() > 1.0:
+        m = m / 255.0
+    return np.clip(m, 0.0, 1.0)
+
+def _feather(mask01: np.ndarray, k: int = 2) -> np.ndarray:
+    """Small Gaussian feather for cleaner edges."""
+    if mask01.ndim == 3:
+        mask01 = mask01[..., 0]
+    k = max(1, int(k) * 2 + 1)
+    m = cv2.GaussianBlur((mask01 * 255.0).astype(np.uint8), (k, k), 0)
+    return (m.astype(np.float32) / 255.0)
+
+def _vertical_gradient(top: Tuple[int,int,int], bottom: Tuple[int,int,int], width: int, height: int) -> np.ndarray:
+    bg = np.zeros((height, width, 3), dtype=np.uint8)
+    for y in range(height):
+        t = y / max(1, height - 1)
+        r = int(top[0] * (1 - t) + bottom[0] * t)
+        g = int(top[1] * (1 - t) + bottom[1] * t)
+        b = int(top[2] * (1 - t) + bottom[2] * t)
+        bg[y, :] = (r, g, b)
+    return bg
 
 # ----------------------------------------------------------------------------
-# PUBLIC 1-LINERS to keep old call-sites working
+# Background creation (kept here to match public API)
 # ----------------------------------------------------------------------------
-# (They're just re-exports from their new homes.)
+def create_professional_background(key_or_cfg: Any, width: int, height: int) -> np.ndarray:
+    """
+    Accepts:
+      - key: str in local preset dict
+      - cfg: {"color": (r,g,b), "gradient": bool}
+    Returns RGB uint8 image (H,W,3).
+    """
+    if isinstance(key_or_cfg, str):
+        cfg = PROFESSIONAL_BACKGROUNDS_LOCAL.get(key_or_cfg, PROFESSIONAL_BACKGROUNDS_LOCAL["office"])
+    elif isinstance(key_or_cfg, dict):
+        cfg = key_or_cfg
+    else:
+        cfg = PROFESSIONAL_BACKGROUNDS_LOCAL["office"]
 
-__all__ = [
-    "segment_person_hq",
-    "segment_person_hq_original",
-    "refine_mask_hq",
-    "replace_background_hq",
-    "create_professional_background",
-    "validate_video_file",
-    "SegmentationError",
-    "MaskRefinementError",
-    "BackgroundReplacementError",
-    "PROFESSIONAL_BACKGROUNDS",
-]
+    color = tuple(int(x) for x in cfg.get("color", (255, 255, 255)))
+    use_grad = bool(cfg.get("gradient", False))
+
+    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
+
+# ----------------------------------------------------------------------------
+# 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.
+    """
+    hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
+
+    lower_green = np.array([40, 40, 40], dtype=np.uint8)
+    upper_green = np.array([80, 255, 255], dtype=np.uint8)
+    green_mask = cv2.inRange(hsv, lower_green, upper_green)
+
+    lower_white = np.array([0, 0, 200], dtype=np.uint8)
+    upper_white = np.array([180, 30, 255], dtype=np.uint8)
+    white_mask = cv2.inRange(hsv, lower_white, upper_white)
+
+    bg_mask = cv2.bitwise_or(green_mask, white_mask)
+    person_mask = cv2.bitwise_not(bg_mask)
+
+    kernel = np.ones((5, 5), np.uint8)
+    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_CLOSE, kernel)
+    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_OPEN, kernel)
+
+    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:
+    """
+    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)
+    """
+    try:
+        if predictor is not None and hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
+            rgb = _ensure_rgb(frame)
+            predictor.set_image(rgb)
+            h, w = rgb.shape[:2]
+            center = np.array([[w // 2, h // 2]])
+            labels = np.array([1])
+            masks, scores, _ = predictor.predict(
+                point_coords=center,
+                point_labels=labels,
+                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
+                idx = int(np.argmax(scores)) if scores is not None else 0
+                m = m[idx]
+            elif m.ndim == 2:  # H,W
+                pass
+            else:
+                raise RuntimeError(f"Unexpected SAM2 mask shape: {m.shape}")
+
+            return _to_mask01(m)
+
+    except Exception as e:
+        logger.warning("SAM2 segmentation failed: %s", e)
+
+    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)
+segment_person_hq_original = segment_person_hq
+
+# ----------------------------------------------------------------------------
+# Refinement (MatAnyOne)
+# ----------------------------------------------------------------------------
+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
+
+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
+
+def _tensor_to_mask01(t: "torch.Tensor") -> np.ndarray:
+    import torch
+    if t.ndim == 4:
+        t = t[0, 0]
+    elif t.ndim == 3:
+        t = t[0]
+    return np.clip(t.detach().float().cpu().numpy(), 0.0, 1.0)
+
+def _simple_mask_refinement(mask01: np.ndarray) -> np.ndarray:
+    m = (mask01 * 255.0).astype(np.uint8)
+    m = cv2.GaussianBlur(m, (5, 5), 0)
+    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:
+    """
+    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].
+    """
+    H, W = frame.shape[:2]
+    mask01 = _to_mask01(mask)
+
+    try:
+        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)
+
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            img_t  = img_t.to(device)
+            mask_t = mask_t.to(device)
+
+            if hasattr(matanyone, "step"):
+                with torch.inference_mode():
+                    out = matanyone.step(
+                        image_tensor=img_t,
+                        mask_tensor=mask_t,
+                        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
+                refined = matanyone.process(frame, mask01)
+                refined = np.asarray(refined).astype(np.float32)
+                return _to_mask01(refined)
+
+            else:
+                logger.warning("MatAnyOne provided but no 'step' or 'process' method found.")
+
+    except Exception as e:
+        logger.warning("MatAnyOne refinement failed: %s", e)
+
+    return _simple_mask_refinement(mask01) if fallback_enabled else mask01
 
 # ----------------------------------------------------------------------------
-# VIDEO VALIDATION (unchanged)
+# Compositing
+# ----------------------------------------------------------------------------
+def replace_background_hq(frame: np.ndarray, mask01: np.ndarray, background: np.ndarray, fallback_enabled: bool = True) -> np.ndarray:
+    """
+    Composite frame over background using feathered mask.
+    Inputs:
+      - frame:      (H,W,3) uint8 (BGR or RGB, doesn't matter for linear blend)
+      - 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)
+    """
+    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)
+        m3 = np.repeat(m[:, :, None], 3, axis=2)
+
+        comp = frame.astype(np.float32) * m3 + background.astype(np.float32) * (1.0 - m3)
+        return np.clip(comp, 0, 255).astype(np.uint8)
+    except Exception as e:
+        if fallback_enabled:
+            logger.warning("Compositing failed (%s) – returning original frame", e)
+            return frame
+        raise
+
+# ----------------------------------------------------------------------------
+# Video validation (detailed)
 # ----------------------------------------------------------------------------
 def validate_video_file(video_path: str) -> Tuple[bool, str]:
     """
@@ -108,4 +338,16 @@ def validate_video_file(video_path: str) -> Tuple[bool, str]:
 
     except Exception as e:
         logger.error(f"validate_video_file: {e}")
-        return False, f"Validation error: {e}"
+        return False, f"Validation error: {e}"
+
+# ----------------------------------------------------------------------------
+# Public symbols
+# ----------------------------------------------------------------------------
+__all__ = [
+    "segment_person_hq",
+    "segment_person_hq_original",
+    "refine_mask_hq",
+    "replace_background_hq",
+    "create_professional_background",
+    "validate_video_file",
+]