Implement state (de)serialization for the mask generators

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -72,8 +72,7 @@ class _AMGBase(ABC):
     """
     """
     def __init__(self):
-        # the data that has to be computed by the 'initialize' method
-        # of the child classes
+        # the state that has to be computed by the 'initialize' method of the child classes
         self._is_initialized = False
         self._crop_list = None
         self._crop_boxes = None
@@ -234,6 +233,17 @@ def _postprocess_masks(self, mask_data, min_mask_region_area, box_nms_thresh, cr
 
         return curr_anns
 
+    def get_state(self):
+        if not self.is_initialized:
+            raise RuntimeError("The state has not been computed yet. Call initialize first.")
+        return {"crop_list": self.crop_list, "crop_boxes": self.crop_boxes, "original_size": self.original_size}
+
+    def set_state(self, state):
+        self._crop_list = state["crop_list"]
+        self._crop_boxes = state["crop_boxes"]
+        self._original_size = state["original_size"]
+        self._is_initialized = True
+
 
 class AutomaticMaskGenerator(_AMGBase):
     """
@@ -434,7 +444,7 @@ def __init__(
         self.use_points = use_points
         self.box_extension = box_extension
 
-        # additional data that is computed by 'initialize'
+        # additional state that is set 'initialize'
         self._initial_segmentation = None
 
     def _compute_initial_segmentation(self):
@@ -558,6 +568,15 @@ def get_initial_segmentation(self):
             raise RuntimeError("AutomaticMaskGenerator has not been initialized. Call initialize first.")
         return self._resize_segmentation(self._initial_segmentation, self.original_size)
 
+    def get_state(self):
+        state = super().get_state()
+        state["initial_segmentation"] = self._initial_segmentation
+        return state
+
+    def set_state(self, state):
+        self._initial_segmentation = state["initial_segmentation"]
+        super().set_state(state)
+
 
 class TiledEmbeddingMaskGenerator(EmbeddingMaskGenerator):
     """
@@ -572,9 +591,13 @@ def __init__(
         super().__init__(predictor=predictor, **kwargs)
         self.n_threads = n_threads
         self.with_background = with_background
-        # additional data for 'initialize'
+
+        # additional state for 'initialize'
         self._tile_shape = None
         self._halo = None
+
+        # state for saving the stitched initial segmentation
+        # (this is quite complex, so we save it to only compute once)
         self._stitched_initial_segmentation = None
 
     def _compute_initial_segmentations(self, image_embeddings, i, n_tiles, verbose):
@@ -747,6 +770,17 @@ def segment_tile(_, tile_id):
         self._stitched_initial_segmentation = initial_segmentation
         return initial_segmentation
 
+    def get_state(self):
+        state = super().get_state()
+        state["tile_shape"] = self._tile_shape
+        state["halo"] = self._halo
+        return state
+
+    def set_state(self, state):
+        self._tile_shape = state["tile_shape"]
+        self._halo = state["halo"]
+        super().set_state(state)
+
 
 #
 # Functional interfaces to run instance segmentation

test/test_instance_segmentation.py

@@ -69,10 +69,19 @@ def test_automatic_mask_generator(self):
         predicted = mask_data_to_segmentation(predicted, image.shape, with_background=True)
         self.assertGreater(matching(predicted, mask, threshold=0.75)["precision"], 0.99)
 
+        # check that regenerating the segmentation works
         predicted2 = amg.generate()
         predicted2 = mask_data_to_segmentation(predicted2, image.shape, with_background=True)
         self.assertTrue(np.array_equal(predicted, predicted2))
 
+        # check that serializing and reserializing the state works
+        state = amg.get_state()
+        amg = AutomaticMaskGenerator(predictor, points_per_side=10, points_per_batch=16)
+        amg.set_state(state)
+        predicted3 = amg.generate()
+        predicted3 = mask_data_to_segmentation(predicted3, image.shape, with_background=True)
+        self.assertTrue(np.array_equal(predicted, predicted3))
+
     def test_embedding_mask_generator(self):
         from micro_sam.instance_segmentation import EmbeddingMaskGenerator, mask_data_to_segmentation
 
@@ -89,11 +98,19 @@ def test_embedding_mask_generator(self):
         initial_seg = amg.get_initial_segmentation()
         self.assertEqual(initial_seg.shape, image.shape)
 
+        # check that regenerating the segmentation works
         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))
 
+        # check that serializing and reserializing the state works
+        state = amg.get_state()
+        amg = EmbeddingMaskGenerator(predictor)
+        amg.set_state(state)
+        predicted3 = amg.generate(pred_iou_thresh=0.96)
+        predicted3 = mask_data_to_segmentation(predicted3, image.shape, with_background=True)
+        self.assertTrue(np.array_equal(predicted, predicted3))
+
     def test_tiled_embedding_mask_generator(self):
         from micro_sam.instance_segmentation import TiledEmbeddingMaskGenerator
 
@@ -112,6 +129,13 @@ def test_tiled_embedding_mask_generator(self):
         predicted2 = amg.generate(pred_iou_thresh=0.96)
         self.assertTrue(np.array_equal(predicted, predicted2))
 
+        # check that serializing and reserializing the state works
+        state = amg.get_state()
+        amg = TiledEmbeddingMaskGenerator(predictor)
+        amg.set_state(state)
+        predicted3 = amg.generate(pred_iou_thresh=0.96)
+        self.assertTrue(np.array_equal(predicted, predicted3))
+
 
 if __name__ == "__main__":
     unittest.main()