Upload app.py

app.py

@@ -63,15 +63,16 @@ def run_sam(image, predictor, selected_points):
     # input_points = np.array([[210, 300]])
     # input_labels = np.array([1])
     masks, _, _ = predictor.predict(
-        point_coords=input_points,
-        point_labels=input_labels,
+        point_coords=np.array(input_points),
+        point_labels=np.array(input_labels),
         multimask_output=False,  # 单对象输出
     )
-    # print(masks.shape, np.unique(masks))
     best_mask = masks[0]
-    # best_mask = masks[np.argmax(scores)]
-    # print(np.unique(best_mask), best_mask.shape)
-    # print(type(best_mask), best_mask.dtype)
+    # dilate
+    if len(input_points) > 1:
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        best_mask = cv2.dilate(best_mask, kernel, iterations=1)
+        best_mask = cv2.erode(best_mask, kernel, iterations=1)
     visible_mask = 255 * best_mask.astype(np.uint8)
     return visible_mask
 
@@ -396,7 +397,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type="numpy", label='Input Occlusion Image', height=300)
+            input_image = gr.Image(type="numpy", label='Input Occlusion Image', interactive=False, height=300)
             with gr.Row():
                 x_input = gr.Number(label="X Coordinate", value=0)
                 y_input = gr.Number(label="Y Coordinate", value=0)
@@ -443,7 +444,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     input_image.upload(
         reset_image,
         [predictor, input_image],
-        [predictor, visible_mask, occlusion_mask, original_image]
+        [predictor, original_image]
     )
     see_button.click(
         see_point,