.

sam2_mask.py

@@ -11,8 +11,11 @@ from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 from gradio_image_prompter import ImagePrompter
 
+# def preprocess_image(image):
+#     return image, gr.State([]), gr.State([]), image
+
 def preprocess_image(image):
-    return image, gr.State([]), gr.State([]), image
+    return image, [], [], image
 
 def get_point(point_type, tracking_points, trackings_input_label, first_frame_path, evt: gr.SelectData):
     print(f"You selected {evt.value} at {evt.index} from {evt.target}")
@@ -75,6 +78,28 @@ def show_points(coords, labels, ax, marker_size=375):
     ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
     ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
 
+
+def show_mask(mask, ax, random_color=False, borders=True):
+    if random_color:
+        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    else:
+        color = np.array([30/255, 144/255, 255/255, 0.6])
+    h, w = mask.shape[-2:]
+    mask = mask.astype(np.uint8)
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    if borders:
+        import cv2
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+        # Try to smooth contours
+        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
+        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
+    ax.imshow(mask_image)
+def show_points(coords, labels, ax, marker_size=375):
+    pos_points = coords[labels == 1]
+    neg_points = coords[labels == 0]
+    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+
 def show_box(box, ax):
     x0, y0 = box[0], box[1]
     w, h = box[2] - box[0], box[3] - box[1]
@@ -153,11 +178,10 @@ def create_sam2_tab():
     
     with gr.Column():
         gr.Markdown("# SAM2 Image Predictor")
-
-        image_input = gr.State()
-        input_image = ImagePrompter(show_label=False)
+        gr.Markdown("1. Upload your image\n2. Click points to mask\n3. Submit")
+        
         points_map = gr.Image(label="Points Map", type="pil", interactive=True)
-        # image_input = gr.Image(type="pil", visible=False)  # Original image
+        input_image = gr.Image(type="pil", visible=False)  # Original image
         
         with gr.Row():
             point_type = gr.Radio(["include", "exclude"], value="include", label="Point Type")
@@ -192,4 +216,3 @@ def create_sam2_tab():
         )
     
     return input_image, points_map, output_image
-