Use normal instance segmentation functionality in 2d annotator and precompute amg state

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -54,6 +54,7 @@ def mask_data_to_segmentation(
     masks: List[Dict[str, Any]],
     shape: tuple[int, ...],
     with_background: bool,
+    min_object_size: int = 0,
 ) -> np.ndarray:
     """Convert the output of the automatic mask generation to an instance segmentation.
 
@@ -63,15 +64,20 @@ def mask_data_to_segmentation(
         shape: The image shape.
         with_background: Whether the segmentation has background. If yes this function assures that the largest
             object in the output will be mapped to zero (the background value).
+        min_object_size: The minimal size of an object in pixels.
     Returns:
         The instance segmentation.
     """
 
     masks = sorted(masks, key=(lambda x: x["area"]), reverse=True)
     segmentation = np.zeros(shape[:2], dtype="uint32")
 
-    for seg_id, mask in enumerate(masks, 1):
+    seg_id = 1
+    for mask in masks:
+        if mask["area"] < min_object_size:
+            continue
         segmentation[mask["segmentation"]] = seg_id
+        seg_id += 1
 
     if with_background:
         seg_ids, sizes = np.unique(segmentation, return_counts=True)

micro_sam/sam_annotator/annotator_2d.py

@@ -1,3 +1,5 @@
+import os
+import pickle
 import warnings
 from typing import Optional, Tuple
 
@@ -38,15 +40,13 @@ def _segment_widget(v: Viewer) -> None:
     v.layers["current_object"].refresh()
 
 
-def _get_amg(is_tiled, with_background, min_initial_size, box_extension=0.05):
+def _get_amg(is_tiled, with_background=True, box_extension=0.05):
     if is_tiled:
-        amg = instance_segmentation.TiledEmbeddingMaskGenerator(
-            PREDICTOR, with_background=with_background, min_initial_size=min_initial_size, box_extension=box_extension,
+        amg = instance_segmentation.TiledAutomaticMaskGenerator(
+            PREDICTOR, with_background=with_background, box_extension=box_extension,
         )
     else:
-        amg = instance_segmentation.EmbeddingMaskGenerator(
-            PREDICTOR, min_initial_size=min_initial_size, box_extension=box_extension,
-        )
+        amg = instance_segmentation.EmbeddingMaskGenerator(PREDICTOR, box_extension=box_extension)
     return amg
 
 
@@ -62,29 +62,27 @@ def _changed_param(amg, **params):
 @magicgui(call_button="Automatic Segmentation")
 def _autosegment_widget(
     v: Viewer,
-    with_background: bool = True,
     pred_iou_thresh: float = 0.88,
     stability_score_thresh: float = 0.95,
-    min_initial_size: int = 10,
-    box_extension: float = 0.05,
+    min_object_size: int = 25,
 ) -> None:
     global AMG
     is_tiled = IMAGE_EMBEDDINGS["input_size"] is None
-    param_changed = _changed_param(
-        AMG, with_background=with_background, min_initial_size=min_initial_size, box_extension=box_extension
-    )
-    if AMG is None or param_changed:
-        if param_changed:
-            print(f"The parameter {param_changed} was changed, so the full instance segmentation has to be recomputed.")
-        AMG = _get_amg(is_tiled, with_background, min_initial_size, box_extension)
+    if AMG is None:
+        AMG = _get_amg(is_tiled)
 
     if not AMG.is_initialized:
         AMG.initialize(v.layers["raw"].data, image_embeddings=IMAGE_EMBEDDINGS, verbose=True)
 
-    seg = AMG.generate(pred_iou_thresh=pred_iou_thresh, stability_score_thresh=stability_score_thresh)
-    if not is_tiled:
-        shape = v.layers["raw"].data.shape[:2]
-        seg = instance_segmentation.mask_data_to_segmentation(seg, shape, with_background)
+    seg = AMG.generate(
+        pred_iou_thresh=pred_iou_thresh, stability_score_thresh=stability_score_thresh,
+        min_mask_region_area=min_object_size
+    )
+
+    shape = v.layers["raw"].data.shape[:2]
+    seg = instance_segmentation.mask_data_to_segmentation(
+        seg, shape, with_background=True, min_object_size=min_object_size
+    )
     assert isinstance(seg, np.ndarray)
 
     v.layers["auto_segmentation"].data = seg
@@ -192,6 +190,25 @@ def _update_viewer(v, raw, show_embeddings, segmentation_result):
     v.layers["current_object"].data = np.zeros(shape, dtype="uint32")
 
 
+def _precompute_amg_state(raw, save_path):
+    global AMG
+
+    is_tiled = IMAGE_EMBEDDINGS["input_size"] is None
+    AMG = _get_amg(is_tiled)
+
+    save_path_amg = os.path.join(save_path, "amg_state.pickle")
+    if os.path.exists(save_path_amg):
+        with open(save_path_amg, "rb") as f:
+            amg_state = pickle.load(f)
+        AMG.set_state(amg_state)
+        return
+
+    print("Precomputing the state for instance segmentation")
+    AMG.initialize(raw, image_embeddings=IMAGE_EMBEDDINGS, verbose=True)
+    with open(save_path_amg, "wb") as f:
+        pickle.dump(AMG.get_state(), f)
+
+
 def annotator_2d(
     raw: np.ndarray,
     embedding_path: Optional[str] = None,
@@ -241,6 +258,8 @@ def annotator_2d(
         PREDICTOR, raw, save_path=embedding_path, ndim=2, tile_shape=tile_shape, halo=halo,
         wrong_file_callback=show_wrong_file_warning
     )
+    if embedding_path is not None:
+        _precompute_amg_state(raw, embedding_path)
 
     # we set the pre-computed image embeddings if we don't use tiling
     # (if we use tiling we cannot directly set it because the tile will be chosen dynamically)