feat: add multi-masking support

app.py

@@ -1,18 +1,17 @@
+import pathlib
+import zipfile
+from typing import Any, Dict, List
+
+import cv2
 import gradio as gr
 import numpy as np
-import cv2
 import torch
-
-
-from typing import Dict, Any, List
-
-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
 
+from src.plot_utils import render_masks
+
 choice_mapping: Dict[str, List[str]] = {
     "tiny": ["sam2_hiera_t.yaml", "assets/checkpoints/sam2_hiera_tiny.pt"],
     "small": ["sam2_hiera_s.yaml", "assets/checkpoints/sam2_hiera_small.pt"],
@@ -27,27 +26,34 @@ def predict(model_choice, annotations: Dict[str, Any]):
     sam2_model = build_sam2(config_file, ckpt_path, device=device)
     predictor = SAM2ImagePredictor(sam2_model)
     predictor.set_image(annotations["image"])
-    coordinates = np.array(
-        [
-            int(annotations["boxes"][0]["xmin"]),
-            int(annotations["boxes"][0]["ymin"]),
-            int(annotations["boxes"][0]["xmax"]),
-            int(annotations["boxes"][0]["ymax"]),
+    coordinates = []
+    for i in range(len(annotations["boxes"])):
+        coordinate = [
+            int(annotations["boxes"][i]["xmin"]),
+            int(annotations["boxes"][i]["ymin"]),
+            int(annotations["boxes"][i]["xmax"]),
+            int(annotations["boxes"][i]["ymax"]),
         ]
-    )
+        coordinates.append(coordinate)
+
     masks, scores, _ = predictor.predict(
         point_coords=None,
         point_labels=None,
-        box=coordinates[None, :],
+        box=np.array(coordinates),
         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)
+    for count, mask in enumerate(masks):
+        mask = mask.transpose(1, 2, 0)  # type:ignore
+        mask_image = (mask * 255).astype(np.uint8)  # Convert to uint8 format
+        cv2.imwrite(f"assets/mask_{count}.png", mask_image)
+        mask_dir = pathlib.Path("assets/")
+        with zipfile.ZipFile("assets/masks.zip", "w") as archive:
+            for mask_file in mask_dir.glob("mask_*.png"):
+                archive.write(mask_file, arcname=mask_file.relative_to(mask_dir))
 
     return [
-        show_masks(annotations["image"], masks, scores, box_coords=coordinates),
-        gr.DownloadButton("Download Mask", value="mask.png", visible=True),
+        render_masks(annotations["image"], masks),
+        gr.DownloadButton("Download Mask", value="assets/masks.zip", visible=True),
     ]
 
 
@@ -77,7 +83,9 @@ with gr.Blocks(delete_cache=(30, 30)) as demo:
         label="Draw a bounding box",
     )
     btn = gr.Button("Get Segmentation Mask")
-    download_btn = gr.DownloadButton("Download Mask", value="mask.png", visible=False)
+    download_btn = gr.DownloadButton(
+        "Download Mask", value="assets/masks.zip", visible=False
+    )
     btn.click(fn=predict, inputs=[model, annotator], outputs=[gr.Plot(), download_btn])
 
 demo.launch()

src/plot_utils.py

@@ -1,90 +1,50 @@
-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,
-    )
-
+from typing import Optional
 
-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)
-    )
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.pyplot import Figure
 
 
-def show_masks(
+def render_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
+    random_color: Optional[bool] = True,
+    smoothen_contours: Optional[bool] = True,
+) -> "Figure":
+    h, w = image.shape[:2]
+    fig, ax = plt.subplots(figsize=(w / 100, h / 100), dpi=100)
+    ax.axis("off")
+    ax.imshow(image)
+
+    for mask in masks:
+        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])
+
+        mask = mask.astype(np.uint8)
+        mask = mask.reshape(h, w)
+        mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+
+        if smoothen_contours:
+            import cv2
+
+            contours, _ = cv2.findContours(
+                mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+            )
+            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, alpha=0.6)
+
+    # Make image occupy the whole figure
+    ax.set_xlim(0, w)
+    ax.set_ylim(h, 0)
+    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+
+    return fig