Update app.py

app.py

@@ -156,6 +156,51 @@ def clear_clicks(video_state):
     video_state["painted_images"] = None
     return Image.fromarray(video_state["origin_images"][0]) if video_state["origin_images"] is not None else None
 
+    
+def load_uploaded_mask_video(video_path, mask_video_path, n_frames, video_state):
+    """Load a user-supplied mask video (white = object) and use it directly,
+    skipping click segmentation and SAM2 tracking."""
+    src = video_state["video_path"] or video_path
+    if mask_video_path is None or src is None:
+        return video_state, None
+
+    # --- load source frames (same logic as track_video) ---
+    vr = VideoReader(src, ctx=cpu(0))
+    images = [vr[i].asnumpy() for i in range(min(len(vr), int(n_frames)))]
+    del vr
+
+    if images[0].shape[0] > images[0].shape[1]:
+        W_ = W
+        H_ = int(W_ * images[0].shape[0] / images[0].shape[1])
+    else:
+        H_ = H
+        W_ = int(H_ * images[0].shape[1] / images[0].shape[0])
+    images = [cv2.resize(img, (W_, H_)) for img in images]
+    video_state["origin_images"] = images
+
+    # --- load mask frames ---
+    mvr = VideoReader(mask_video_path, ctx=cpu(0))
+    mask_frames = []
+    for i in range(min(len(mvr), len(images))):
+        m = mvr[i].asnumpy()
+        if m.ndim == 3:
+            m = cv2.cvtColor(m, cv2.COLOR_RGB2GRAY)
+        m = cv2.resize(m, (W_, H_), interpolation=cv2.INTER_NEAREST)
+        m = (m > 127).astype(np.float32)
+        # match the (H_, W_, 3) shape that track_video produces
+        mask_frames.append(np.repeat(m[..., None], 3, axis=2))
+    del mvr
+
+    while len(mask_frames) < len(images):      # pad if mask video is shorter
+        mask_frames.append(mask_frames[-1].copy())
+    video_state["masks"] = mask_frames
+
+    # preview: overlay first mask on first frame
+    color = np.array(COLOR_PALETTE[0], dtype=np.float32) / 255.0
+    frame = images[0].astype(np.float32) / 255.0
+    painted = (1 - mask_frames[0] * 0.5) * frame + mask_frames[0] * 0.5 * color
+    preview = np.uint8(np.clip(painted * 255, 0, 255))
+    return video_state, Image.fromarray(preview)
 
 def preprocess_for_removal(images, masks):
     out_images = []
@@ -402,6 +447,16 @@ with gr.Blocks() as demo:
             inputs=[dilation_slider, inference_steps_slider, video_state],
             outputs=remove_video
         )
+
+        mask_video_input = gr.Video(label="(Optional) Upload Mask Video — white = object to remove")
+        use_mask_btn = gr.Button("Use Uploaded Mask (skips clicking & tracking)", elem_id="my-btn")
+        
+        use_mask_btn.click(
+            load_uploaded_mask_video,
+            inputs=[video_input, mask_video_input, n_frames_slider, video_state],
+            outputs=[video_state, image_output],
+        )
+        
         get_info_btn.click(get_video_info, inputs=[video_input, video_state], \
                        outputs=image_output)
         image_output.select(fn=segment_frame, inputs=[point_prompt, video_state], outputs=image_output)