Update pipeline/video_pipeline.py

pipeline/video_pipeline.py

@@ -105,51 +105,85 @@ def _normalize_input(inp, work_dir: Path) -> str:
         raise TypeError(f"Unsupported input: {type(inp)}")
     return str(target)
 
-# --- SAM2 Mask Generation (now returns mask at ORIGINAL video size) ---
-def generate_first_frame_mask(video_path, predictor):
-    """Generate segmentation mask for the first frame using SAM2, then upsample to original size."""
+# --- SAM2 Mask Generation (multi-frame, returns mask at ORIGINAL video size) ---
+def generate_first_frame_mask(video_path, predictor, num_frames: int = 3):
+    """
+    Build a robust seed mask by running SAM2 on the first N frames (default 3),
+    upsampling each mask back to the ORIGINAL video resolution, and combining
+    them by majority vote. Output is a uint8 mask in {0, 255} at (orig_h, orig_w).
+    """
+    import math
+
     cap = cv2.VideoCapture(video_path)
-    ret, frame = cap.read()
-    cap.release()
+    frames = []
+
+    # Grab the first frame to establish original size
+    ret, first = cap.read()
     if not ret:
+        cap.release()
         raise ValueError("Failed to read video frame")
+    orig_h, orig_w = first.shape[:2]
+    frames.append(first)
 
-    orig_h, orig_w = frame.shape[:2]
-
-    # Downscale for inference if needed (≤1080 on the longer side)
-    scaled = frame
-    if max(orig_h, orig_w) > 1080:
-        scale = 1080 / max(orig_h, orig_w)
-        new_w, new_h = int(orig_w * scale), int(orig_h * scale)
-        scaled = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
-        logger.info(f"[sam2] Resized for SAM2: {orig_w}x{orig_h} -> {new_w}x{new_h}")
-    else:
-        new_h, new_w = orig_h, orig_w
-        logger.info(f"[sam2] Using original size for SAM2: {orig_w}x{orig_h}")
+    # Read up to (num_frames-1) more initial frames
+    for _ in range(1, max(1, num_frames)):
+        ret, f = cap.read()
+        if not ret:
+            break
+        frames.append(f)
+    cap.release()
 
-    # Predict mask on the (possibly) downscaled frame
+    masks_fullres = []
     with torch.inference_mode(), torch.autocast(
         "cuda", dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
     ):
-        predictor.set_image(cv2.cvtColor(scaled, cv2.COLOR_BGR2RGB))
-        h_s, w_s = scaled.shape[:2]
-        masks, scores, _ = predictor.predict(
-            point_coords=np.array([[w_s // 2, h_s // 2]]),
-            point_labels=np.array([1]),
-            multimask_output=True,
-        )
-
-    # Pick best mask and convert to 0/255 uint8
-    mask_small = (masks[np.argmax(scores)] * 255).astype(np.uint8)
-
-    # Upsample back to ORIGINAL video resolution (nearest-neighbor to preserve edges)
-    if (new_w, new_h) != (orig_w, orig_h):
-        mask_full = cv2.resize(mask_small, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
-        logger.info(f"[sam2] Upsampled mask to original: {new_w}x{new_h} -> {orig_w}x{orig_h}")
-    else:
-        mask_full = mask_small
+        for idx, frame in enumerate(frames):
+            h, w = frame.shape[:2]
+            # Downscale for inference if needed (≤1080 on the long side)
+            if max(h, w) > 1080:
+                scale = 1080 / max(h, w)
+                new_w, new_h = int(w * scale), int(h * scale)
+                scaled = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
+                logger.info(f"[sam2] f{idx}: resized for SAM2 {w}x{h} -> {new_w}x{new_h}")
+            else:
+                new_h, new_w = h, w
+                scaled = frame
+                logger.info(f"[sam2] f{idx}: using original size for SAM2 {w}x{h}")
+
+            # Run SAM2 on the (possibly) downscaled frame
+            predictor.set_image(cv2.cvtColor(scaled, cv2.COLOR_BGR2RGB))
+            masks, scores, _ = predictor.predict(
+                point_coords=np.array([[new_w // 2, new_h // 2]]),
+                point_labels=np.array([1]),
+                multimask_output=True,
+            )
+            mask_small = masks[np.argmax(scores)]  # float 0..1
+
+            # Upsample the mask back to ORIGINAL resolution with NEAREST (no blur)
+            if (new_w, new_h) != (orig_w, orig_h):
+                mask_full = cv2.resize(
+                    mask_small.astype(np.float32),
+                    (orig_w, orig_h),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+                logger.info(f"[sam2] f{idx}: upsampled mask -> {orig_w}x{orig_h}")
+            else:
+                mask_full = mask_small.astype(np.float32)
+
+            masks_fullres.append(mask_full)
+
+    if not masks_fullres:
+        raise RuntimeError("SAM2 produced no masks")
+
+    # Majority vote across frames
+    stack = np.stack(masks_fullres, axis=0)  # (N, H, W), values 0..1
+    required = (len(masks_fullres) + 1) // 2  # ceil(N/2)
+    vote = (np.sum(stack > 0.5, axis=0) >= required).astype(np.uint8) * 255
+
+    logger.info(f"[sam2] multi-frame seed: N={len(masks_fullres)}, "
+                f"orig_size={orig_w}x{orig_h}, majority={required}/{len(masks_fullres)}")
 
-    return mask_full
+    return vote
 
 # --- Temporal Smoothing ---
 def smooth_alpha_video(alpha_path, output_path, window_size=5):
@@ -244,7 +278,7 @@ def stage1_create_transparent_video(input_file, sam2_predictor, matanyone_proces
                 logger.warning("[stage1] Audio extraction failed, continuing without audio")
                 audio_path = None
 
-            # 2) First-frame mask via SAM2 (now at original size)
+            # 2) First-frame mask via SAM2 (now multi-frame at original size)
             mask = generate_first_frame_mask(input_path, sam2_predictor)
             mask_path = str(temp_dir / "mask.png")
             ok = cv2.imwrite(mask_path, mask)
@@ -255,7 +289,7 @@ def stage1_create_transparent_video(input_file, sam2_predictor, matanyone_proces
             # 3) MatAnyone in a thread with timeout (prevents silent hangs)
             logger.info(
                 f"[stage1] Starting MatAnyone.process_video "
-                f"(input_path={input_path}, mask_path={mask_path}, output_path={temp_dir}, max_size=720)"
+                f"(input_path={input_path}, mask_path={mask_path}, output_path={temp_dir,}, max_size=720)"
             )
 
             result_holder = {"ok": False, "fg": None, "alpha": None, "exc": None}