Update app.py

app.py

@@ -409,52 +409,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.
 
-    Strategy:
-    - Image → CHW (3,H,W), no batch/time dims.
-    - Seed mask → try **2D prob** (H,W) first to avoid 5 vs 6-D concat, then fall back to 1xHxW.
-    - Call InferenceCore.step with prob as a **positional** argument.
-    - Try with/without `matting=True` (some builds don't accept it).
-    - Subsequent frames call step(image) with no seed.
+    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.
     """
     def __init__(self):
         self.core = None
         self.initialized = False
 
     # ----- tensor helpers -----
-    def _to_chw_float(self, img01: np.ndarray) -> "torch.Tensor":
-        """img01: HxWx3 in [0,1] -> torch float 3xHxW on DEVICE"""
+    def _to_bchw_float(self, img01: np.ndarray) -> "torch.Tensor":
+        """img01: HxWx3 in [0,1] -> torch float (1,3,H,W) on DEVICE"""
         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_hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
-        """mask_u8: HxW uint8 -> torch float (H,W) on DEVICE, resized if needed"""
+    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"""
         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)  # (H,W)
+        prob = (mask_u8.astype(np.float32) / 255.0)[None, None, ...]  # (1,1,H,W)
         t = torch.from_numpy(prob).contiguous().float()
         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":
-        """mask_u8: HxW uint8 -> torch float (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":
+        """Optional fallback: (H,W) prob (kept on DEVICE)"""
         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)
+        prob = (mask_u8.astype(np.float32) / 255.0)
         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 (K,H,W) by taking the first channel if needed.
-        Also handles MatAnyone tuples/lists like (indices, probs) by taking the 2nd item.
+        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.
         """
         if isinstance(alpha_like, (list, tuple)) and len(alpha_like) > 1:
-            alpha_like = alpha_like[1]  # handle (indices, probs)
+            alpha_like = alpha_like[1]
 
         if isinstance(alpha_like, torch.Tensor):
             t = alpha_like.detach()
@@ -466,11 +468,11 @@ 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 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:
@@ -513,33 +515,31 @@ def initialize(self) -> bool:
                 return False
 
     # ----- robust call helpers ------------------------------------------------
-    def _call_step_seed(self, img_chw: "torch.Tensor",
-                        prob_hw: "torch.Tensor",
-                        prob_1hw: "torch.Tensor"):
+    def _call_step_seed(self,
+                        img_bchw: "torch.Tensor",
+                        prob_b1hw: "torch.Tensor",
+                        prob_hw: "torch.Tensor"):
         """
-        Try a few safe permutations to satisfy different MatAnyone builds.
-        Order chosen to avoid the 5-vs-6D concat error seen in group_modules.py.
+        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)
         """
-        # 1) image (3,H,W), prob (H,W)
         try:
-            return self.core.step(img_chw, prob_hw)
+            return self.core.step(img_bchw, prob_b1hw)
         except (TypeError, RuntimeError):
             pass
-
-        # 2) image (3,H,W), prob (1,H,W)
         try:
-            return self.core.step(img_chw, prob_1hw)
+            return self.core.step(img_bchw, prob_b1hw, matting=True)
         except (TypeError, RuntimeError):
             pass
-
-        # 3) image (3,H,W), prob (H,W), matting kw
         try:
-            return self.core.step(img_chw, prob_hw, matting=True)
+            return self.core.step(img_bchw, prob_hw)
         except (TypeError, RuntimeError):
             pass
-
-        # 4) image (3,H,W), prob (1,H,W), matting kw
-        return self.core.step(img_chw, prob_1hw, matting=True)
+        return self.core.step(img_bchw, prob_hw, matting=True)
 
     # ----- video matting using first-frame PROB mask (PATCHED) ----------------
     def process_video(self, input_path: str, mask_path: str, output_path: str) -> str:
@@ -547,8 +547,7 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         Produce a single-channel alpha mp4 matching input fps & size.
 
         First frame:
-          - Build both (H,W) and (1,H,W) soft prob tensors from SAM2 mask.
-          - Call step(...) via _call_step_seed (tries 2D first to dodge 6D concat).
+          - 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.
         """
@@ -567,13 +566,14 @@ 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
         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)
+        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)
 
         # temp frames
         tmp_dir = TEMP_DIR / f"ma_{int(time.time())}_{random.randint(1000,9999)}"
@@ -589,10 +589,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_chw = self._to_chw_float(frame_rgb01)  # (3,H,W)
+        img_bchw = self._to_bchw_float(frame_rgb01)  # (1,3,H,W)
 
         with torch.no_grad():
-            out_prob = self._call_step_seed(img_chw, prob_hw, prob_1hw)
+            out_prob = self._call_step_seed(img_bchw, prob_b1hw, prob_hw)
 
         alpha_u8 = self._alpha_to_u8_hw(out_prob)
         cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
@@ -605,14 +605,14 @@ 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_bchw = self._to_bchw_float(frame_rgb01)
 
             with torch.no_grad():
                 try:
-                    out_prob = self.core.step(img_chw)           # simplest path
+                    out_prob = self.core.step(img_bchw)
                 except TypeError:
-                    # Extremely old/new variants: try permissive kw
-                    out_prob = self.core.step(img_chw, matting=True)
+                    # very old/new variants: try permissive kw
+                    out_prob = self.core.step(img_bchw, matting=True)
 
             alpha_u8 = self._alpha_to_u8_hw(out_prob)
             cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)