support for multi-iamge

app.py

@@ -63,31 +63,38 @@ def segment_reference(image, click):
 
     return masks
 
-def segment_target(target_image, ref_image, ref_mask):
-    target_image = np.array(target_image)
+def segment_target(target_images, ref_image, ref_mask):
+    target_images = [np.array(target_image) for target_image in target_images]
     ref_image = np.array(ref_image)
-    state = sam_utils.load_masks(sam2_vid, [target_image], ref_image, ref_mask)
-    out = sam_utils.propagate_masks(sam2_vid, state)[-1]['segmentation']
-    return out  # Just for placeholder demo
+    state = sam_utils.load_masks(sam2_vid, target_images, ref_image, ref_mask)
+    out = sam_utils.propagate_masks(sam2_vid, state)[1:]
+    return [mask['segmentation'] for mask in out]
 
 def on_reference_upload(img):
     global click_coords
     click_coords = []  # clear the clicks
     return "Click Info: Cleared (new image uploaded)"
 
-def visualize_segmentation(image, masks, target_image, target_mask):
+def visualize_segmentation(image, masks, target_images, target_masks):
     # Visualize the segmentation result
-    fig, ax = plt.subplots(1, 2, figsize=(12, 6))
-    ax[0].imshow(image.convert("L"), cmap='gray')
+    num_tgt = len(target_images)
+    fig, ax = plt.subplots(2, num_tgt, figsize=(6*num_tgt, 12))
+    if num_tgt == 1:
+        ax = np.expand_dims(ax, axis=1)
+    ax[0][0].imshow(image.convert("L"), cmap='gray')
     for i, mask in enumerate(masks):
-        sam_utils.show_mask(mask, ax[0], obj_id=i, alpha=0.75)
-    ax[0].axis('off')
-    ax[0].set_title("Reference Image with Expert Segmentation")
-    ax[1].imshow(target_image.convert("L"), cmap='gray')
-    for i, mask in enumerate(target_mask):
-        sam_utils.show_mask(mask, ax[1], obj_id=i, alpha=0.75)
-    ax[1].axis('off')
-    ax[1].set_title("Target Image with Inferred Segmentation")
+        sam_utils.show_mask(mask, ax[0][0], obj_id=i, alpha=0.75)
+    ax[0][0].axis('off')
+    ax[0][0].set_title("Reference Image with Expert Segmentation")
+    for i in range(1, num_tgt):
+        # set the rest to empty
+        ax[0][i].axis('off')
+    for i in range(num_tgt):
+        ax[1][i].imshow(target_images[i].convert("L"), cmap='gray')
+        for j, mask in enumerate(target_masks[i]):
+            sam_utils.show_mask(mask, ax[1][i], obj_id=j, alpha=0.75)
+        ax[1][i].axis('off')
+        ax[1][i].set_title("Target Image with Inferred Segmentation")
     # save it to buffer
     plt.tight_layout()
     buf = BytesIO()
@@ -106,12 +113,18 @@ def record_click(img, evt: gr.SelectData):
     click_coords.append([evt.index[0], evt.index[1]])
     return f"Clicked at: {click_coords}"
 
-def generate(reference_image, target_image):
+def generate(reference_image, target_images):
+    global click_coords
     if not click_coords:
         return None, "Click on the reference image first!"
+    
+    target_images = [Image.open(f.name).convert("RGB").resize((1024,1024)) for f in target_images]
+
     ref_mask = segment_reference(reference_image, click_coords)
-    tgt_mask = segment_target(target_image, reference_image, ref_mask)
-    vis = visualize_segmentation(reference_image, ref_mask, target_image, tgt_mask)
+    tgt_masks = segment_target(target_images, reference_image, ref_mask)
+    vis = visualize_segmentation(reference_image, ref_mask, target_images, tgt_masks)
+    # clear the clicks
+    click_coords = []
     return vis, "Done!"
 
 with gr.Blocks() as demo:
@@ -119,7 +132,7 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         reference_img = gr.Image(type="pil", label="Reference Image")
-        target_img = gr.Image(type="pil", label="Target Image")
+        target_img = gr.File(file_types=["image"], file_count="multiple", label="Target Images")
     
     click_info = gr.Textbox(label="Click Info")
     generate_btn = gr.Button("Generate")