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

Improve automatic 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()
 
 

examples/sam_annotator_2d.py

@@ -2,22 +2,30 @@
 from micro_sam.sam_annotator import annotator_2d
 
 
-# TODO describe how to get the data and don't use hard-coded system path
+# TODO describe how to get the data and don't use hard-coded path
 def livecell_annotator():
     im = imageio.imread(
         "/home/pape/Work/data/incu_cyte/livecell/images/livecell_test_images/A172_Phase_C7_1_01d04h00m_4.tif"
     )
     embedding_path = "./embeddings/embeddings-livecell_cropped.zarr"
-    annotator_2d(im, embedding_path, show_embeddings=True)
+    annotator_2d(im, embedding_path, show_embeddings=False)
+
+
+# This runs interactive 2d annotation for data from the cell tracking challenge:
+# It uses the training data for the HeLA dataset. You can download the data from
+# http://data.celltrackingchallenge.net/training-datasets/DIC-C2DH-HeLa.zip
+def hela_2d_annotator():
+    im = imageio.imread("./data/DIC-C2DH-HeLa/train/01/t011.tif")
+    embedding_path = "./embeddings/embeddings-hela2d.zarr"
+    annotator_2d(im, embedding_path, show_embeddings=False)
 
 
 def main():
     # 2d annotator for livecell data
-    livecell_annotator()
+    # livecell_annotator()
 
-    # TODO
     # 2d annotator for cell tracking challenge hela data
-    # hela_2d_annotator()
+    hela_2d_annotator()
 
 
 if __name__ == "__main__":

examples/sam_annotator_tracking.py

@@ -3,7 +3,7 @@
 
 
 # This runs the interactive tracking annotator for data from the cell tracking challenge:
-# It uses the training data for the HeLA dataset. You can download the data via
+# It uses the training data for the HeLA dataset. You can download the data from
 # http://data.celltrackingchallenge.net/training-datasets/DIC-C2DH-HeLa.zip
 def track_ctc_data():
     path = "./data/DIC-C2DH-HeLa/train/01"

micro_sam/sam_annotator/annotator_2d.py

@@ -5,8 +5,8 @@
 from napari import Viewer
 
 from .. import util
+from .. import segment_instances
 from ..visualization import project_embeddings_for_visualization
-from ..segment_instances import segment_from_embeddings
 from ..segment_from_prompts import segment_from_points
 from .util import (
     commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points, toggle_label, LABEL_COLOR_CYCLE
@@ -21,20 +21,26 @@ def segment_wigdet(v: Viewer):
     v.layers["current_object"].refresh()
 
 
-# TODO enable choosing setting the segmentation method and setting other params
-@magicgui(call_button="Segment All Objects")
-def autosegment_widget(v: Viewer):
-    # choose if we segment with/without tiling based on the image shape
-    seg = segment_from_embeddings(PREDICTOR, IMAGE_EMBEDDINGS)
+# TODO expose more parameters
+@magicgui(call_button="Segment All Objects", method={"choices": ["default", "sam", "embeddings"]})
+def autosegment_widget(v: Viewer, method: str = "default"):
+    if method in ("default", "sam"):
+        print("Run automatic segmentation with SAM. This can take a few minutes ...")
+        image = v.layers["raw"].data
+        seg = segment_instances.segment_instances_sam(SAM, image)
+    elif method == "embeddings":
+        seg = segment_instances.segment_instances_from_embeddings(PREDICTOR, IMAGE_EMBEDDINGS)
+    else:
+        raise ValueError
     v.layers["auto_segmentation"].data = seg
     v.layers["auto_segmentation"].refresh()
 
 
 def annotator_2d(raw, embedding_path=None, show_embeddings=False, segmentation_result=None):
     # for access to the predictor and the image embeddings in the widgets
-    global PREDICTOR, IMAGE_EMBEDDINGS
+    global PREDICTOR, IMAGE_EMBEDDINGS, SAM
 
-    PREDICTOR = util.get_sam_model()
+    PREDICTOR, SAM = util.get_sam_model(return_sam=True)
     IMAGE_EMBEDDINGS = util.precompute_image_embeddings(PREDICTOR, raw, save_path=embedding_path, ndim=2)
     util.set_precomputed(PREDICTOR, IMAGE_EMBEDDINGS)
 

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,

test/test_segment_instances.py

@@ -0,0 +1,67 @@
+import unittest
+
+import micro_sam.util as util
+import numpy as np
+
+from elf.evaluation.matching import matching
+from skimage.draw import disk
+
+
+class TestSegmentInstances(unittest.TestCase):
+
+    # create an input image with three objects
+    def _get_input(self, shape=(128, 128)):
+        mask = np.zeros(shape, dtype="uint8")
+
+        def write_object(center, radius):
+            circle = disk(center, radius, shape=shape)
+            mask[circle] = 1
+
+        center = tuple(sh // 4 for sh in shape)
+        write_object(center, radius=10)
+
+        center = tuple(sh // 2 for sh in shape)
+        write_object(center, radius=9)
+
+        center = tuple(3 * sh // 4 for sh in shape)
+        write_object(center, radius=11)
+
+        image = mask * 255
+        return mask, image
+
+    def _get_model(self):
+        predictor, sam = util.get_sam_model(model_type="vit_b", return_sam=True)
+        return predictor, sam
+
+    @unittest.skip("This test takes very long.")
+    def test_segment_instances_sam(self):
+        from micro_sam.segment_instances import segment_instances_sam
+
+        mask, image = self._get_input()
+        _, sam = self._get_model()
+
+        predicted = segment_instances_sam(sam, image)
+        self.assertGreater(matching(predicted, mask, threshold=0.75)["precision"], 0.99)
+
+    @unittest.skip("Needs some more debugging.")
+    def test_segment_instances_from_embeddings(self):
+        from micro_sam.segment_instances import segment_instances_from_embeddings
+
+        mask, image = self._get_input()
+        predictor, _ = self._get_model()
+
+        image_embeddings = util.precompute_image_embeddings(predictor, image)
+        util.set_precomputed(predictor, image_embeddings)
+
+        predicted = segment_instances_from_embeddings(predictor, image_embeddings)
+        # import napari
+        # v = napari.Viewer()
+        # v.add_image(image)
+        # v.add_labels(mask)
+        # v.add_labels(predicted)
+        # napari.run()
+        self.assertGreater(matching(predicted, mask, threshold=0.75)["precision"], 0.99)
+
+
+if __name__ == "__main__":
+    unittest.main()