Update app.py

app.py

@@ -407,56 +407,54 @@ def create_mask(self, image_rgb: np.ndarray) -> Optional[np.ndarray]:
 
 # =============================================================================
 # CHAPTER 6: MATANYONE HANDLER (First-frame PROB mask)
-# ============================================================================
-
+# ==============================================================================
 
 class MatAnyoneHandler:
     """
-    MatAnyone loader + inference adapter.
-
-    Key decisions:
-    - Image is fed as Bx3xHxW (B=1).
-    - First-frame seed is *soft probability* Bx1xHxW in [0,1].
-    - We try multiple safe call patterns (with/without `matting=True`, with prob
-      as Bx1xHxW first, then HW) to satisfy different MatAnyone builds.
-    - Subsequent frames: step(image) with no seed.
+    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).
     """
     def __init__(self):
         self.core = None
         self.initialized = False
 
     # ----- tensor helpers -----
-    def _to_bchw_float(self, img01: np.ndarray) -> "torch.Tensor":
-        """img01: HxWx3 in [0,1] -> torch float (1,3,H,W) on DEVICE"""
+    def _to_chw_float(self, img01: np.ndarray) -> "torch.Tensor":
+        """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)
-        t = t.unsqueeze(0)  # (1,3,H,W)
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
-    def _prob_b1hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
-        """mask_u8: HxW uint8 -> torch float (1,1,H,W) on DEVICE, resized if needed"""
+    def _prob_hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
+        """mask_u8: HxW -> torch float (H,W) in [0,1] on DEVICE (no batch, no channel)."""
         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, None, ...]  # (1,1,H,W)
+        prob = (mask_u8.astype(np.float32) / 255.0)  # (H,W)
         t = torch.from_numpy(prob).contiguous().float()
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
-    def _prob_hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
-        """Optional fallback: (H,W) prob (kept on DEVICE)"""
+    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)."""
         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)
+        prob = (mask_u8.astype(np.float32) / 255.0)[None, ...]  # (1,H,W)
         t = torch.from_numpy(prob).contiguous().float()
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
     def _alpha_to_u8_hw(self, alpha_like) -> np.ndarray:
         """
-        Accepts torch Tensor or numpy-like. Returns uint8 HxW (0..255).
-        Handles (H,W), (1,H,W), or (B,1,H,W)/(B,H,W) by squeezing → (H,W).
-        Also handles tuples/lists (indices, probs) by taking the 2nd item.
+        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.
         """
         if isinstance(alpha_like, (list, tuple)) and len(alpha_like) > 1:
-            alpha_like = alpha_like[1]
+            alpha_like = alpha_like[1]  # (indices, probs) -> take probs
 
         if isinstance(alpha_like, torch.Tensor):
             t = alpha_like.detach()
@@ -468,11 +466,13 @@ def _alpha_to_u8_hw(self, alpha_like) -> np.ndarray:
             a = np.clip(a, 0, 1)
 
         a = np.squeeze(a)
-        # Accept (H,W) only at this point
-        if a.ndim == 3 and a.shape[0] >= 1:
-            a = a[0]
+        # If still 3D like (1,H,W) or (H,W,1) after np.squeeze it should be (H,W)
         if a.ndim != 2:
-            raise ValueError(f"Alpha must be HxW; got {a.shape}")
+            # 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}")
         return np.clip(a * 255.0, 0, 255).astype(np.uint8)
 
     def initialize(self) -> bool:
@@ -514,42 +514,39 @@ def initialize(self) -> bool:
                 state.matanyone_error = f"MatAnyone init error: {e}"
                 return False
 
-    # ----- robust call helpers ------------------------------------------------
+    # ----- robust call helpers (UNBATCHED) -----------------------------------
     def _call_step_seed(self,
-                        img_bchw: "torch.Tensor",
-                        prob_b1hw: "torch.Tensor",
-                        prob_hw: "torch.Tensor"):
+                        img_chw: "torch.Tensor",
+                        prob_hw: "torch.Tensor",
+                        prob_1hw: "torch.Tensor"):
         """
-        Try permutations that keep ranks aligned and avoid 5-vs-6D mismatches.
-        Order:
-          1) step(img[B,3,H,W], prob[B,1,H,W])
-          2) step(img[B,3,H,W], prob[B,1,H,W], matting=True)
-          3) step(img[B,3,H,W], prob[H,W])
-          4) step(img[B,3,H,W], prob[H,W], matting=True)
+        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:
-            return self.core.step(img_bchw, prob_b1hw)
+            return self.core.step(img_chw, prob_hw)
         except (TypeError, RuntimeError):
             pass
         try:
-            return self.core.step(img_bchw, prob_b1hw, matting=True)
+            return self.core.step(img_chw, prob_hw, matting=True)
         except (TypeError, RuntimeError):
             pass
         try:
-            return self.core.step(img_bchw, prob_hw)
+            return self.core.step(img_chw, prob_1hw)
         except (TypeError, RuntimeError):
             pass
-        return self.core.step(img_bchw, prob_hw, matting=True)
+        return self.core.step(img_chw, prob_1hw, matting=True)
 
-    # ----- video matting using first-frame PROB mask (PATCHED) ----------------
+    # ----- 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:
-          - Build (1,1,H,W) seed prob from SAM2 and pass to step via _call_step_seed.
-        Remaining frames:
-          - Call step(image) with no seed.
+        First frame: pass a soft seed prob (HxW) alongside CHW image.
+        Remaining frames: call step(image) only.
         """
         if not self.initialized or self.core is None:
             raise RuntimeError("MatAnyone not initialized")
@@ -566,14 +563,13 @@ 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
+        # soft seed prob (unbatched)
         seed_mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
         if seed_mask is None:
             cap.release()
             raise RuntimeError("Seed mask read failed")
-
-        prob_b1hw = self._prob_b1hw_from_mask_u8(seed_mask, w, h)  # (1,1,H,W)
-        prob_hw   = self._prob_hw_from_mask_u8(seed_mask, w, h)    # (H,W)
+        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
 
         # temp frames
         tmp_dir = TEMP_DIR / f"ma_{int(time.time())}_{random.randint(1000,9999)}"
@@ -589,10 +585,10 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
             raise RuntimeError("Empty first frame")
 
         frame_rgb01 = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
-        img_bchw = self._to_bchw_float(frame_rgb01)  # (1,3,H,W)
+        img_chw = self._to_chw_float(frame_rgb01)  # (3,H,W)
 
         with torch.no_grad():
-            out_prob = self._call_step_seed(img_bchw, prob_b1hw, prob_hw)
+            out_prob = self._call_step_seed(img_chw, prob_hw, prob_1hw)
 
         alpha_u8 = self._alpha_to_u8_hw(out_prob)
         cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
@@ -605,14 +601,13 @@ 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_bchw = self._to_bchw_float(frame_rgb01)
+            img_chw = self._to_chw_float(frame_rgb01)
 
             with torch.no_grad():
                 try:
-                    out_prob = self.core.step(img_bchw)
+                    out_prob = self.core.step(img_chw)
                 except TypeError:
-                    # very old/new variants: try permissive kw
-                    out_prob = self.core.step(img_bchw, matting=True)
+                    out_prob = self.core.step(img_chw, matting=True)
 
             alpha_u8 = self._alpha_to_u8_hw(out_prob)
             cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
@@ -640,6 +635,8 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         return str(alpha_path)
 
 
+
+
 # =============================================================================
 # CHAPTER 7: AI BACKGROUNDS
 # =============================================================================