Merge pull request #21 from computational-cell-analytics/seg-with-bg

Implement instance segmentation with background

Author: constantinpape

micro_sam/sam_annotator/annotator_2d.py

@@ -23,11 +23,11 @@ def segment_wigdet(v: Viewer):
 
 # TODO expose more parameters
 @magicgui(call_button="Segment All Objects", method={"choices": ["default", "sam", "embeddings"]})
-def autosegment_widget(v: Viewer, method: str = "default"):
+def autosegment_widget(v: Viewer, method: str = "default", with_background: bool = True):
     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)
+        seg = segment_instances.segment_instances_sam(SAM, image, with_background=with_background)
     elif method == "embeddings":
         seg = segment_instances.segment_instances_from_embeddings(PREDICTOR, IMAGE_EMBEDDINGS)
     else:

micro_sam/segment_instances.py

@@ -19,7 +19,7 @@
 #
 
 
-def segment_instances_sam(sam, image, **kwargs):
+def segment_instances_sam(sam, image, with_background=False, **kwargs):
     segmentor = SamAutomaticMaskGenerator(sam, **kwargs)
 
     image_ = util._to_image(image)
@@ -30,6 +30,12 @@ def segment_instances_sam(sam, image, **kwargs):
     for seg_id, mask in enumerate(masks, 1):
         segmentation[mask["segmentation"]] = seg_id
 
+    seg_ids, sizes = np.unique(segmentation, return_counts=True)
+    if with_background and 0 not in seg_ids:
+        bg_id = seg_ids[np.argmax(seg_ids)]
+        segmentation[segmentation == bg_id] = 0
+        vigra.analysis.relabelConsecutive(segmentation, out=segmentation)
+
     return segmentation
 
 

test/test_segment_instances.py

@@ -10,21 +10,21 @@
 class TestSegmentInstances(unittest.TestCase):
 
     # create an input image with three objects
-    def _get_input(self, shape=(128, 128)):
+    def _get_input(self, shape=(96, 96)):
         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)
+        write_object(center, radius=8)
 
         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)
+        write_object(center, radius=7)
 
         image = mask * 255
         return mask, image