Enable calling generate multiple times in mask generators

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -1,6 +1,7 @@
 import multiprocessing as mp
 from abc import ABC
 from concurrent import futures
+from copy import deepcopy
 from typing import List, Optional
 
 import numpy as np
@@ -373,10 +374,12 @@ def generate(
     ):
         if not self.is_initialized:
             raise RuntimeError("AutomaticMaskGenerator has not been initialized. Call initialize first.")
+
         data = amg_utils.MaskData()
         for data_, crop_box in zip(self.crop_list, self.crop_boxes):
             crop_data = self._postprocess_batch(
-                data=data_, crop_box=crop_box, original_size=self.original_size,
+                data=deepcopy(data_),
+                crop_box=crop_box, original_size=self.original_size,
                 pred_iou_thresh=pred_iou_thresh,
                 stability_score_thresh=stability_score_thresh,
                 stability_score_offset=stability_score_offset,
@@ -527,7 +530,8 @@ def generate(
             raise RuntimeError("AutomaticMaskGenerator has not been initialized. Call initialize first.")
 
         data = self._postprocess_batch(
-            data=self.crop_list[0], crop_box=self.crop_boxes[0], original_size=self.original_size,
+            data=deepcopy(self.crop_list[0]), crop_box=self.crop_boxes[0],
+            original_size=self.original_size,
             pred_iou_thresh=pred_iou_thresh,
             stability_score_thresh=stability_score_thresh,
             stability_score_offset=stability_score_offset,
@@ -541,12 +545,12 @@ def generate(
     def _resize_segmentation(self, segmentation, shape):
         longest_size = max(shape)
         longest_shape = (longest_size, longest_size)
-        segmentation_ = resize(
+        resized_segmentation = resize(
             segmentation, longest_shape, order=0, preserve_range=True, anti_aliasing=False
         ).astype(segmentation.dtype)
         crop = tuple(slice(0, sh) for sh in shape)
-        segmentation_ = segmentation_[crop]
-        return segmentation_
+        resized_segmentation = resized_segmentation[crop]
+        return resized_segmentation
 
     def get_initial_segmentation(self):
         if not self.is_initialized:
@@ -667,7 +671,7 @@ def generate(
 
         def segment_tile(_, tile_id):
             tile = tiling.getBlockWithHalo(tile_id, list(self._halo)).outerBlock
-            mask_data = self._crop_list[tile_id]
+            mask_data = deepcopy(self._crop_list[tile_id])
             crop_box = self.crop_boxes[tile_id]
             this_tile_shape = tuple(end - beg for beg, end in zip(tile.begin, tile.end))
             mask_data = self._postprocess_batch(

test/test_instance_segmentation.py

@@ -64,11 +64,15 @@ def test_automatic_mask_generator(self):
 
         amg = AutomaticMaskGenerator(predictor, points_per_side=10, points_per_batch=16)
         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)
 
+        predicted2 = amg.generate()
+        predicted2 = mask_data_to_segmentation(predicted2, image.shape, with_background=True)
+        self.assertTrue(np.array_equal(predicted, predicted2))
+
     def test_embedding_mask_generator(self):
         from micro_sam.instance_segmentation import EmbeddingMaskGenerator, mask_data_to_segmentation
 
@@ -85,6 +89,11 @@ def test_embedding_mask_generator(self):
         initial_seg = amg.get_initial_segmentation()
         self.assertEqual(initial_seg.shape, image.shape)
 
+        predicted2 = amg.generate(pred_iou_thresh=0.96)
+        predicted2 = mask_data_to_segmentation(predicted2, image.shape, with_background=True)
+
+        self.assertTrue(np.array_equal(predicted, predicted2))
+
     def test_tiled_embedding_mask_generator(self):
         from micro_sam.instance_segmentation import TiledEmbeddingMaskGenerator
 
@@ -100,6 +109,9 @@ def test_tiled_embedding_mask_generator(self):
         self.assertGreater(matching(predicted, mask, threshold=0.75)["precision"], 0.99)
         self.assertEqual(initial_seg.shape, image.shape)
 
+        predicted2 = amg.generate(pred_iou_thresh=0.96)
+        self.assertTrue(np.array_equal(predicted, predicted2))
+
 
 if __name__ == "__main__":
     unittest.main()