Update utils/cv_processing.py

utils/cv_processing.py

@@ -1,14 +1,6 @@
 #!/usr/bin/env python3
 """
-cv_processing.py · slim orchestrator layer (self-contained, backward-compatible)
-──────────────────────────────────────────────────────────────────────────────
-Public API (unchanged):
-  - 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)
+cv_processing.py · FIXED VERSION with proper SAM2 handling
 """
 
 from __future__ import annotations
@@ -23,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 # ----------------------------------------------------------------------------
-# Background presets (local copy; safe defaults)
+# Background presets
 # ----------------------------------------------------------------------------
 PROFESSIONAL_BACKGROUNDS_LOCAL: Dict[str, Dict[str, Any]] = {
     "office":   {"color": (240, 248, 255), "gradient": True},
@@ -33,7 +25,7 @@
     "white":    {"color": (255, 255, 255), "gradient": False},
     "black":    {"color": (0, 0, 0),       "gradient": False},
 }
-PROFESSIONAL_BACKGROUNDS = PROFESSIONAL_BACKGROUNDS_LOCAL  # alias for callers
+PROFESSIONAL_BACKGROUNDS = PROFESSIONAL_BACKGROUNDS_LOCAL
 
 # ----------------------------------------------------------------------------
 # Helpers
@@ -42,7 +34,6 @@ def _ensure_rgb(img: np.ndarray) -> np.ndarray:
     if img is None:
         return img
     if img.ndim == 3 and img.shape[2] == 3:
-        # Assume OpenCV BGR → convert to RGB
         return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     return img
 
@@ -73,16 +64,8 @@ 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 (RGB)
+# Background creation
 # ----------------------------------------------------------------------------
 def create_professional_background(key_or_cfg: Any, width: int, height: int) -> np.ndarray:
     if isinstance(key_or_cfg, str):
@@ -102,213 +85,262 @@ def create_professional_background(key_or_cfg: Any, width: int, height: int) ->
     return _vertical_gradient(dark, color, width, height)
 
 # ----------------------------------------------------------------------------
-# Segmentation
+# Improved Segmentation
 # ----------------------------------------------------------------------------
 def _simple_person_segmentation(frame_bgr: np.ndarray) -> np.ndarray:
+    """Basic fallback segmentation using color detection"""
+    h, w = frame_bgr.shape[:2]
+    
+    # Convert to HSV for better color detection
     hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
-
+    
+    # Detect skin tones (basic person detection)
+    lower_skin = np.array([0, 20, 70], dtype=np.uint8)
+    upper_skin = np.array([20, 255, 255], dtype=np.uint8)
+    skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
+    
+    # Also detect non-green/non-white areas as potential person
     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)
-
+    
+    # Assume person is NOT green screen
+    person_mask = cv2.bitwise_not(green_mask)
+    
+    # Combine with skin detection
+    person_mask = cv2.bitwise_or(person_mask, skin_mask)
+    
+    # Clean up the mask
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_OPEN, kernel, iterations=1)
+    
+    # Find largest contour (assume it's the person)
+    contours, _ = cv2.findContours(person_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if contours:
+        largest_contour = max(contours, key=cv2.contourArea)
+        person_mask = np.zeros_like(person_mask)
+        cv2.drawContours(person_mask, [largest_contour], -1, 255, -1)
+    
     return (person_mask.astype(np.float32) / 255.0)
 
 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:
-        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)
-            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
-            )
-            m = np.array(masks)
-            if m.ndim == 3:
-                idx = int(np.argmax(scores)) if scores is not None else 0
-                m = m[idx]
-            elif m.ndim != 2:
-                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)
-
-segment_person_hq_original = segment_person_hq  # back-compat alias
-
-# ----------------------------------------------------------------------------
-# MatAnyOne helpers
-# ----------------------------------------------------------------------------
-def _to_tensor_chw(img_uint8_bgr: np.ndarray) -> "torch.Tensor":
-    import torch
-    rgb = cv2.cvtColor(img_uint8_bgr, cv2.COLOR_BGR2RGB)
-    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
-    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
-    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 _remap_harden(mask01: np.ndarray, inside: float = 0.70, outside: float = 0.35) -> np.ndarray:
     """
-    Pull the mask toward {0,1} to avoid 'ghost' translucency.
-    Values <= outside -> 0; >= inside -> 1; linear in between.
+    High-quality person segmentation with proper SAM2 handling
     """
-    m = mask01.astype(np.float32)
-    if inside <= outside:
-        return m
-    m = (m - outside) / max(1e-6, (inside - outside))
-    return np.clip(m, 0.0, 1.0)
-
-def _pad_and_smooth_edges(mask01: np.ndarray, dilate_px: int = 6, edge_blur_px: int = 2) -> np.ndarray:
-    m = (mask01 * 255.0).astype(np.uint8)
-    if dilate_px > 0:
-        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_px, dilate_px))
-        m = cv2.dilate(m, k, iterations=1)
-    if edge_blur_px > 0:
-        ksize = edge_blur_px * 2 + 1
-        m = cv2.GaussianBlur(m, (ksize, ksize), 0)
-    return (m.astype(np.float32) / 255.0)
-
-def _try_matanyone_refine(
-    matanyone: Any,
-    frame_bgr: np.ndarray,
-    mask01: np.ndarray
-) -> Optional[np.ndarray]:
-    """
-    Try several MatAnyOne interfaces:
-      1) InferenceCore.infer(PIL_image, PIL_mask)
-      2) .step(image_tensor=NCHW, mask_tensor=NCHW)
-      3) .process(image_np, mask_np)
-      4) callable(image_tensor, mask_tensor) → tensor
-    Returns refined mask01 (np.ndarray) or None if not usable.
-    """
-    try:
-        # --- (1) PIL infer path ------------------------------------------------
-        if hasattr(matanyone, "infer"):
-            try:
-                from PIL import Image
-                img_pil  = Image.fromarray(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
-                m_pil    = Image.fromarray((mask01 * 255.0).astype(np.uint8))
-                out_pil  = matanyone.infer(img_pil, m_pil)
-                out_np   = np.asarray(out_pil).astype(np.float32)
-                return _to_mask01(out_np)
-            except Exception as e:
-                logger.debug("MatAnyOne.infer path failed: %s", e)
-
-        # --- (2) tensor .step path --------------------------------------------
-        if hasattr(matanyone, "step"):
-            import torch
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            img_t  = _to_tensor_chw(frame_bgr).unsqueeze(0).to(device)   # (1,3,H,W)
-            mask_t = _mask_to_tensor01(mask01).to(device)                # (1,1,H,W)
-            with torch.inference_mode():
-                out = matanyone.step(
-                    image_tensor=img_t,
-                    mask_tensor=mask_t,
-                    objects=None,
-                    first_frame_pred=True
+    h, w = frame.shape[:2]
+    
+    # Skip SAM2 if explicitly disabled
+    if use_sam2 is False:
+        return _simple_person_segmentation(frame)
+    
+    # Try SAM2 if available
+    if predictor is not None:
+        try:
+            # Ensure we have the right methods
+            if hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
+                # Convert to RGB for SAM2
+                rgb = _ensure_rgb(frame)
+                
+                # Set the image
+                predictor.set_image(rgb)
+                
+                # Generate multiple prompt points for better coverage
+                points = []
+                labels = []
+                
+                # Add center point
+                points.append([w // 2, h // 2])
+                labels.append(1)  # Foreground
+                
+                # Add points for head area (upper center)
+                points.append([w // 2, h // 4])
+                labels.append(1)
+                
+                # Add body points
+                points.append([w // 2, h // 2 + h // 8])
+                labels.append(1)
+                
+                # Convert to numpy arrays
+                point_coords = np.array(points, dtype=np.float32)
+                point_labels = np.array(labels, dtype=np.int32)
+                
+                # Predict with multiple masks
+                result = predictor.predict(
+                    point_coords=point_coords,
+                    point_labels=point_labels,
+                    multimask_output=True
                 )
-            if hasattr(matanyone, "output_prob_to_mask"):
-                out = matanyone.output_prob_to_mask(out)
-            return _tensor_to_mask01(out)
-
-        # --- (3) numpy .process path ------------------------------------------
-        if hasattr(matanyone, "process"):
-            out = matanyone.process(frame_bgr, mask01)
-            return _to_mask01(np.asarray(out))
-
-        # --- (4) callable / nn.Module path ------------------------------------
-        if callable(matanyone):
-            import torch
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            img_t  = _to_tensor_chw(frame_bgr).unsqueeze(0).to(device)
-            mask_t = _mask_to_tensor01(mask01).to(device)
-            with torch.inference_mode():
-                out = matanyone(img_t, mask_t)
-            return _tensor_to_mask01(out)
-
-    except Exception as e:
-        logger.warning("MatAnyOne refine error: %s", e)
+                
+                # Extract masks and scores
+                if isinstance(result, dict):
+                    masks = result.get("masks", None)
+                    scores = result.get("scores", None)
+                elif isinstance(result, tuple) and len(result) >= 2:
+                    masks, scores = result[0], result[1]
+                else:
+                    masks = result
+                    scores = None
+                
+                # Validate and process masks
+                if masks is not None:
+                    masks = np.array(masks)
+                    
+                    if masks.size > 0:  # Check if not empty
+                        # Handle different mask shapes
+                        if masks.ndim == 3 and masks.shape[0] > 0:
+                            # Multiple masks - choose best one
+                            if scores is not None and len(scores) > 0:
+                                best_idx = np.argmax(scores)
+                                mask = masks[best_idx]
+                            else:
+                                # Use first mask if no scores
+                                mask = masks[0]
+                        elif masks.ndim == 2:
+                            # Single mask
+                            mask = masks
+                        else:
+                            logger.warning(f"Unexpected mask shape from SAM2: {masks.shape}")
+                            mask = None
+                        
+                        if mask is not None:
+                            # Convert to proper format
+                            mask = _to_mask01(mask)
+                            
+                            # Validate mask has actual content
+                            if mask.max() > 0.1:  # At least 10% confidence somewhere
+                                return mask
+                            else:
+                                logger.warning("SAM2 mask too weak, using fallback")
+                else:
+                    logger.warning("SAM2 returned no masks")
+                    
+        except Exception as e:
+            logger.warning(f"SAM2 segmentation error: {e}")
+    
+    # Fallback to simple segmentation
+    if fallback_enabled:
+        logger.debug("Using fallback segmentation")
+        return _simple_person_segmentation(frame)
+    else:
+        # Return full mask if no fallback
+        return np.ones((h, w), dtype=np.float32)
 
-    return None
+segment_person_hq_original = segment_person_hq
 
 # ----------------------------------------------------------------------------
-# Refinement (MatAnyOne)
+# MatAnyone Refinement (Fixed)
 # ----------------------------------------------------------------------------
 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:
     """
-    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, ...)
+    Refine mask with MatAnyone - with proper handling
     """
-    # tolerate legacy order: refine_mask_hq(mask, frame, ...)
-    if _looks_like_mask(frame) and _looks_like_mask(mask) and mask.ndim == 3 and mask.shape[2] == 3:
-        frame, mask = mask, frame  # swap
-
+    # Convert mask to proper format
     mask01 = _to_mask01(mask)
-
-    # Use MatAnyOne when possible
-    if use_matanyone is not False and matanyone is not None:
-        refined = _try_matanyone_refine(matanyone, frame, mask01)
-        if refined is not None:
-            # Hardening + edge handling to avoid translucent body/halo
-            refined = _remap_harden(refined, inside=0.70, outside=0.35)
-            refined = _pad_and_smooth_edges(refined, dilate_px=4, edge_blur_px=1)
-            return refined
-        else:
-            logger.warning("MatAnyOne provided but no usable interface found; falling back.")
-
-    # Simple refinement fallback
-    m = (mask01 * 255.0).astype(np.uint8)
-    m = cv2.GaussianBlur(m, (5, 5), 0)
-    m = cv2.bilateralFilter(m, 9, 75, 75)
-    m = (m.astype(np.float32) / 255.0)
-    m = _remap_harden(m, inside=0.68, outside=0.40)
-    m = _pad_and_smooth_edges(m, dilate_px=3, edge_blur_px=1)
-    return m if fallback_enabled else mask01
+    
+    # Skip MatAnyone if explicitly disabled
+    if use_matanyone is False:
+        return mask01
+    
+    # Try MatAnyone if available
+    if matanyone is not None:
+        try:
+            # Try different MatAnyone interfaces
+            refined = None
+            
+            # Method 1: Direct callable
+            if callable(matanyone):
+                try:
+                    refined = matanyone(frame, mask01)
+                    if refined is not None:
+                        refined = _to_mask01(np.array(refined))
+                except Exception as e:
+                    logger.debug(f"MatAnyone callable failed: {e}")
+            
+            # Method 2: step method
+            if refined is None and hasattr(matanyone, 'step'):
+                try:
+                    refined = matanyone.step(frame, mask01)
+                    if refined is not None:
+                        refined = _to_mask01(np.array(refined))
+                except Exception as e:
+                    logger.debug(f"MatAnyone step failed: {e}")
+            
+            # Method 3: process method
+            if refined is None and hasattr(matanyone, 'process'):
+                try:
+                    refined = matanyone.process(frame, mask01)
+                    if refined is not None:
+                        refined = _to_mask01(np.array(refined))
+                except Exception as e:
+                    logger.debug(f"MatAnyone process failed: {e}")
+            
+            # Use refined mask if successful
+            if refined is not None and refined.max() > 0.1:
+                # Apply post-processing
+                refined = _postprocess_mask(refined)
+                return refined
+            else:
+                logger.warning("MatAnyone refinement failed or produced empty mask")
+                
+        except Exception as e:
+            logger.warning(f"MatAnyone error: {e}")
+    
+    # Fallback refinement
+    if fallback_enabled:
+        return _fallback_refine(mask01)
+    else:
+        return mask01
+
+def _postprocess_mask(mask01: np.ndarray) -> np.ndarray:
+    """Post-process mask to clean edges and remove artifacts"""
+    # Convert to uint8
+    mask_uint8 = (mask01 * 255).astype(np.uint8)
+    
+    # Remove small holes
+    kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel_close)
+    
+    # Smooth edges
+    mask_uint8 = cv2.GaussianBlur(mask_uint8, (3, 3), 0)
+    
+    # Threshold to clean up
+    _, mask_uint8 = cv2.threshold(mask_uint8, 127, 255, cv2.THRESH_BINARY)
+    
+    # Final smooth
+    mask_uint8 = cv2.GaussianBlur(mask_uint8, (5, 5), 1)
+    
+    return mask_uint8.astype(np.float32) / 255.0
+
+def _fallback_refine(mask01: np.ndarray) -> np.ndarray:
+    """Simple fallback refinement"""
+    mask_uint8 = (mask01 * 255).astype(np.uint8)
+    
+    # Bilateral filter for edge-preserving smoothing
+    mask_uint8 = cv2.bilateralFilter(mask_uint8, 9, 75, 75)
+    
+    # Morphological operations
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
+    mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_OPEN, kernel)
+    
+    # Edge feathering
+    mask_uint8 = cv2.GaussianBlur(mask_uint8, (5, 5), 1)
+    
+    return mask_uint8.astype(np.float32) / 255.0
 
 # ----------------------------------------------------------------------------
 # Compositing
@@ -320,21 +352,31 @@ def replace_background_hq(
     fallback_enabled: bool = True,
     **_compat,
 ) -> np.ndarray:
+    """High-quality background replacement with alpha blending"""
     try:
         H, W = frame.shape[:2]
+        
+        # Resize background if needed
         if background.shape[:2] != (H, W):
             background = cv2.resize(background, (W, H), interpolation=cv2.INTER_LANCZOS4)
-
+        
+        # Ensure mask is properly formatted
         m = _to_mask01(mask01)
-        # Very light feather to hide stair-steps; most shaping already done
+        
+        # Apply slight feather for smooth edges
         m = _feather(m, k=1)
+        
+        # Convert to 3-channel for multiplication
         m3 = np.repeat(m[:, :, None], 3, axis=2)
-
+        
+        # Alpha blending
         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)
+            logger.warning(f"Compositing failed ({e}) – returning original frame")
             return frame
         raise
 
@@ -350,30 +392,30 @@ def validate_video_file(video_path: str) -> Tuple[bool, str]:
         if size == 0:
             return False, "File is empty"
         if size > 2 * 1024 * 1024 * 1024:
-            return False, "File > 2 GB — too large for the Space quota"
+            return False, "File > 2 GB"
 
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            return False, "OpenCV cannot read the file"
+            return False, "Cannot read file"
 
         n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps      = cap.get(cv2.CAP_PROP_FPS)
-        w        = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        h        = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         cap.release()
 
         if n_frames == 0:
             return False, "No frames detected"
         if fps <= 0 or fps > 120:
-            return False, f"Suspicious FPS: {fps}"
+            return False, f"Invalid FPS: {fps}"
         if w <= 0 or h <= 0:
-            return False, "Zero resolution"
+            return False, "Invalid resolution"
         if w > 4096 or h > 4096:
-            return False, f"Resolution {w}×{h} too high (max 4 096²)"
+            return False, f"Resolution {w}×{h} too high"
         if (n_frames / fps) > 300:
             return False, "Video longer than 5 minutes"
 
-        return True, f"OK → {w}×{h}, {fps:.1f} fps, {n_frames/fps:.1f} s"
+        return True, f"OK → {w}×{h}, {fps:.1f} fps, {n_frames/fps:.1f}s"
 
     except Exception as e:
         logger.error(f"validate_video_file: {e}")
@@ -390,4 +432,4 @@ def validate_video_file(video_path: str) -> Tuple[bool, str]:
     "create_professional_background",
     "validate_video_file",
     "PROFESSIONAL_BACKGROUNDS",
-]
+]