Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -209,16 +209,27 @@ def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr
                      ("[Optional]", "Click image"), (": Try to click the image shown in step2 if you want to generate more masks.\n", None)]
     return painted_image, video_state, interactive_state, operation_log, operation_log
 
-def add_multi_mask(video_state, interactive_state, mask_dropdown):
+def add_multi_mask(video_state, interactive_state, mask_dropdown, mask_input_mode, template_frame):
     try:
-        mask = video_state["masks"][video_state["select_frame_number"]]
+        if mask_input_mode == "Click":
+            # 🟢 Стандартная логика SAM
+            mask = video_state["masks"][video_state["select_frame_number"]]
+        else:
+            # ⚪ Рисованная маска
+            drawn_mask = template_frame  # это уже numpy-массив из gr.Sketchpad
+            if drawn_mask is None:
+                raise ValueError("No mask was drawn.")
+            mask = (np.array(drawn_mask)[:, :, 0] > 127).astype(np.uint8)  # бинаризация
+
         interactive_state["multi_mask"]["masks"].append(mask)
         interactive_state["multi_mask"]["mask_names"].append("mask_{:03d}".format(len(interactive_state["multi_mask"]["masks"])))
         mask_dropdown.append("mask_{:03d}".format(len(interactive_state["multi_mask"]["masks"])))
+
         select_frame, _, _ = show_mask(video_state, interactive_state, mask_dropdown)
         operation_log = [("",""),("Added a mask, use the mask select for target tracking or inpainting.","Normal")]
-    except:
-        operation_log = [("Please click the image in step2 to generate masks.", "Error"), ("","")]
+    except Exception as e:
+        operation_log = [(f"Error: {str(e)}", "Error"), ("","")]
+
     return interactive_state, gr.update(choices=interactive_state["multi_mask"]["mask_names"], value=mask_dropdown), select_frame, [[],[]], operation_log, operation_log
 
 def clear_click(video_state, click_state):
@@ -639,13 +650,20 @@ 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.ImageEditor(type="numpy",label="Draw mask manually or click to add",interactive=True,visible=False,elem_id="template_frame",elem_classes="image",brush=gr.Brush(colors=["#FFFFFF"], default_color="#FFFFFF", default_size=2))
+                template_frame = gr.Sketchpad(label="Draw mask manually or click to add",brush_radius=10,shape=(720, 1280),interactive=True,visible=False,type="numpy",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=[("",""), ("Use click or draw to select mask, then track or inpaint.", "Normal")],
                                                  color_map={"Normal": "green", "Error": "red", "Clear clicks": "gray", "Add mask": "green", "Remove mask": "red"},
                                                  visible=False)
+                mask_input_mode = gr.Radio(
+                    choices=["Click", "Draw"],
+                    value="Click",
+                    label="Select Mask Input Mode",
+                    interactive=True,
+                    visible=True
+                )
                 mask_mode_selector = gr.Radio(
                     choices=["Click", "Draw"],
                     value="Click",