Update app.py

app.py

@@ -409,16 +409,22 @@ def create_mask(self, image_rgb: np.ndarray) -> Optional[np.ndarray]:
 # CHAPTER 6: MATANYONE HANDLER (First-frame PROB mask)
 # ==============================================================================
 
+# =============================================================================
+# CHAPTER 6: MATANYONE HANDLER (Robust unbatched calls + fallbacks)
+# =============================================================================
+
 class MatAnyoneHandler:
     """
-    MatAnyone loader + inference adapter (unbatched I/O).
-
-    This build of MatAnyone appears to add its own batch dimension internally.
-    To avoid 5D tensors, we feed:
-      - image as CHW (3,H,W)
-      - first-frame seed as HW (H,W)  (soft probabilities in [0,1])
-    We try a few safe call signatures to handle minor API differences
-    (with/without `matting=True`, with prob as HW, then 1xHxW).
+    Robust MatAnyone loader + inference adapter.
+
+    What this does:
+      - Prefers UNBATCHED inputs:
+          image: (3, H, W)  float32 in [0,1]
+          prob : (H, W)     float32 in [0,1]  (soft seed from first frame)
+      - Falls back to other safe permutations that some MatAnyone builds expect:
+          prob 1xHxW, numpy HxWx3 + HxW, etc.
+      - Never uses idx_mask/objects (your build asserts on idx mask path).
+      - Squeezes model outputs back to HxW uint8.
     """
     def __init__(self):
         self.core = None
@@ -426,7 +432,7 @@ def __init__(self):
 
     # ----- tensor helpers -----
     def _to_chw_float(self, img01: np.ndarray) -> "torch.Tensor":
-        """img01: HxWx3 in [0,1] -> torch float (3,H,W) on DEVICE (no batch!)."""
+        """img01: HxWx3 in [0,1] -> torch float (3,H,W) on DEVICE (no batch)."""
         assert img01.ndim == 3 and img01.shape[2] == 3, f"Expected HxWx3, got {img01.shape}"
         t = torch.from_numpy(img01.transpose(2, 0, 1)).contiguous().float()  # (3,H,W)
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
@@ -440,7 +446,7 @@ def _prob_hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.T
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
     def _prob_1hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
-        """backup: 1xHxW tensor if a variant expects a leading channel (still unbatched)."""
+        """Optional: 1xHxW (channel-first, still unbatched)."""
         if mask_u8.shape[0] != h or mask_u8.shape[1] != w:
             mask_u8 = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST)
         prob = (mask_u8.astype(np.float32) / 255.0)[None, ...]  # (1,H,W)
@@ -450,8 +456,7 @@ def _prob_1hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.
     def _alpha_to_u8_hw(self, alpha_like) -> np.ndarray:
         """
         Accepts torch / numpy / tuple(list) outputs.
-        Returns uint8 HxW (0..255).
-        Squeezes (1,H,W), (B,1,H,W) etc. down to (H,W) when possible.
+        Returns uint8 HxW (0..255). Squeezes common shapes down to HxW.
         """
         if isinstance(alpha_like, (list, tuple)) and len(alpha_like) > 1:
             alpha_like = alpha_like[1]  # (indices, probs) -> take probs
@@ -466,13 +471,11 @@ def _alpha_to_u8_hw(self, alpha_like) -> np.ndarray:
             a = np.clip(a, 0, 1)
 
         a = np.squeeze(a)
-        # If still 3D like (1,H,W) or (H,W,1) after np.squeeze it should be (H,W)
+        if a.ndim == 3 and a.shape[0] >= 1:  # (1,H,W) -> (H,W)
+            a = a[0]
         if a.ndim != 2:
-            # Try common forms: (1,H,W) or (B,H,W) -> pick first
-            if a.ndim == 3 and a.shape[0] >= 1:
-                a = a[0]
-            if a.ndim != 2:
-                raise ValueError(f"Alpha must be HxW; got {a.shape}")
+            raise ValueError(f"Alpha must be HxW; got {a.shape}")
+
         return np.clip(a * 255.0, 0, 255).astype(np.uint8)
 
     def initialize(self) -> bool:
@@ -514,38 +517,73 @@ def initialize(self) -> bool:
                 state.matanyone_error = f"MatAnyone init error: {e}"
                 return False
 
-    # ----- robust call helpers (UNBATCHED) -----------------------------------
-    def _call_step_seed(self,
-                        img_chw: "torch.Tensor",
-                        prob_hw: "torch.Tensor",
-                        prob_1hw: "torch.Tensor"):
+    # ----- robust call helpers ------------------------------------------------
+    def _try_step_variants_seed(self,
+                                img_chw_t: "torch.Tensor",
+                                img_hwc_np: np.ndarray,
+                                prob_hw_t: "torch.Tensor",
+                                prob_1hw_t: "torch.Tensor",
+                                prob_hw_np: np.ndarray,
+                                prob_hwc1_np: np.ndarray):
         """
-        Try signatures that keep inputs UNBATCHED:
-          1) step(img[3,H,W], prob[H,W])
-          2) step(img[3,H,W], prob[H,W], matting=True)
-          3) step(img[3,H,W], prob[1,H,W])
-          4) step(img[3,H,W], prob[1,H,W], matting=True)
+        Try multiple MatAnyone.step() signatures in a safe order.
+        We avoid idx_mask/objects because this build asserts on idx path.
+
+        Order (from most to least strict about tensors):
+          1) step(CHW, HW)
+          2) step(CHW, HW, matting=True)
+          3) step(CHW, 1HW)
+          4) step(CHW, 1HW, matting=True)
+          5) step(HWC, HW)                 # numpy fallbacks
+          6) step(HWC, HW, matting=True)
+          7) step(HWC, HWC1)
+          8) step(HWC, HWC1, matting=True)
         """
-        try:
-            return self.core.step(img_chw, prob_hw)
-        except (TypeError, RuntimeError):
-            pass
-        try:
-            return self.core.step(img_chw, prob_hw, matting=True)
-        except (TypeError, RuntimeError):
-            pass
-        try:
-            return self.core.step(img_chw, prob_1hw)
-        except (TypeError, RuntimeError):
-            pass
-        return self.core.step(img_chw, prob_1hw, matting=True)
+        trials = [
+            ( (img_chw_t, prob_hw_t), {} ),
+            ( (img_chw_t, prob_hw_t), {"matting": True} ),
+            ( (img_chw_t, prob_1hw_t), {} ),
+            ( (img_chw_t, prob_1hw_t), {"matting": True} ),
+            ( (img_hwc_np, prob_hw_np), {} ),
+            ( (img_hwc_np, prob_hw_np), {"matting": True} ),
+            ( (img_hwc_np, prob_hwc1_np), {} ),
+            ( (img_hwc_np, prob_hwc1_np), {"matting": True} ),
+        ]
+        last_err = None
+        for (args, kwargs) in trials:
+            try:
+                return self.core.step(*args, **kwargs)
+            except Exception as e:
+                last_err = e
+                # Keep trying next variant
+        raise last_err  # bubble up the most informative final error
+
+    def _try_step_variants_noseed(self,
+                                  img_chw_t: "torch.Tensor",
+                                  img_hwc_np: np.ndarray):
+        """
+        Variants when no seed is provided on subsequent frames.
+        """
+        trials = [
+            ( (img_chw_t,), {} ),
+            ( (img_chw_t,), {"matting": True} ),
+            ( (img_hwc_np,), {} ),
+            ( (img_hwc_np,), {"matting": True} ),
+        ]
+        last_err = None
+        for (args, kwargs) in trials:
+            try:
+                return self.core.step(*args, **kwargs)
+            except Exception as e:
+                last_err = e
+        raise last_err
 
     # ----- video matting using first-frame PROB mask --------------------------
     def process_video(self, input_path: str, mask_path: str, output_path: str) -> str:
         """
         Produce a single-channel alpha mp4 matching input fps & size.
 
-        First frame: pass a soft seed prob (HxW) alongside CHW image.
+        First frame: pass a soft seed prob (~HW) alongside the image.
         Remaining frames: call step(image) only.
         """
         if not self.initialized or self.core is None:
@@ -563,13 +601,16 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-        # soft seed prob (unbatched)
+        # soft seed prob
         seed_mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
         if seed_mask is None:
             cap.release()
             raise RuntimeError("Seed mask read failed")
-        prob_hw  = self._prob_hw_from_mask_u8(seed_mask, w, h)   # (H,W)
-        prob_1hw = self._prob_1hw_from_mask_u8(seed_mask, w, h)  # (1,H,W) fallback
+
+        prob_hw_t  = self._prob_hw_from_mask_u8(seed_mask, w, h)    # (H,W) torch
+        prob_1hw_t = self._prob_1hw_from_mask_u8(seed_mask, w, h)   # (1,H,W) torch
+        prob_hw_np = (seed_mask.astype(np.float32) / 255.0)         # (H,W) numpy
+        prob_hwc1_np = prob_hw_np[:, :, None]                        # (H,W,1) numpy
 
         # temp frames
         tmp_dir = TEMP_DIR / f"ma_{int(time.time())}_{random.randint(1000,9999)}"
@@ -583,12 +624,17 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         if not ok or frame_bgr is None:
             cap.release()
             raise RuntimeError("Empty first frame")
-
         frame_rgb01 = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
-        img_chw = self._to_chw_float(frame_rgb01)  # (3,H,W)
+
+        img_chw_t   = self._to_chw_float(frame_rgb01)     # (3,H,W) torch
+        img_hwc_np  = frame_rgb01                          # (H,W,3) numpy
 
         with torch.no_grad():
-            out_prob = self._call_step_seed(img_chw, prob_hw, prob_1hw)
+            out_prob = self._try_step_variants_seed(
+                img_chw_t, img_hwc_np,
+                prob_hw_t, prob_1hw_t,
+                prob_hw_np, prob_hwc1_np
+            )
 
         alpha_u8 = self._alpha_to_u8_hw(out_prob)
         cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
@@ -601,13 +647,11 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
                 break
 
             frame_rgb01 = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
-            img_chw = self._to_chw_float(frame_rgb01)
+            img_chw_t  = self._to_chw_float(frame_rgb01)
+            img_hwc_np = frame_rgb01
 
             with torch.no_grad():
-                try:
-                    out_prob = self.core.step(img_chw)
-                except TypeError:
-                    out_prob = self.core.step(img_chw, matting=True)
+                out_prob = self._try_step_variants_noseed(img_chw_t, img_hwc_np)
 
             alpha_u8 = self._alpha_to_u8_hw(out_prob)
             cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
@@ -635,8 +679,6 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         return str(alpha_path)
 
 
-
-
 # =============================================================================
 # CHAPTER 7: AI BACKGROUNDS
 # =============================================================================