Update app.py

app.py

@@ -36,6 +36,57 @@ def filter_kwargs(cls, kwargs):
     valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
     return {k: v for k, v in kwargs.items() if k in valid_params}
 
+
+def load_mask_video(mask_video_path, target_frame_count, target_size, target_fps):
+    """
+    Load a user-supplied mask video.
+
+    Convention:
+      - white / bright pixels = area to inpaint
+      - black pixels = keep unchanged
+      - grayscale masks are supported
+      - the mask video is resized to the source video's resolution
+      - if the mask has fewer frames than the source, an error is raised
+    """
+    if not mask_video_path:
+        return None
+
+    cap = cv2.VideoCapture(mask_video_path)
+    if not cap.isOpened():
+        raise ValueError(f"Could not open mask video: {mask_video_path}")
+
+    masks = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        if (gray.shape[1], gray.shape[0]) != target_size:
+            gray = cv2.resize(
+                gray,
+                target_size,
+                interpolation=cv2.INTER_NEAREST
+            )
+
+        # Normalize to a binary mask. Any non-black pixel is treated as masked.
+        mask = (gray > 127).astype(np.uint8)
+        masks.append(mask)
+
+    cap.release()
+
+    if not masks:
+        raise ValueError("The supplied mask video contains no readable frames.")
+
+    if len(masks) < target_frame_count:
+        raise ValueError(
+            f"Mask video has {len(masks)} frames, but the source video has "
+            f"{target_frame_count} frames."
+        )
+
+    # Extra mask frames are ignored so both videos stay aligned.
+    return masks[:target_frame_count]
+
 from huggingface_hub import snapshot_download
 
 def download_component_subfolder(repo_id, subfolder):
@@ -128,7 +179,7 @@ def get_prompt(click_state, click_input):
 
 @spaces.GPU
 # extract frames from upload video
-def get_frames_from_video(video_input, video_state):
+def get_frames_from_video(video_input, mask_video_input, video_state):
     """
     Args:
         video_path:str
@@ -348,9 +399,11 @@ def inpaint_video(video_state, *_):
     # import pdb;pdb.set_trace()
     frames = video_state["origin_images"]
     masks = video_state["masks"]
-    # masks = masks * 255
     fps = int(video_state["fps"])
 
+    if masks is None or len(masks) != len(frames):
+        raise gr.Error("Mask video is missing or does not match the source video frame count.")
+
     total_frames = len(frames)
     target_frame_count = (total_frames - 1) // 16 * 16 + 1  
     frames = frames[:target_frame_count]  
@@ -456,7 +509,7 @@ description = r"""
 <center></center>
 <b>Official Gradio demo</b> for <a href='https://github.com/Kunbyte-AI/ROSE' target='_blank'><b>Remove Objects with Side Effects in Videos</b></a>.<br>
 🔥 ROSE is a robust inpainting algorithm.<br>
-🤗 Try to drop your video, add the masks and get the the inpainting results!<br>
+🤗 Drop your source video and, optionally, a pre-made mask video, then get the inpainting result!<br>
 """
 
 css = """
@@ -518,9 +571,21 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as iface:
         # input video
         gr.Markdown("## Step1: Upload video")
         with gr.Row(equal_height=True):
-            with gr.Column(scale=2):      
-                video_input = gr.Video(elem_classes="video")
-                extract_frames_button = gr.Button(value="Get video info", interactive=True, variant="primary") 
+            with gr.Column(scale=2):
+                video_input = gr.Video(label="Source video", elem_classes="video")
+                mask_video_input = gr.Video(
+                    label="Mask video (optional)",
+                    elem_classes="video"
+                )
+                gr.Markdown(
+                    "**Mask convention:** white = remove/inpaint, black = keep. "
+                    "The mask video must be frame-aligned with the source video."
+                )
+                extract_frames_button = gr.Button(
+                    value="Get video info",
+                    interactive=True,
+                    variant="primary"
+                )
             with gr.Column(scale=2):
                 run_status = gr.HighlightedText(value=[("",""), ("Try to upload your video and click the Get video info button to get started!", "Normal")],
                                                 color_map={"Normal": "green", "Error": "red", "Clear clicks": "gray", "Add mask": "green", "Remove mask": "red"})
@@ -553,20 +618,20 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as iface:
                 mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask selection", info=".", visible=False)
             
         # output video
-        step3_title = gr.Markdown("---\n## Step3: Track masks and get the inpainting result", visible=False)
+        step3_title = gr.Markdown("---\n## Step3: Get the inpainting result", visible=False)
         with gr.Row(equal_height=True):
             with gr.Column(scale=2):
                 tracking_video_output = gr.Video(visible=False, elem_classes="video")
-                tracking_video_predict_button = gr.Button(value="1. Tracking", visible=False, elem_classes="margin_center")
+                tracking_video_predict_button = gr.Button(value="1. Tracking (optional)", visible=False, elem_classes="margin_center")
             with gr.Column(scale=2):
                 inpaiting_video_output = gr.Video(visible=False, elem_classes="video")
-                inpaint_video_predict_button = gr.Button(value="2. Inpainting", visible=False, elem_classes="margin_center")
+                inpaint_video_predict_button = gr.Button(value="Inpainting", visible=False, elem_classes="margin_center")
 
     # first step: get the video information 
     extract_frames_button.click(
         fn=get_frames_from_video,
         inputs=[
-            video_input, video_state
+            video_input, mask_video_input, video_state
         ],
         outputs=[video_state, video_info, template_frame,
                  image_selection_slider, track_pause_number_slider,point_prompt, clear_button_click, Add_mask_button, template_frame,
@@ -608,14 +673,14 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as iface:
         outputs=[tracking_video_output, video_state, interactive_state, run_status, run_status2]
     )
 
-    # inpaint video from select image and mask
+    # inpaint directly from the supplied mask video
     inpaint_video_predict_button.click(
         fn=inpaint_video,
-        #inputs=[video_state, resize_ratio_number, dilate_radius_number, raft_iter_number, subvideo_length_number, neighbor_length_number, ref_stride_number, mask_dropdown],
-        inputs=[video_state, mask_dropdown],
+        inputs=[video_state],
         outputs=[inpaiting_video_output, run_status, run_status2]
     )
 
+
     # click to get mask
     mask_dropdown.change(
         fn=show_mask,