unsamv2

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -8,7 +8,7 @@ import os
 import sys
 import threading
 from pathlib import Path
-from typing import List, Optional, Sequence, Tuple
+from typing import List, Optional, Sequence
 
 import cv2
 import gradio as gr
@@ -28,9 +28,11 @@ if SAM2_REPO.exists():
 from sam2.build_sam import build_sam2  # noqa: E402
 from sam2.sam2_image_predictor import SAM2ImagePredictor  # noqa: E402
 
-logging.basicConfig(level=os.getenv("UNSAMV2_LOGLEVEL", "INFO"))
+logging.basicConfig(level=logging.INFO)
 LOGGER = logging.getLogger("unsamv2-gradio")
 
+USE_M2M_REFINEMENT = True
+
 CONFIG_PATH = os.getenv("UNSAMV2_CONFIG", "configs/unsamv2_small.yaml")
 CKPT_PATH = Path(
     os.getenv("UNSAMV2_CKPT", SAM2_REPO / "checkpoints" / "unsamv2_plus_ckpt.pt")
@@ -53,6 +55,22 @@ POINT_COLORS_BGR = {
 MASK_COLOR_BGR = (0, 196, 255)
 OUTLINE_COLOR_BGR = (0, 165, 255)
 
+DEFAULT_IMAGE_PATH = REPO_ROOT / "demo" / "bird.webp"
+
+
+def _load_default_image() -> Optional[np.ndarray]:
+    if not DEFAULT_IMAGE_PATH.exists():
+        LOGGER.warning("Default image missing at %s", DEFAULT_IMAGE_PATH)
+        return None
+    img_bgr = cv2.imread(str(DEFAULT_IMAGE_PATH), cv2.IMREAD_COLOR)
+    if img_bgr is None:
+        LOGGER.warning("Could not read default image at %s", DEFAULT_IMAGE_PATH)
+        return None
+    return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+
+
+DEFAULT_IMAGE = _load_default_image()
+
 
 class ModelManager:
     """Keeps SAM2 models on each device and spawns lightweight predictors."""
@@ -82,7 +100,7 @@ class ModelManager:
         return self._models[key]
 
     def make_predictor(self, device: torch.device) -> SAM2ImagePredictor:
-        return SAM2ImagePredictor(self.get_model(device))
+        return SAM2ImagePredictor(self.get_model(device), mask_threshold=-1.0)
 
 
 MODEL_MANAGER = ModelManager()
@@ -119,6 +137,41 @@ def build_granularity_tensor(value: float, device: torch.device) -> torch.Tensor
     return tensor
 
 
+def apply_m2m_refinement(
+    predictor,
+    point_coords,
+    point_labels,
+    granularity,
+    logits,
+    best_mask_idx,
+    use_m2m: bool = True,
+):
+    """Optionally run a second M2M pass using the best mask's logits."""
+    if not use_m2m:
+        return None
+
+    logging.info("Applying M2M refinement...")
+    try:
+        if logits is None:
+            raise ValueError("logits must be provided for M2M refinement.")
+
+        low_res_logits = logits[best_mask_idx : best_mask_idx + 1]
+        refined_masks, refined_scores, _ = predictor.predict(
+            point_coords=point_coords,
+            point_labels=point_labels,
+            multimask_output=False,
+            gra=granularity,
+            mask_input=low_res_logits,
+        )
+        refined_mask = refined_masks[0]
+        refined_score = float(refined_scores[0])
+        logging.info("M2M refinement completed with score: %.3f", refined_score)
+        return refined_mask, refined_score
+    except Exception as exc:  # pragma: no cover - logging only
+        logging.error("M2M refinement failed: %s, using original mask", exc)
+        return None
+
+
 def draw_overlay(
     image: np.ndarray,
     mask: Optional[np.ndarray],
@@ -150,37 +203,21 @@ def draw_overlay(
     return cv2.cvtColor(canvas_bgr, cv2.COLOR_BGR2RGB)
 
 
-def points_table(points: Sequence[Sequence[float]], labels: Sequence[int]) -> List[List[str]]:
-    table = []
-    for idx, ((x, y), lbl) in enumerate(zip(points, labels), start=1):
-        table.append([
-            idx,
-            round(float(x), 1),
-            round(float(y), 1),
-            "fg" if lbl == 1 else "bg",
-        ])
-    return table
-
-
 def handle_image_upload(image: Optional[np.ndarray]):
     img = ensure_uint8(image)
     if img is None:
         return (
-            None,
             None,
             None,
             [],
             [],
-            [],
             "Upload an image to start adding clicks.",
         )
     return (
         img,
-        None,
         img,
         [],
         [],
-        [],
         "Image loaded. Choose click type, then tap on the image.",
     )
 
@@ -195,20 +232,16 @@ def handle_click(
     if image is None:
         return (
             gr.update(),
-            None,
             pts,
             lbls,
-            points_table(pts, lbls),
             "Upload an image first.",
         )
     coord = evt.index  # (x, y)
     if coord is None:
         return (
             gr.update(),
-            None,
             pts,
             lbls,
-            points_table(pts, lbls),
             "Couldn't read click position.",
         )
     x, y = coord
@@ -217,29 +250,27 @@ def handle_click(
     lbls = lbls + [label]
     overlay = draw_overlay(image, None, pts, lbls)
     status = f"Added {'positive' if label == 1 else 'negative'} click at ({int(x)}, {int(y)})."
-    return overlay, None, pts, lbls, points_table(pts, lbls), status
+    return overlay, pts, lbls, status
 
 
 def undo_last_click(image: Optional[np.ndarray], pts: List[Sequence[float]], lbls: List[int]):
     if not pts:
         return (
             gr.update(),
-            None,
             pts,
             lbls,
-            points_table(pts, lbls),
             "No clicks to undo.",
         )
     pts = pts[:-1]
     lbls = lbls[:-1]
     overlay = draw_overlay(image, None, pts, lbls) if image is not None else None
     status = "Removed the last click."
-    return overlay, None, pts, lbls, points_table(pts, lbls), status
+    return overlay, pts, lbls, status
 
 
 def clear_clicks(image: Optional[np.ndarray]):
     overlay = image if image is not None else None
-    return overlay, None, [], [], [], "Cleared all clicks."
+    return overlay, [], [], "Cleared all clicks."
 
 
 def _run_segmentation(
@@ -250,9 +281,9 @@ def _run_segmentation(
 ):
     img = ensure_uint8(image)
     if img is None:
-        return None, None, "Upload an image to segment."
+        return None, "Upload an image to segment."
     if not pts:
-        return draw_overlay(img, None, [], []), None, "Add at least one click before running segmentation."
+        return draw_overlay(img, None, [], []), "Add at least one click before running segmentation."
 
     device = choose_device()
     predictor = MODEL_MANAGER.make_predictor(device)
@@ -262,7 +293,7 @@ def _run_segmentation(
     labels = np.asarray(lbls, dtype=np.int32)
     gran_tensor = build_granularity_tensor(granularity, predictor.device)
 
-    masks, scores, _ = predictor.predict(
+    masks, scores, logits = predictor.predict(
         point_coords=coords,
         point_labels=labels,
         multimask_output=True,
@@ -271,10 +302,27 @@ def _run_segmentation(
     )
     best_idx = int(np.argmax(scores))
     best_mask = masks[best_idx].astype(bool)
+    status = (
+        f"Best mask #{best_idx + 1} IoU score: {float(scores[best_idx]):.3f} | "
+        f"granularity={granularity:.2f}"
+    )
+
+    refinement = apply_m2m_refinement(
+        predictor=predictor,
+        point_coords=coords,
+        point_labels=labels,
+        granularity=float(granularity),
+        logits=logits,
+        best_mask_idx=best_idx,
+        use_m2m=USE_M2M_REFINEMENT,
+    )
+    if refinement is not None:
+        refined_mask, refined_score = refinement
+        best_mask = refined_mask.astype(bool)
+        status += f" | M2M IoU: {refined_score:.3f}"
+
     overlay = draw_overlay(img, best_mask, pts, lbls)
-    mask_vis = (best_mask.astype(np.uint8) * 255)
-    status = f"Best mask #{best_idx + 1} IoU score: {float(scores[best_idx]):.3f} | granularity={granularity:.2f}"
-    return overlay, mask_vis, status
+    return overlay, status
 
 
 if spaces is not None and ZERO_GPU_ENABLED:
@@ -287,30 +335,20 @@ def build_demo() -> gr.Blocks:
     with gr.Blocks(title="UnSAMv2 Interactive Segmentation", theme=gr.themes.Soft()) as demo:
         gr.Markdown(
             """## UnSAMv2 · Interactive Granularity Control
-Upload an image, add positive/negative clicks, tune granularity, and run segmentation.
-ZeroGPU automatically pulls a GPU when available; otherwise the app falls back to CPU."""
+Upload an image, add positive/negative clicks, tune granularity, and run segmentation."""
         )
 
-        image_state = gr.State()
+        image_state = gr.State(DEFAULT_IMAGE)
         points_state = gr.State([])
         labels_state = gr.State([])
 
-        with gr.Row():
-            image_input = gr.Image(
-                label="1 · Upload image & click to add prompts",
-                type="numpy",
-                height=480,
-            )
-            overlay_output = gr.Image(
-                label="Segmentation preview",
-                interactive=False,
-                height=480,
-            )
-            mask_output = gr.Image(
-                label="Binary mask",
-                interactive=False,
-                height=480,
-            )
+        image_input = gr.Image(
+            label="Image · clicks & mask",
+            type="numpy",
+            height=480,
+            value=DEFAULT_IMAGE,
+            sources=["upload"],
+        )
 
         with gr.Row():
             point_mode = gr.Radio(
@@ -322,34 +360,26 @@ ZeroGPU automatically pulls a GPU when available; otherwise the app falls back t
                 minimum=GRANULARITY_MIN,
                 maximum=GRANULARITY_MAX,
                 value=0.2,
-                step=0.05,
+                step=0.01,
                 label="Granularity",
                 info="Lower = finer details, Higher = coarser regions",
             )
-            segment_button = gr.Button("3 · Segment", variant="primary")
+            segment_button = gr.Button("Segment", variant="primary")
 
         with gr.Row():
             undo_button = gr.Button("Undo last click")
             clear_button = gr.Button("Clear clicks")
 
-        points_table_output = gr.Dataframe(
-            headers=["#", "x", "y", "type"],
-            datatype=["number", "number", "number", "str"],
-            interactive=False,
-            label="2 · Click history",
-        )
         status_markdown = gr.Markdown(" Ready.")
 
         image_input.upload(
             handle_image_upload,
             inputs=[image_input],
             outputs=[
-                overlay_output,
-                mask_output,
+                image_input,
                 image_state,
                 points_state,
                 labels_state,
-                points_table_output,
                 status_markdown,
             ],
         )
@@ -358,12 +388,10 @@ ZeroGPU automatically pulls a GPU when available; otherwise the app falls back t
             handle_image_upload,
             inputs=[image_input],
             outputs=[
-                overlay_output,
-                mask_output,
+                image_input,
                 image_state,
                 points_state,
                 labels_state,
-                points_table_output,
                 status_markdown,
             ],
         )
@@ -377,11 +405,9 @@ ZeroGPU automatically pulls a GPU when available; otherwise the app falls back t
                 image_state,
             ],
             outputs=[
-                overlay_output,
-                mask_output,
+                image_input,
                 points_state,
                 labels_state,
-                points_table_output,
                 status_markdown,
             ],
         )
@@ -390,11 +416,9 @@ ZeroGPU automatically pulls a GPU when available; otherwise the app falls back t
             undo_last_click,
             inputs=[image_state, points_state, labels_state],
             outputs=[
-                overlay_output,
-                mask_output,
+                image_input,
                 points_state,
                 labels_state,
-                points_table_output,
                 status_markdown,
             ],
         )
@@ -403,11 +427,9 @@ ZeroGPU automatically pulls a GPU when available; otherwise the app falls back t
             clear_clicks,
             inputs=[image_state],
             outputs=[
-                overlay_output,
-                mask_output,
+                image_input,
                 points_state,
                 labels_state,
-                points_table_output,
                 status_markdown,
             ],
         )
@@ -415,7 +437,7 @@ ZeroGPU automatically pulls a GPU when available; otherwise the app falls back t
         segment_button.click(
             segment_fn,
             inputs=[image_state, points_state, labels_state, granularity_slider],
-            outputs=[overlay_output, mask_output, status_markdown],
+            outputs=[image_input, status_markdown],
         )
 
         demo.queue(max_size=8)