Create utils/mask_bridge.py

utils/mask_bridge.py

@@ -1,60 +1,42 @@
-#!/usr/bin/env python3
-"""
-mask_bridge.py - SAM2 to MatAnyone mask conversion
-Handles shape/dtype/device normalization between models
-"""
-
+# utils/mask_bridge.py
+from __future__ import annotations
 import torch
-import math
-from typing import Optional, Tuple
 
 def log_shape(tag: str, t: torch.Tensor) -> None:
-    """Debug logging for tensor shapes and values"""
-    mn = float(t.min()) if t.numel() else math.nan
-    mx = float(t.max()) if t.numel() else math.nan
-    print(f"{tag}: shape={tuple(t.shape)} dtype={t.dtype} device={t.device} range=[{mn:.3f},{mx:.3f}]")
+    try:
+        mn = float(t.min()) if t.numel() else float("nan")
+        mx = float(t.max()) if t.numel() else float("nan")
+        print(f"[mask_bridge] {tag}: shape={tuple(t.shape)} dtype={t.dtype} device={t.device} "
+              f"range=[{mn:.4f},{mx:.4f}]")
+    except Exception:
+        pass
 
 def sam2_to_matanyone_mask(
-    sam2_masks: torch.Tensor,         # shape: (B, M, H, W) from SAM2 post_process_masks
-    iou_scores: Optional[torch.Tensor] = None,  # optional, (B, M)
-    threshold: float = 0.5,           # binarization for hard mask
-    return_mode: str = "single",      # "single" → (1,H,W); "multi" → (C,H,W)
-    keep_soft: bool = False,          # if True → soft [0,1] alpha channel
+    sam2_masks: torch.Tensor,         # (B,M,H,W) after post_process
+    iou_scores: torch.Tensor | None,  # (B,M) or None
+    threshold: float = 0.5,
+    return_mode: str = "single",      # "single"→(1,H,W) or "multi"→(C,H,W)
+    keep_soft: bool = False,
 ) -> torch.Tensor:
-    """
-    Convert SAM2 output masks to MatAnyone-ready format.
-    
-    Returns a MatAnyone-ready tensor on the same device:
-      - "single": (1,H,W) float32 in [0,1]
-      - "multi":  (C,H,W) float32 in [0,1]
-    """
     assert sam2_masks.ndim == 4, f"Expect (B,M,H,W). Got {tuple(sam2_masks.shape)}"
     B, M, H, W = sam2_masks.shape
-    assert B == 1, "We pass one frame to build first-frame mask."
-
-    masks = sam2_masks[0]  # (M,H,W)
+    assert B == 1, "Bridge expects B=1 for first-frame bootstrapping"
 
-    # Choose best mask
-    if iou_scores is not None and iou_scores.ndim == 2:
+    candidates = sam2_masks[0]  # (M,H,W)
+    if iou_scores is not None and iou_scores.ndim == 2 and iou_scores.shape[0] == 1:
         best_idx = int(torch.argmax(iou_scores[0]).item())
     else:
-        # Fallback: pick the mask with largest foreground area
-        areas = masks.sum(dim=(1,2))
+        areas = candidates.sum(dim=(-2,-1))
         best_idx = int(torch.argmax(areas).item())
 
     if return_mode == "multi":
-        out = masks
+        out = candidates  # (M,H,W) treat as (C,H,W)
     else:
-        out = masks[best_idx:best_idx+1]  # (1,H,W)
+        out = candidates[best_idx:best_idx+1]  # (1,H,W)
 
-    # Ensure float32 [0,1]
-    out = out.to(dtype=torch.float32)
+    out = out.to(torch.float32)
     if not keep_soft:
         out = (out >= threshold).float()
-
-    # Final sanity: contiguous, shapes
-    out = out.contiguous()
-    assert out.ndim == 3, f"Expect (C,H,W); got {tuple(out.shape)}"
-    assert out.shape[0] >= 1, f"Need at least 1 channel; got {out.shape[0]}"
-    
-    return out  # (C,H,W) float32 [0,1]
+    out = out.clamp_(0.0, 1.0).contiguous()
+    log_shape("sam2→mat.mask", out)
+    return out