Add test for tiled embedding mask generator and add some type annotations

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -320,7 +320,13 @@ def _process_crop(self, image, crop_box, crop_layer_idx, verbose, precomputed_em
         return data
 
     @torch.no_grad()
-    def initialize(self, image: np.ndarray, image_embeddings=None, i=None, verbose=False):
+    def initialize(
+        self,
+        image: np.ndarray,
+        image_embeddings=None,
+        i: Optional[int] = None,
+        verbose: bool = False
+    ):
         """
         """
         original_size = image.shape[:2]
@@ -403,14 +409,14 @@ class EmbeddingMaskGenerator(_AMGBase):
     def __init__(
         self,
         predictor: SamPredictor,
-        offsets=None,
-        min_initial_size=0,
-        distance_type="l2",
-        bias=0.0,
-        use_box=True,
-        use_mask=True,
-        use_points=False,
-        box_extension=0.05,
+        offsets: Optional[List[List[int]]] = None,
+        min_initial_size: int = 0,
+        distance_type: str = "l2",
+        bias: float = 0.0,
+        use_box: bool = True,
+        use_mask: bool = True,
+        use_points: bool = False,
+        box_extension: float = 0.05,
     ):
         super().__init__()
 
@@ -475,7 +481,13 @@ def _compute_mask_data(self, initial_segmentation, original_size, verbose):
         return mask_data
 
     @torch.no_grad()
-    def initialize(self, image: np.ndarray, image_embeddings=None, i=None, verbose=False):
+    def initialize(
+        self,
+        image: np.ndarray,
+        image_embeddings=None,
+        i: Optional[int] = None,
+        verbose: bool = False
+    ):
         """
         """
         original_size = image.shape[:2]
@@ -545,10 +557,17 @@ def get_initial_segmentation(self):
 class TiledEmbeddingMaskGenerator(EmbeddingMaskGenerator):
     """
     """
-    def __init__(self, n_threads=mp.cpu_count(), with_background=True, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        predictor: SamPredictor,
+        n_threads: int = mp.cpu_count(),
+        with_background: bool = True,
+        **kwargs
+    ):
+        super().__init__(predictor=predictor, **kwargs)
         self.n_threads = n_threads
         self.with_background = with_background
+        # additional data for 'initialize'
         self._tile_shape = None
         self._halo = None
         self._stitched_initial_segmentation = None
@@ -599,14 +618,17 @@ def initialize(
         tile_shape: List[int],
         halo: List[int],
         image_embeddings=None,
-        i=None,
-        verbose=False,
+        i: Optional[int] = None,
+        verbose: bool = False,
+        embedding_save_path: Optional[str] = None,
     ):
         """
         """
         original_size = image.shape[:2]
         if image_embeddings is None:
-            image_embeddings = util.precompute_image_embeddings(self.predictor, image, tile_shape=tile_shape, halo=halo)
+            image_embeddings = util.precompute_image_embeddings(
+                self.predictor, image, tile_shape=tile_shape, halo=halo, save_path=embedding_save_path
+            )
 
         tiling = blocking([0, 0], original_size, tile_shape)
         n_tiles = tiling.numberOfBlocks

test/test_instance_segmentation.py

@@ -1,4 +1,6 @@
+import os
 import unittest
+from shutil import rmtree
 
 import micro_sam.util as util
 import numpy as np
@@ -9,6 +11,8 @@
 
 
 class TestInstanceSegmentation(unittest.TestCase):
+    embedding_path = "./tmp_embeddings.zarr"
+
     # create an input image with three objects
     @staticmethod
     def _get_input(shape=(256, 256)):
@@ -32,9 +36,12 @@ def write_object(center, radius):
         return mask, image
 
     @staticmethod
-    def _get_model(image):
+    def _get_model(image, tile_shape=None, halo=None, save_path=None):
         predictor = util.get_sam_model(model_type="vit_b")
-        image_embeddings = util.precompute_image_embeddings(predictor, image)
+
+        image_embeddings = util.precompute_image_embeddings(
+            predictor, image, tile_shape=tile_shape, halo=halo, save_path=save_path
+        )
         return predictor, image_embeddings
 
     # we compute the default mask and predictor once for the class
@@ -44,6 +51,11 @@ def setUpClass(cls):
         cls.mask, cls.image = cls._get_input()
         cls.predictor, cls.image_embeddings = cls._get_model(cls.image)
 
+    # remove temp embeddings if any
+    def tearDown(self):
+        if os.path.exists(self.embedding_path):
+            rmtree(self.embedding_path)
+
     def test_automatic_mask_generator(self):
         from micro_sam.instance_segmentation import AutomaticMaskGenerator, mask_data_to_segmentation
 
@@ -73,6 +85,21 @@ def test_embedding_mask_generator(self):
         initial_seg = amg.get_initial_segmentation()
         self.assertEqual(initial_seg.shape, image.shape)
 
+    def test_tiled_embedding_mask_generator(self):
+        from micro_sam.instance_segmentation import TiledEmbeddingMaskGenerator
+
+        tile_shape, halo = (576, 576), (64, 64)
+        mask, image = self._get_input(shape=(1024, 1024))
+        predictor, image_embeddings = self._get_model(image, tile_shape, halo, self.embedding_path)
+
+        amg = TiledEmbeddingMaskGenerator(predictor)
+        amg.initialize(image, image_embeddings=image_embeddings, tile_shape=tile_shape, halo=halo)
+        predicted = amg.generate(pred_iou_thresh=0.96)
+        initial_seg = amg.get_initial_segmentation()
+
+        self.assertGreater(matching(predicted, mask, threshold=0.75)["precision"], 0.99)
+        self.assertEqual(initial_seg.shape, image.shape)
+
 
 if __name__ == "__main__":
     unittest.main()