feat: initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,92 @@
+import gradio as gr
+import numpy as np
+import cv2
+
+from PIL import Image
+
+from src.plot_utils import show_masks
+from gradio_image_annotation import image_annotator
+
+
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+choice_mapping = {
+    "tiny": ["sam2_hiera_t.yaml", "assets/checkpoints/sam2_hiera_tiny.pt"],
+    "small": ["sam2_hiera_s.yaml", "assets/checkpoints/sam2_hiera_small.pt"],
+    "base_plus": ["sam2_hiera_b+.yaml", "assets/checkpoints/sam2_hiera_base_plus.pt"],
+    "large": ["sam2_hiera_l.yaml", "assets/checkpoints/sam2_hiera_large.pt"],
+}
+
+
+def predict(model_choice: str, annotations, image):
+    config_file, ckpt_path = choice_mapping[str(model_choice)]
+    sam2_model = build_sam2(config_file, ckpt_path, device="cpu")
+    predictor = SAM2ImagePredictor(sam2_model)
+    predictor.set_image(image)
+    coordinates = np.array(
+        [
+            int(annotations["boxes"][0]["xmin"]),
+            int(annotations["boxes"][0]["ymin"]),
+            int(annotations["boxes"][0]["xmax"]),
+            int(annotations["boxes"][0]["ymax"]),
+        ]
+    )
+    masks, scores, _ = predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=coordinates[None, :],
+        multimask_output=False,
+    )
+    mask = masks.transpose(1, 2, 0)
+    mask_image = (mask * 255).astype(np.uint8)  # Convert to uint8 format
+    cv2.imwrite("mask.png", mask_image)
+
+    return [
+        show_masks(image, masks, scores, box_coords=coordinates),
+        gr.DownloadButton("Download Mask", value="mask.png", visible=True),
+    ]
+
+
+with gr.Blocks(delete_cache=(30, 30)) as demo:
+    gr.Markdown(
+        """
+        # 1. Choose Model Checkpoint
+        """
+    )
+    with gr.Row():
+        model = gr.Dropdown(
+            choices=["tiny", "small", "base_plus", "large"],
+            value="tiny",
+            label="Model Checkpoint",
+            info="Which model checkpoint to load?",
+        )
+
+    gr.Markdown(
+        """
+        # 2. Upload an Image
+        """
+    )
+
+    with gr.Row():
+        img = gr.Image(value="./assets/img.png", type="numpy", label="Input Image")
+
+    gr.Markdown(
+        """
+        # 3. Draw Bounding Box
+        """
+    )
+
+    annotator = image_annotator(
+        value={"image": img.value["path"]},
+        disable_edit_boxes=True,
+        single_box=True,
+        label="Draw a bounding box",
+    )
+    btn = gr.Button("Get Segmentation Mask")
+    download_btn = gr.DownloadButton("Download Mask", value="mask.png", visible=False)
+    btn.click(
+        fn=predict, inputs=[model, annotator, img], outputs=[gr.Plot(), download_btn]
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+gradio_image_annotation
+-e segment-anything-2/

segment-anything-2

@@ -0,0 +1 @@
+Subproject commit 82b026cd5578af78757323ab99a0b5c8dc456cff

src/plot_utils.py

@@ -0,0 +1,90 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+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]
+    ax.add_patch(
+        plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2)
+    )
+
+
+def show_masks(
+    image,
+    masks,
+    scores,
+    point_coords=None,
+    box_coords=None,
+    input_labels=None,
+    borders=True,
+):
+    num_masks = len(masks)
+    num_cols = num_masks  # Number of columns is equal to the number of masks
+
+    fig, axes = plt.subplots(1, num_cols, figsize=(5 * num_cols, 5))
+
+    if num_masks == 1:
+        axes = [axes]  # Ensure axes is iterable when there's only one mask
+
+    for i, (mask, score) in enumerate(zip(masks, scores)):
+        ax = axes[i]
+
+        ax.imshow(image)
+        show_mask(mask, ax, borders=borders)
+        if point_coords is not None:
+            assert input_labels is not None
+            show_points(point_coords, input_labels, ax)
+        if box_coords is not None:
+            show_box(box_coords, ax)
+        if len(scores) > 1:
+            ax.set_title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
+        ax.axis("off")
+
+    plt.tight_layout()
+    return plt