Update pipeline/video_pipeline.py

pipeline/video_pipeline.py

@@ -142,12 +142,12 @@ def generate_first_frame_mask(video_path, predictor):
         frame = cv2.resize(frame, (int(w * scale), int(h * scale)))
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
         predictor.set_image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-        masks, _, _ = predictor.predict(
+        masks, scores, _ = predictor.predict(
             point_coords=np.array([[w//2, h//2]]),
             point_labels=np.array([1]),
             multimask_output=True
         )
-    return (masks[np.argmax(predictor.get_mask_scores())] * 255).astype(np.uint8)
+    return (masks[np.argmax(scores)] * 255).astype(np.uint8)
 # --- Temporal Smoothing ---
 def smooth_alpha_video(alpha_path, output_path, window_size=5):
     """Apply temporal smoothing to alpha masks"""