Add SAM instance segmentation

Author: constantinpape

development/instance_segmentation.py

@@ -2,7 +2,7 @@
 import napari
 
 from elf.io import open_file
-from micro_sam.segment_instances import segment_from_embeddings
+from micro_sam.segment_instances import segment_instances_from_embeddings, segment_instances_sam
 from micro_sam.visualization import compute_pca
 
 
@@ -11,15 +11,23 @@ def mito_segmentation():
     with open_file(input_path) as f:
         raw = f["*.png"][-1, :768, :768]
 
-    predictor = util.get_sam_model()
+    predictor, sam = util.get_sam_model(return_sam=True)
+
+    print("Run SAM prediction ...")
+    seg_sam = segment_instances_sam(sam, raw)
+
     image_embeddings = util.precompute_image_embeddings(predictor, raw, "../examples/embeddings/embeddings-mito2d.zarr")
     embedding_pca = compute_pca(image_embeddings["features"])
 
-    seg, initial_seg = segment_from_embeddings(predictor, image_embeddings=image_embeddings, return_initial_seg=True)
+    print("Run prediction from embeddings ...")
+    seg, initial_seg = segment_instances_from_embeddings(
+        predictor, image_embeddings=image_embeddings, return_initial_seg=True
+    )
 
     v = napari.Viewer()
     v.add_image(raw)
     v.add_image(embedding_pca, scale=(12, 12))
+    v.add_labels(seg_sam)
     v.add_labels(seg)
     v.add_labels(initial_seg)
     napari.run()
@@ -32,21 +40,27 @@ def cell_segmentation():
 
     frame = 11
 
-    predictor = util.get_sam_model()
+    predictor, sam = util.get_sam_model(return_sam=True)
+
+    print("Run prediction from embeddings ...")
     image_embeddings = util.precompute_image_embeddings(
         predictor, timeseries, "../examples/embeddings/embeddings-ctc.zarr"
     )
     embedding_pca = compute_pca(image_embeddings["features"][frame])
 
-    seg, initial_seg = segment_from_embeddings(
+    seg, initial_seg = segment_instances_from_embeddings(
         predictor, image_embeddings=image_embeddings, i=frame, return_initial_seg=True
     )
 
+    print("Run SAM prediction ...")
+    seg_sam = segment_instances_sam(sam, timeseries[frame])
+
     v = napari.Viewer()
     v.add_image(timeseries[frame])
     v.add_image(embedding_pca, scale=(8, 8))
     v.add_labels(seg)
     v.add_labels(initial_seg)
+    v.add_labels(seg_sam)
     napari.run()
 
 

micro_sam/segment_instances.py

@@ -2,7 +2,9 @@
 import vigra
 
 from elf.segmentation import embeddings as embed
+from segment_anything import SamAutomaticMaskGenerator
 from skimage.transform import resize
+
 try:
     from napari.utils import progress as tqdm
 except ImportError:
@@ -17,8 +19,18 @@
 #
 
 
-# TODO implement automatic instance segmentation based on the functionalities from segment anything:
-# https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/automatic_mask_generator.py
+def segment_instances_sam(sam, image, **kwargs):
+    segmentor = SamAutomaticMaskGenerator(sam, **kwargs)
+
+    image_ = util._to_image(image)
+    masks = segmentor.generate(image_)
+    masks = sorted(masks, key=(lambda x: x["area"]), reverse=True)
+
+    segmentation = np.zeros(image.shape[:2], dtype="uint32")
+    for seg_id, mask in enumerate(masks, 1):
+        segmentation[mask["segmentation"]] = seg_id
+
+    return segmentation
 
 
 #
@@ -58,7 +70,7 @@ def _refine_initial_segmentation(predictor, initial_seg, image_embeddings, i, ve
 # - Can we get intermediate, larger embeddings from SAM?
 # - Can we run the encoder in a sliding window and somehow stitch the embeddings?
 # - Or: run the encoder in a sliding window and stitch the initial segmentation result.
-def segment_from_embeddings(
+def segment_instances_from_embeddings(
     predictor, image_embeddings, size_threshold=10, i=None,
     offsets=[[-1, 0], [0, -1], [-3, 0], [0, -3]], distance_type="l2", bias=0.0,
     verbose=True, return_initial_seg=False,