Enable passing pre-computed embeddings to AutomaticMaskGenerator

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -288,12 +288,14 @@ def _process_batch(self, points, im_size):
 
         return data
 
-    def _process_crop(self, image, crop_box, crop_layer_idx, verbose):
+    def _process_crop(self, image, crop_box, crop_layer_idx, verbose, precomputed_embeddings):
         # crop the image and calculate embeddings
         x0, y0, x1, y1 = crop_box
         cropped_im = image[y0:y1, x0:x1, :]
         cropped_im_size = cropped_im.shape[:2]
-        self.predictor.set_image(cropped_im)
+
+        if not precomputed_embeddings:
+            self.predictor.set_image(cropped_im)
 
         # get the points for this crop
         points_scale = np.array(cropped_im_size)[None, ::-1]
@@ -312,23 +314,39 @@ def _process_crop(self, image, crop_box, crop_layer_idx, verbose):
             data.cat(batch_data)
             del batch_data
 
-        self.predictor.reset_image()
+        if not precomputed_embeddings:
+            self.predictor.reset_image()
+
         return data
 
-    # TODO enable initializeing with embeddings
-    # (which can be done for only a single crop box)
     @torch.no_grad()
-    def initialize(self, image: np.ndarray, verbose=False):
+    def initialize(self, image: np.ndarray, image_embeddings=None, i=None, embedding_path=None, verbose=False):
         """
         """
-        image = util._to_image(image)
         original_size = image.shape[:2]
         crop_boxes, layer_idxs = amg_utils.generate_crop_boxes(
             original_size, self.crop_n_layers, self.crop_overlap_ratio
         )
+
+        # we can set fixed image embeddings if we only have a single crop box
+        # (which is the default setting)
+        # otherwise we have to recompute the embeddings for each crop and can't precompute
+        if len(crop_boxes) == 1:
+            if image_embeddings is None:
+                image_embeddings = util.precompute_image_embeddings(self.predictor, image, save_path=embedding_path)
+            util.set_precomputed(self.predictor, image_embeddings, i=i)
+            precomputed_embeddings = True
+        else:
+            precomputed_embeddings = False
+
+        # we need to cast to the image representation that is compatible with SAM
+        image = util._to_image(image)
+
         crop_list = []
         for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
-            crop_data = self._process_crop(image, crop_box, layer_idx, verbose=verbose)
+            crop_data = self._process_crop(
+                image, crop_box, layer_idx, verbose=verbose, precomputed_embeddings=precomputed_embeddings
+            )
             crop_list.append(crop_data)
 
         self._is_initialized = True

test/test_instance_segmentation.py

@@ -9,9 +9,9 @@
 
 
 class TestInstanceSegmentation(unittest.TestCase):
-
     # create an input image with three objects
-    def _get_input(self, shape=(512, 512)):
+    @staticmethod
+    def _get_input(shape=(256, 256)):
         mask = np.zeros(shape, dtype="uint8")
 
         def write_object(center, radius):
@@ -31,30 +31,40 @@ def write_object(center, radius):
         mask = label(mask)
         return mask, image
 
-    def _get_model(self):
-        return util.get_sam_model(model_type="vit_b", return_sam=False)
+    @staticmethod
+    def _get_model(image):
+        predictor = util.get_sam_model(model_type="vit_b")
+        image_embeddings = util.precompute_image_embeddings(predictor, image)
+        return predictor, image_embeddings
+
+    # we compute the default mask and predictor once for the class
+    # so that we don't have to precompute it every time
+    @classmethod
+    def setUpClass(cls):
+        cls.mask, cls.image = cls._get_input()
+        cls.predictor, cls.image_embeddings = cls._get_model(cls.image)
 
     def test_automatic_mask_generator(self):
         from micro_sam.instance_segmentation import AutomaticMaskGenerator, mask_data_to_segmentation
 
-        mask, image = self._get_input(shape=(256, 256))
-        predictor = self._get_model()
+        mask, image = self.mask, self.image
+        predictor, image_embeddings = self.predictor, self.image_embeddings
 
         amg = AutomaticMaskGenerator(predictor, points_per_side=10, points_per_batch=16)
-        amg.initialize(image, verbose=False)
+        amg.initialize(image, image_embeddings=image_embeddings, verbose=False)
         predicted = amg.generate()
         predicted = mask_data_to_segmentation(predicted, image.shape, with_background=True)
 
         self.assertGreater(matching(predicted, mask, threshold=0.75)["precision"], 0.99)
 
-    def test_embedding_based_mask_generator(self):
+    def test_embedding_mask_generator(self):
         from micro_sam.instance_segmentation import EmbeddingMaskGenerator, mask_data_to_segmentation
 
-        mask, image = self._get_input()
-        predictor = self._get_model()
+        mask, image = self.mask, self.image
+        predictor, image_embeddings = self.predictor, self.image_embeddings
 
         amg = EmbeddingMaskGenerator(predictor)
-        amg.initialize(image, verbose=False)
+        amg.initialize(image, image_embeddings=image_embeddings, verbose=False)
         predicted = amg.generate(pred_iou_thresh=0.96)
         predicted = mask_data_to_segmentation(predicted, image.shape, with_background=True)