Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -155,6 +155,23 @@ def get_end_number(track_pause_number_slider, video_state, interactive_state):
 
     return video_state["painted_images"][track_pause_number_slider],interactive_state, operation_log, operation_log
 
+def draw_refine(video_state, template_frame):
+    if isinstance(template_frame, dict) and "mask" in template_frame:
+        mask = template_frame["mask"]
+        if mask is not None:
+            mask = (np.array(mask) > 127).astype(np.uint8)
+            painted_image = mask_painter(
+                video_state["origin_images"][video_state["select_frame_number"]].copy(),
+                mask,
+                mask_color=1
+            )
+            video_state["masks"][video_state["select_frame_number"]] = mask
+            video_state["painted_images"][video_state["select_frame_number"]] = painted_image
+            operation_log = [("Manual mask added by drawing", "Normal")]
+            return painted_image, video_state, operation_log, operation_log
+    operation_log = [("No mask drawn", "Error")]
+    return video_state["painted_images"][video_state["select_frame_number"]], video_state, operation_log, operation_log
+
 # use sam to get the mask
 def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr.SelectData):
     """
@@ -622,13 +639,17 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as iface:
         step2_title = gr.Markdown("---\n## Step2: Add masks", visible=False)
         with gr.Row(equal_height=True):
             with gr.Column(scale=2):
-                template_frame = gr.Image(type="pil",interactive=True, elem_id="template_frame", visible=False, elem_classes="image")
+                template_frame = gr.Image(type="numpy",tool="sketch",brush_radius=4,label="Draw mask manually or click to add",interactive=True,visible=False,elem_id="template_frame",elem_classes="image")
                 image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track start frame", visible=False)
                 track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track end frame", visible=False)
             with gr.Column(scale=2, elem_classes="jc_center"):
-                run_status2 = gr.HighlightedText(value=[("",""), ("Try to upload your video and click the Get video info button to get started! (Kindly ensure that the uploaded video consists of fewer than 500 frames in total)", "Normal")],
-                                                 color_map={"Normal": "green", "Error": "red", "Clear clicks": "gray", "Add mask": "green", "Remove mask": "red"},
-                                                 visible=False)
+                mask_mode_selector = gr.Radio(
+                choices=["Click", "Draw"],
+                value="Click",
+                label="Mask input mode",
+                interactive=True,
+                visible=False
+                )
                 with gr.Column():
                     point_prompt = gr.Radio(
                         choices=["Positive", "Negative"],
@@ -675,11 +696,12 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as iface:
     
     # click select image to get mask using sam
     template_frame.select(
-        fn=sam_refine,
-        inputs=[video_state, point_prompt, click_state, interactive_state],
-        outputs=[template_frame, video_state, interactive_state, run_status, run_status2]
+        fn=lambda *args: sam_refine(*args) if args[4] == "Click" else draw_refine(args[1], args[0]),
+        inputs=[template_frame, video_state, point_prompt, click_state, mask_mode_selector, interactive_state],
+        outputs=[template_frame, video_state, run_status, run_status2]
     )
 
+
     # add different mask
     Add_mask_button.click(
         fn=add_multi_mask,