Update utils/segmentation.py

utils/segmentation.py

@@ -51,6 +51,58 @@ class SegmentationError(Exception):
     "SegmentationError",
 ]
 
+# ============================================================================
+# SAM2 TO MATANYONE MASK BRIDGE
+# ============================================================================
+def _sam2_to_matanyone_mask(masks: Any, scores: Any = None) -> np.ndarray:
+    """
+    Convert SAM2 multi-mask output to single best mask for MatAnyone.
+    SAM2 returns (N, H, W) where N is typically 3 masks.
+    We need to return a single (H, W) mask.
+    """
+    if masks is None or len(masks) == 0:
+        raise SegmentationError("No masks returned from SAM2")
+    
+    # Handle torch tensors
+    if isinstance(masks, torch.Tensor):
+        masks = masks.cpu().numpy()
+    if scores is not None and isinstance(scores, torch.Tensor):
+        scores = scores.cpu().numpy()
+    
+    # Ensure we have the right shape
+    if masks.ndim == 4:  # (B, N, H, W)
+        masks = masks[0]  # Take first batch
+    if masks.ndim != 3:  # Should be (N, H, W)
+        raise SegmentationError(f"Unexpected mask shape: {masks.shape}")
+    
+    # Select best mask
+    if scores is not None and len(scores) > 0:
+        best_idx = int(np.argmax(scores))
+    else:
+        # Fallback: pick mask with largest area
+        areas = [np.sum(m > 0.5) for m in masks]
+        best_idx = int(np.argmax(areas))
+    
+    mask = masks[best_idx]
+    
+    # Convert to uint8 binary mask
+    if mask.dtype in (np.float32, np.float64):
+        mask = (mask > 0.5).astype(np.uint8) * 255
+    elif mask.dtype != np.uint8:
+        mask = mask.astype(np.uint8)
+    
+    # Ensure single channel
+    if mask.ndim == 3:
+        mask = mask[:, :, 0] if mask.shape[2] > 1 else mask.squeeze()
+    
+    # Binary threshold
+    _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+    
+    # Verify output shape
+    assert mask.ndim == 2, f"Output mask must be 2D, got shape {mask.shape}"
+    
+    return mask
+
 # ============================================================================
 # MAIN API
 # ============================================================================
@@ -124,10 +176,13 @@ def segment_person_hq_original(image: np.ndarray, predictor: Any, fallback_enabl
                     point_labels=labels,
                     multimask_output=True,
                 )
+            
+            # Use the bridge function to get single best mask
             if masks is not None and len(masks):
-                mask = _process_mask(masks[int(np.argmax(scores))])
+                mask = _sam2_to_matanyone_mask(masks, scores)
                 if _validate_mask_quality(mask, image.shape[:2]):
                     return mask
+                    
         if fallback_enabled:
             return _classical_segmentation_cascade(image)
         raise RuntimeError("SAM2 failed and fallback disabled")
@@ -164,10 +219,9 @@ def _sam2_predict(image: np.ndarray, predictor: Any,
             point_labels=labels,
             multimask_output=True,
         )
-    if masks is None or len(masks) == 0:
-        raise RuntimeError("SAM2 produced no masks")
-    best = masks[int(np.argmax(scores))] if scores is not None else masks[0]
-    return _process_mask(best)
+    
+    # Use the bridge function to convert multi-mask to single mask
+    return _sam2_to_matanyone_mask(masks, scores)
 
 
 def _generate_smart_prompts(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
@@ -275,6 +329,7 @@ def _validate_mask_quality(mask: np.ndarray, shape: Tuple[int,int]) -> bool:
     return MIN_AREA_RATIO <= ratio <= MAX_AREA_RATIO
 
 def _process_mask(mask: np.ndarray) -> np.ndarray:
+    """Legacy mask processor - kept for compatibility but mostly replaced by _sam2_to_matanyone_mask"""
     if mask.dtype in (np.float32, np.float64):
         if mask.max() <= 1.0:
             mask = (mask*255).astype(np.uint8)