Merge pull request #58 from computational-cell-analytics/instance-seg-plus

Implement tool for interactive instacne segmentation

Author: constantinpape

development/annotator_2d_tiled.py

@@ -16,7 +16,7 @@ def annotator_with_tiling():
     # napari.run()
 
     embedding_path = "./embeddings/embeddings-tiled.zarr"
-    annotator_2d(im, embedding_path, tile_shape=(1024, 1024), halo=(256, 256))
+    annotator_2d(im, embedding_path, tile_shape=(1024, 1024), halo=(256, 256), show_embeddings=False)
 
 
 def debug():

development/annotator_3d_tiled.py

@@ -6,7 +6,7 @@ def annotator_with_tiling():
     with z5py.File("/home/pape/Work/data/cremi/sampleA.n5", "r") as f:
         raw = f["volumes/raw/s0"][:25]
     embedding_path = "./embeddings/embeddings-tiled_3d.zarr"
-    annotator_3d(raw, embedding_path, tile_shape=(512, 512), halo=(64, 64))
+    annotator_3d(raw, embedding_path, tile_shape=(512, 512), halo=(64, 64), show_embeddings=False)
 
 
 def segment_tiled():
@@ -25,8 +25,8 @@ def segment_tiled():
 
 
 def main():
-    # annotator_with_tiling()
-    segment_tiled()
+    annotator_with_tiling()
+    # segment_tiled()
 
 
 main()

examples/sam_annotator_tracking.py

@@ -10,7 +10,7 @@ def track_ctc_data():
     example_data_directory = "./data"
     with open_file(str(fetch_example_data(example_data_directory)), mode="r") as f:
         timeseries = f["*.tif"]
-    annotator_tracking(timeseries, embedding_path="./embeddings/embeddings-ctc.zarr")
+    annotator_tracking(timeseries, embedding_path="./embeddings/embeddings-ctc.zarr", show_embeddings=False)
 
 
 def fetch_example_data(save_directory):

micro_sam/sam_annotator/__init__.py

@@ -2,3 +2,4 @@
 from .annotator_3d import annotator_3d
 from .annotator_tracking import annotator_tracking
 from .image_series_annotator import image_folder_annotator, image_series_annotator
+from .interactive_instance_segmentation import interactive_instance_segmentation

micro_sam/sam_annotator/annotator_2d.py

@@ -39,23 +39,34 @@ def segment_wigdet(v: Viewer):
     v.layers["current_object"].refresh()
 
 
-# TODO expose more parameters:
-# - min initial size
-# - advanced params???
 @magicgui(call_button="Automatic Segmentation")
 def autosegment_widget(
-    v: Viewer, with_background: bool = True, box_extension: float = 0.1, pred_iou_thresh: float = 0.88
+    v: Viewer,
+    with_background: bool = True,
+    pred_iou_thresh: float = 0.88,
+    stability_score_thresh: float = 0.95,
+    min_initial_size: int = 10,
+    use_box: bool = True,
+    use_mask: bool = True,
+    use_points: bool = False,
+    box_extension: float = 0.1,
 ):
     is_tiled = IMAGE_EMBEDDINGS["input_size"] is None
     if is_tiled:
         seg = segment_instances.segment_instances_from_embeddings_with_tiling(
             PREDICTOR, IMAGE_EMBEDDINGS, with_background=with_background,
             box_extension=box_extension, pred_iou_thresh=pred_iou_thresh,
+            stability_score_thresh=stability_score_thresh,
+            min_initial_size=min_initial_size,
+            use_box=use_box, use_points=use_points, use_mask=use_mask,
         )
     else:
         seg = segment_instances.segment_instances_from_embeddings(
             PREDICTOR, IMAGE_EMBEDDINGS, with_background=with_background,
             box_extension=box_extension, pred_iou_thresh=pred_iou_thresh,
+            stability_score_thresh=stability_score_thresh,
+            min_initial_size=min_initial_size,
+            use_box=use_box, use_points=use_points, use_mask=use_mask,
         )
     v.layers["auto_segmentation"].data = seg
     v.layers["auto_segmentation"].refresh()
@@ -90,11 +101,9 @@ def _initialize_viewer(raw, segmentation_result, tile_shape, show_embeddings):
     v.add_labels(data=np.zeros(shape, dtype="uint32"), name="current_object")
 
     # show the PCA of the image embeddings
-    if show_embeddings and tile_shape is None:
-        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS["features"], shape)
+    if show_embeddings:
+        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS)
         v.add_image(embedding_vis, name="embeddings", scale=scale)
-    elif show_embeddings:
-        warnings.warn("Embeddings cannot be shown for tiled prediction.")
 
     labels = ["positive", "negative"]
     prompts = v.add_points(
@@ -205,8 +214,6 @@ def annotator_2d(
 
 
 def main():
-    import warnings
-
     parser = _initialize_parser(description="Run interactive segmentation for an image.")
     args = parser.parse_args()
     raw = util.load_image_data(args.input, ndim=2, key=args.key)

micro_sam/sam_annotator/annotator_3d.py

@@ -215,7 +215,7 @@ def annotator_3d(
 
     # show the PCA of the image embeddings
     if show_embeddings:
-        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS["features"], raw.shape)
+        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS)
         v.add_image(embedding_vis, name="embeddings", scale=scale)
 
     labels = ["positive", "negative"]

micro_sam/sam_annotator/annotator_tracking.py

@@ -381,7 +381,7 @@ def annotator_tracking(
 
     # show the PCA of the image embeddings
     if show_embeddings:
-        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS["features"], raw.shape)
+        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS)
         v.add_image(embedding_vis, name="embeddings", scale=scale)
 
     #

micro_sam/sam_annotator/interactive_instance_segmentation.py

@@ -0,0 +1,91 @@
+import napari
+import numpy as np
+from napari import Viewer
+from magicgui import magicgui
+
+from .annotator_2d import _get_shape
+from .util import _initialize_parser
+from ..import util
+from ..import segment_instances
+from ..visualization import project_embeddings_for_visualization
+
+
+@magicgui(call_button="Automatic Segmentation [S]")
+def autosegment_widget(
+    v: Viewer,
+    pred_iou_thresh: float = 0.88,
+    stability_score_thresh: float = 0.95,
+    min_initial_size: int = 10,
+    box_extension: float = 0.1,
+    with_background: bool = True,
+    use_box: bool = True,
+    use_mask: bool = True,
+    use_points: bool = False,
+):
+    is_tiled = IMAGE_EMBEDDINGS["input_size"] is None
+    if is_tiled:
+        seg, initial_seg = segment_instances.segment_instances_from_embeddings_with_tiling(
+            PREDICTOR, IMAGE_EMBEDDINGS, with_background=with_background,
+            box_extension=box_extension, pred_iou_thresh=pred_iou_thresh,
+            stability_score_thresh=stability_score_thresh,
+            min_initial_size=min_initial_size, return_initial_segmentation=True, verbose=2,
+            use_box=use_box, use_mask=use_mask, use_points=use_points,
+        )
+    else:
+        seg, initial_seg = segment_instances.segment_instances_from_embeddings(
+            PREDICTOR, IMAGE_EMBEDDINGS, with_background=with_background,
+            box_extension=box_extension, pred_iou_thresh=pred_iou_thresh,
+            stability_score_thresh=stability_score_thresh,
+            min_initial_size=min_initial_size, return_initial_segmentation=True, verbose=2,
+            use_box=use_box, use_mask=use_mask, use_points=use_points,
+        )
+
+    v.layers["auto_segmentation"].data = seg
+    v.layers["auto_segmentation"].refresh()
+
+    v.layers["initial_segmentation"].data = initial_seg
+    v.layers["initial_segmentation"].refresh()
+
+
+def interactive_instance_segmentation(
+    raw, embedding_path=None, model_type="vit_h", tile_shape=None, halo=None, checkpoint=None,
+):
+    """Visualizing and debugging automatic instance segmentation.
+    """
+    global PREDICTOR, IMAGE_EMBEDDINGS
+
+    PREDICTOR = util.get_sam_model(model_type=model_type, checkpoint_path=checkpoint)
+    IMAGE_EMBEDDINGS = util.precompute_image_embeddings(
+        PREDICTOR, raw, save_path=embedding_path, ndim=2, tile_shape=tile_shape, halo=halo,
+    )
+
+    shape = _get_shape(raw)
+
+    v = napari.Viewer()
+
+    v.add_image(raw)
+    embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS)
+    v.add_image(embedding_vis, name="embeddings", scale=scale, visible=False)
+    v.add_labels(data=np.zeros(shape, dtype="uint32"), name="initial_segmentation")
+    v.add_labels(data=np.zeros(shape, dtype="uint32"), name="auto_segmentation")
+
+    v.window.add_dock_widget(autosegment_widget)
+
+    napari.run()
+
+
+def main():
+    parser = _initialize_parser(
+        description="Run the automatic instance segmentation in interactive mode"
+        "to determine the optimal segmentation parameters.",
+        with_segmentation_result=False, with_show_embeddings=False,
+    )
+    parser.add_argument(
+        "-c", "--checkpoint", default=None,
+        help="Path to alternative checkpoint instead of the standard model."
+    )
+    args = parser.parse_args()
+    raw = util.load_image_data(args.input, ndim=2, key=args.key)
+    interactive_instance_segmentation(
+        raw, args.embedding_path, args.model_type, args.tile_shape, args.halo, args.checkpoint
+    )

micro_sam/sam_annotator/util.py

@@ -323,7 +323,7 @@ def toggle_label(prompts):
     prompts.refresh_colors()
 
 
-def _initialize_parser(description, with_segmentation_result=True):
+def _initialize_parser(description, with_segmentation_result=True, with_show_embeddings=True):
     parser = argparse.ArgumentParser(description=description)
 
     parser.add_argument(
@@ -336,6 +336,7 @@ def _initialize_parser(description, with_segmentation_result=True):
         help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
         "for a image series it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
     )
+
     parser.add_argument(
         "-e", "--embedding_path",
         help="The filepath for saving/loading the pre-computed image embeddings. "
@@ -356,10 +357,11 @@ def _initialize_parser(description, with_segmentation_result=True):
             help="The key for opening the segmentation data. Same rules as for 'key' apply."
         )
 
-    parser.add_argument(
-        "--show_embeddings", action="store_true",
-        help="Visualize the embeddings computed by SegmentAnything. This can be helpful for debugging."
-    )
+    if with_show_embeddings:
+        parser.add_argument(
+            "--show_embeddings", action="store_true",
+            help="Visualize the embeddings computed by SegmentAnything. This can be helpful for debugging."
+        )
     parser.add_argument(
         "--model_type", default="vit_h", help="The segment anything model that will be used, one of vit_h,l,b."
     )