Update models/loaders/matanyone_loader.py

models/loaders/matanyone_loader.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
 """
-MatAnyone Model Loader (Hardened)
-- Prevents 5D (B,T,C,H,W) tensors from ever reaching conv2d.
+MatAnyone Model Loader (Hardened v2)
+- Prevents 5D (B,T,C,H,W) tensors from reaching conv2d.
 - Normalizes images to BCHW [B,C,H,W] and masks to B1HW [B,1,H,W].
-- If idx_mask=True, converts masks to integer labels (long) safely.
-- Tries unbatched then batched calls for maximum compatibility.
-- Resizes masks with 'nearest' to preserve label integrity.
-- Includes a debug_shapes() helper for quick diagnostics.
+- idx_mask=True -> integer label map, but final output still a 2-D [H,W] mask for OpenCV.
+- ALWAYS returns a 2-D, contiguous, float32 mask [H,W] to downstream code.
+- Tries unbatched then batched calls; resizes masks with NEAREST to preserve labels.
+- Includes debug_shapes() for quick diagnostics.
 """
 
 import os
@@ -24,9 +24,6 @@
 # ------------------------------- Utilities -------------------------------- #
 
 def _select_device(pref: str) -> str:
-    """
-    Resolve a safe device string. If CUDA not available, fall back to CPU.
-    """
     pref = (pref or "").lower()
     if pref.startswith("cuda"):
         return "cuda" if torch.cuda.is_available() else "cpu"
@@ -36,7 +33,6 @@ def _select_device(pref: str) -> str:
 
 
 def _as_tensor_on_device(x, device: str) -> torch.Tensor:
-    """Convert ndarray or Tensor to torch.Tensor on device."""
     if isinstance(x, torch.Tensor):
         return x.to(device)
     return torch.from_numpy(np.asarray(x)).to(device)
@@ -45,11 +41,7 @@ def _as_tensor_on_device(x, device: str) -> torch.Tensor:
 def _to_bchw(x, device: str, is_mask: bool = False) -> torch.Tensor:
     """
     Normalize input to BCHW (image) or B1HW (mask).
-
     Accepts: HWC, CHW, BCHW, BHWC, BTCHW, BTHWC, TCHW, THWC, HW.
-    - Collapses any time/clip dimension T if present (takes t=0 if T>1).
-    - Images returned as float32 in [0,1], shape [B,C,H,W] (C=3 or 4; C=1 expanded to 3).
-    - Masks  returned as float32 in [0,1], shape [B,1,H,W].
     """
     x = _as_tensor_on_device(x, device)
 
@@ -59,43 +51,39 @@ def _to_bchw(x, device: str, is_mask: bool = False) -> torch.Tensor:
     elif x.dtype in (torch.int16, torch.int32, torch.int64):
         x = x.float()
 
-    # 5D: [B,T,C,H,W] or [B,T,H,W,C]  -> take first frame
+    # 5D: [B,T,C,H,W] or [B,T,H,W,C] -> take first frame
     if x.ndim == 5:
         B, T = x.shape[0], x.shape[1]
-        x = x[:, 0] if T > 0 else x.squeeze(1)  # -> [B,C,H,W] or [B,H,W,C]
+        x = x[:, 0] if T > 0 else x.squeeze(1)
 
     # 4D
     if x.ndim == 4:
-        # If BHWC, permute to BCHW
         if x.shape[-1] in (1, 3, 4) and x.shape[1] not in (1, 3, 4):
-            x = x.permute(0, 3, 1, 2).contiguous()
+            x = x.permute(0, 3, 1, 2).contiguous()  # BHWC -> BCHW
 
     # 3D
     elif x.ndim == 3:
-        # HWC -> CHW
-        if x.shape[-1] in (1, 3, 4):
+        if x.shape[-1] in (1, 3, 4):               # HWC -> CHW
             x = x.permute(2, 0, 1).contiguous()
-        x = x.unsqueeze(0)  # -> BCHW
+        x = x.unsqueeze(0)                          # -> BCHW
 
     # 2D
     elif x.ndim == 2:
         if is_mask:
-            x = x.unsqueeze(0).unsqueeze(0)  # -> B1HW
+            x = x.unsqueeze(0).unsqueeze(0)         # -> B1HW
         else:
-            x = x.unsqueeze(0).unsqueeze(0)  # 1,1,H,W
-            x = x.repeat(1, 3, 1, 1)        # 1,3,H,W
+            x = x.unsqueeze(0).unsqueeze(0)         # 1,1,H,W
+            x = x.repeat(1, 3, 1, 1)                # 1,3,H,W
 
     else:
         raise ValueError(f"Unsupported tensor ndim={x.ndim} for normalization")
 
-    # Now x should be BCHW
+    # Finalize channels / clamp
     if is_mask:
-        # Ensure single-channel
         if x.shape[1] > 1:
             x = x[:, :1]
         x = x.clamp_(0.0, 1.0).to(torch.float32)
     else:
-        # Ensure reasonable channels
         C = x.shape[1]
         if C == 1:
             x = x.repeat(1, 3, 1, 1)
@@ -107,9 +95,7 @@ def _to_bchw(x, device: str, is_mask: bool = False) -> torch.Tensor:
 
 
 def _resize_mask_to(img_bchw: torch.Tensor, mask_b1hw: torch.Tensor) -> torch.Tensor:
-    """
-    Ensure mask spatial dims match image. Use NEAREST to keep labels crisp.
-    """
+    """Ensure mask spatial dims match image. Use NEAREST to keep labels crisp."""
     if img_bchw.shape[-2:] == mask_b1hw.shape[-2:]:
         return mask_b1hw
     import torch.nn.functional as F
@@ -117,9 +103,7 @@ def _resize_mask_to(img_bchw: torch.Tensor, mask_b1hw: torch.Tensor) -> torch.Te
 
 
 def debug_shapes(tag: str, image, mask) -> None:
-    """
-    Quick diagnostics: logs shape/dtype/min/max for image/mask.
-    """
+    """Quick diagnostics: logs shape/dtype/min/max for image/mask."""
     def _info(name, t):
         try:
             tt = torch.as_tensor(t)
@@ -129,17 +113,69 @@ def _info(name, t):
                         f"min={mn:.4f} max={mx:.4f}")
         except Exception as e:
             logger.info(f"[{tag}:{name}] type={type(t)} err={e}")
-
     _info("image", image)
     _info("mask",  mask)
 
 
+def _to_2d_numpy_mask(x) -> np.ndarray:
+    """
+    Convert any tensor/ndarray mask to a 2-D contiguous float32 array [H,W] in [0,1].
+    Handles inputs like: B1HW, BCHW, 1HW, CHW, HWC, HW, etc.
+    """
+    if isinstance(x, torch.Tensor):
+        t = x.detach()
+    else:
+        t = torch.as_tensor(x)
+
+    # Bring to float in [0,1] if likely 0..255
+    if t.dtype == torch.uint8:
+        t = t.float().div_(255.0)
+    elif t.dtype in (torch.int16, torch.int32, torch.int64):
+        t = t.float()
+    else:
+        t = t.float()
+
+    # Reduce dimensions to [H,W]
+    if t.ndim == 4:           # e.g., [B, C, H, W]
+        if t.shape[0] > 1:
+            t = t[0]
+        # now [C,H,W]
+        if t.shape[0] > 1:    # multiple channels -> take first (or could mean)
+            t = t[0]
+        else:
+            t = t[0]          # squeeze channel -> [H,W]
+    elif t.ndim == 3:
+        # Could be [1,H,W], [C,H,W], or [H,W,1]
+        if t.shape[0] in (1, 3, 4):     # CHW/1HW
+            t = t[0]                    # -> [H,W] (first channel)
+        elif t.shape[-1] == 1:          # HWC with single channel
+            t = t[..., 0]               # -> [H,W]
+        else:
+            # Unknown 3D -> take first slice
+            t = t[0]
+    elif t.ndim == 2:
+        pass  # already [H,W]
+    else:
+        # Any other: try to squeeze to 2-D
+        t = t.squeeze()
+        if t.ndim != 2:
+            # fallback to a tiny neutral mask
+            h = int(t.shape[-2]) if t.ndim >= 2 else 512
+            w = int(t.shape[-1]) if t.ndim >= 2 else 512
+            t = torch.full((h, w), 0.5, dtype=torch.float32)
+
+    # Clamp and convert to contiguous numpy
+    t = t.clamp_(0.0, 1.0)
+    m = t.cpu().numpy().astype(np.float32)
+    return np.ascontiguousarray(m)
+
+
 # --------------------------- Boundary Wrapper ------------------------------ #
 
 class _MatAnyoneWrapper:
     """
     Thin, defensive wrapper around the MatAnyone InferenceCore.
-    Normalizes inputs at the boundary so the core never sees >4D tensors.
+    Normalizes inputs at the boundary and always outputs a 2-D mask for OpenCV.
     """
 
     def __init__(self, core: Any, device: str):
@@ -153,12 +189,6 @@ def __init__(self, core: Any, device: str):
         except Exception as e:
             logger.debug(f"MatAnyone core .to({self.device}) not applied: {e}")
 
-    @staticmethod
-    def _to_numpy(x):
-        if isinstance(x, torch.Tensor):
-            return x.detach().cpu().numpy()
-        return np.asarray(x)
-
     def _normalize_pair(
         self, image, mask, idx_mask: bool
     ) -> Tuple[torch.Tensor, torch.Tensor, bool]:
@@ -169,23 +199,21 @@ def _normalize_pair(
 
     def __call__(self, image, mask, idx_mask: bool = False, **kwargs):
         """
-        Preferred entry: handles normalization and robust call patterns.
+        Entry point: returns a 2-D float32 mask [H,W] for downstream OpenCV.
         """
         img_bchw, msk_b1hw, idx_mask = self._normalize_pair(image, mask, idx_mask)
 
-        # Special handling for idx_mask: convert to integer label map.
+        # idx_mask path -> integer labels; still output 2-D for downstream
         if idx_mask:
-            # Threshold -> {0,1} long; squeeze channel
             m_bhw = (msk_b1hw > 0.5).long()[:, 0]  # [B,H,W]
-            # Try unbatched first if B==1
+            # Try unbatched if B==1
             if img_bchw.shape[0] == 1:
                 img_chw = img_bchw[0]  # [C,H,W]
                 m_hw   = m_bhw[0]      # [H,W]
-                # Prefer step(image, mask, idx_mask=True)
                 try:
                     if hasattr(self.core, "step"):
                         out = self.core.step(image=img_chw, mask=m_hw, idx_mask=True, **kwargs)
-                        return self._to_numpy(out)
+                        return _to_2d_numpy_mask(out)
                 except Exception as e_unbatched_idx:
                     logger.debug(f"MatAnyone unbatched idx_mask step() failed: {e_unbatched_idx}")
             # Batched fallback
@@ -194,21 +222,20 @@ def __call__(self, image, mask, idx_mask: bool = False, **kwargs):
                     if hasattr(self.core, method_name):
                         method = getattr(self.core, method_name)
                         out = method(image=img_bchw, mask=m_bhw, idx_mask=True, **kwargs)
-                        return self._to_numpy(out)
+                        return _to_2d_numpy_mask(out)
                 except Exception as e_batched_idx:
                     logger.debug(f"MatAnyone {method_name} idx_mask batched call failed: {e_batched_idx}")
 
             logger.warning("MatAnyone idx_mask calls failed; returning integer mask as fallback.")
-            return self._to_numpy(m_bhw if m_bhw.shape[0] > 1 else m_bhw[0])
+            return _to_2d_numpy_mask(m_bhw)
 
-        # Non-index soft/binary mask path
+        # Non-index mask path (soft/binary)
         try:
-            # Try unbatched first (common CHW / 1HW)
             if hasattr(self.core, "step") and img_bchw.shape[0] == 1:
                 img_chw = img_bchw[0]        # [C,H,W]
                 m_1hw   = msk_b1hw[0]        # [1,H,W]
                 out = self.core.step(image=img_chw, mask=m_1hw, idx_mask=False, **kwargs)
-                return self._to_numpy(out)
+                return _to_2d_numpy_mask(out)
         except Exception as e_unbatched:
             logger.debug(f"MatAnyone unbatched step() failed: {e_unbatched}")
 
@@ -218,13 +245,13 @@ def __call__(self, image, mask, idx_mask: bool = False, **kwargs):
                 if hasattr(self.core, method_name):
                     method = getattr(self.core, method_name)
                     out = method(image=img_bchw, mask=msk_b1hw, idx_mask=False, **kwargs)
-                    return self._to_numpy(out)
+                    return _to_2d_numpy_mask(out)
             except Exception as e_batched:
                 logger.debug(f"MatAnyone {method_name} batched call failed: {e_batched}")
 
         logger.warning("MatAnyone calls failed; returning input mask as fallback.")
-        return self._to_numpy(msk_b1hw.squeeze(1))  # [B,H,W] or [H,W] if squeezed
-
+        # Return a valid 2-D mask even on total failure
+        return _to_2d_numpy_mask(msk_b1hw)
 
 # ------------------------------- Loader ----------------------------------- #
 
@@ -243,8 +270,7 @@ def __init__(self, device: str = "cuda", cache_dir: str = "./checkpoints/matanyo
     def load(self) -> Optional[Any]:
         """
         Load MatAnyone model and return a callable wrapper.
-        Returns:
-            _MatAnyoneWrapper or None
+        Returns: _MatAnyoneWrapper or None
         """
         logger.info(f"Loading MatAnyone model: {self.model_id} (device={self.device})")
 
@@ -272,9 +298,7 @@ def load(self) -> Optional[Any]:
         return None
 
     def _load_official(self) -> Optional[Any]:
-        """
-        Load using the official MatAnyone API and wrap with boundary normalizer.
-        """
+        """Load using the official MatAnyone API and wrap with boundary normalizer."""
         try:
             from matanyone import InferenceCore  # type: ignore
         except Exception as e:
@@ -287,28 +311,15 @@ def _load_official(self) -> Optional[Any]:
 
     def _load_fallback(self) -> Optional[Any]:
         """Create a minimal fallback that smooths/returns the mask."""
-
         class _FallbackCore:
             def step(self, image, mask, idx_mask: bool = False, **kwargs):
-                # Convert mask to numpy
-                if isinstance(mask, torch.Tensor):
-                    mask_np = mask.detach().cpu().numpy()
-                else:
-                    mask_np = np.asarray(mask)
+                # Convert to 2-D numpy mask as final step
+                m2d = _to_2d_numpy_mask(mask)
                 try:
                     import cv2
-                    if mask_np.ndim == 2:
-                        return cv2.GaussianBlur(mask_np, (5, 5), 1.0)
-                    if mask_np.ndim == 3:
-                        # Handle CHW-style smoothing (per-channel)
-                        if mask_np.shape[0] in (1, 3, 4):
-                            sm = np.empty_like(mask_np)
-                            for i in range(mask_np.shape[0]):
-                                sm[i] = cv2.GaussianBlur(mask_np[i], (5, 5), 1.0)
-                            return sm
-                    return mask_np
+                    return cv2.GaussianBlur(m2d, (5, 5), 1.0)
                 except Exception:
-                    return mask_np
+                    return m2d
 
             def process(self, image, mask, **kwargs):
                 return self.step(image, mask, **kwargs)
@@ -320,7 +331,6 @@ def process(self, image, mask, **kwargs):
     # --------------------------- Housekeeping --------------------------- #
 
     def cleanup(self):
-        """Clean up resources."""
         if self.model:
             try:
                 del self.model
@@ -331,7 +341,6 @@ def cleanup(self):
             torch.cuda.empty_cache()
 
     def get_info(self) -> Dict[str, Any]:
-        """Get loader information."""
         return {
             "loaded": self.model is not None,
             "model_id": self.model_id,
@@ -340,7 +349,6 @@ def get_info(self) -> Dict[str, Any]:
             "model_type": type(self.model).__name__ if self.model else None,
         }
 
-    # Optional: instance-level shape debugging hook
+    # Optional: instance-level shape debugging
     def debug_shapes(self, image, mask, tag: str = ""):
         debug_shapes(tag, image, mask)
-