Fix state precomputation for instance segmentation (#552)

Fix state precomputation for instance segmentation

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -887,7 +887,7 @@ def initialize(
             pbar_update: Callback to update an external progress bar.
         """
         _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update)
-        pbar_init(1, "Initialize instannce segmentation with decoder")
+        pbar_init(1, "Initialize instance segmentation with decoder")
 
         if image_embeddings is None:
             image_embeddings = util.precompute_image_embeddings(self._predictor, image)
@@ -1071,7 +1071,7 @@ def initialize(
         tiling = blocking([0, 0], original_size, tile_shape)
 
         _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update)
-        pbar_init(tiling.numberOfBlocks, "Initialize tiled instannce segmentation with decoder")
+        pbar_init(tiling.numberOfBlocks, "Initialize tiled instance segmentation with decoder")
 
         foreground = np.zeros(original_size, dtype="float32")
         center_distances = np.zeros(original_size, dtype="float32")

micro_sam/precompute_state.py

@@ -4,6 +4,7 @@
 import os
 import pickle
 
+from functools import partial
 from glob import glob
 from pathlib import Path
 from typing import Optional, Tuple, Union, List
@@ -159,15 +160,33 @@ def _precompute_state_for_file(
     else:
         image_data = util.load_image_data(input_path, key)
 
+    # Precompute the image embeddings.
     output_path = Path(output_path).with_suffix(".zarr")
     embeddings = util.precompute_image_embeddings(
         predictor, image_data, output_path, ndim=ndim, tile_shape=tile_shape, halo=halo,
     )
+
+    # Precompute the state for automatic instance segmnetaiton (AMG or AIS).
     if precompute_amg_state:
         if decoder is None:
-            cache_amg_state(predictor, image_data, embeddings, output_path, verbose=True)
+            cache_function = partial(
+                cache_amg_state, predictor=predictor, image_embeddings=embeddings, save_path=output_path
+            )
+        else:
+            cache_function = partial(
+                cache_is_state, predictor=predictor, decoder=decoder,
+                image_embeddings=embeddings, save_path=output_path
+            )
+
+        if ndim is None:
+            ndim = image_data.ndim
+
+        if ndim == 2:
+            cache_function(raw=image_data, verbose=True)
         else:
-            cache_is_state(predictor, decoder, image_data, embeddings, output_path, verbose=True)
+            n = image_data.shape[0]
+            for i in tqdm(range(n), total=n, desc="Precompute instance segmentation state"):
+                cache_function(raw=image_data, i=i, verbose=False)
 
 
 def _precompute_state_for_files(

test/test_sam_annotator/test_cli.py

@@ -11,7 +11,7 @@
 
 
 class TestCLI(unittest.TestCase):
-    model_type = "vit_t" if util.VIT_T_SUPPORT else "vit_b"
+    model_type = "vit_t_lm" if util.VIT_T_SUPPORT else "vit_b_lm"
     tmp_folder = "tmp-files"
 
     def setUp(self):
@@ -53,31 +53,37 @@ def test_precompute_embeddings(self):
         emb_path1 = os.path.join(self.tmp_folder, "embedddings1.zarr")
         run([
             "micro_sam.precompute_embeddings", "-i", im_path, "-e", emb_path1,
-            "-m", self.model_type
+            "-m", self.model_type, "--precompute_amg_state"
         ])
         self.assertTrue(os.path.exists(emb_path1))
         with zarr.open(emb_path1, "r") as f:
             self.assertIn("features", f)
+        ais_path = os.path.join(emb_path1, "is_state.h5")
+        self.assertTrue(os.path.exists(ais_path))
 
         # Test precomputation with image stack.
         emb_path2 = os.path.join(self.tmp_folder, "embedddings2.zarr")
         run([
             "micro_sam.precompute_embeddings", "-i", self.tmp_folder, "-e", emb_path2,
-            "-m", self.model_type, "-k", "*.tif"
+            "-m", self.model_type, "-k", "*.tif", "--precompute_amg_state"
         ])
         self.assertTrue(os.path.exists(emb_path2))
         with zarr.open(emb_path2, "r") as f:
             self.assertIn("features", f)
             self.assertEqual(f["features"].shape[0], 3)
+        ais_path = os.path.join(emb_path2, "is_state.h5")
+        self.assertTrue(os.path.exists(ais_path))
 
         # Test precomputation with pattern to process multiple image.
         emb_path3 = os.path.join(self.tmp_folder, "embedddings3")
         run([
             "micro_sam.precompute_embeddings", "-i", self.tmp_folder, "-e", emb_path3,
-            "-m", self.model_type, "--pattern", "*.tif"
+            "-m", self.model_type, "--pattern", "*.tif", "--precompute_amg_state"
         ])
         for i in range(3):
             self.assertTrue(os.path.exists(os.path.join(emb_path3, f"image-{i}.zarr")))
+            ais_path = os.path.join(emb_path3, f"image-{i}.zarr", "is_state.h5")
+            self.assertTrue(os.path.exists(ais_path))
 
 
 if __name__ == "__main__":