Update state precomputation and extend test (#551)

Update state precomputation and extend test

Author: constantinpape

micro_sam/precompute_state.py

@@ -93,8 +93,9 @@ def cache_is_state(
     save_path: Union[str, os.PathLike],
     verbose: bool = True,
     i: Optional[int] = None,
+    skip_load: bool = False,
     **kwargs,
-) -> instance_segmentation.AMGBase:
+) -> Optional[instance_segmentation.AMGBase]:
     """Compute and cache or load the state for the automatic mask generator.
 
     Args:
@@ -105,6 +106,7 @@ def cache_is_state(
         save_path: The embedding save path. The AMG state will be stored in 'save_path/amg_state.pickle'.
         verbose: Whether to run the computation verbose.
         i: The index for which to cache the state.
+        skip_load: Skip loading the state if it is precomputed.
         kwargs: The keyword arguments for the amg class.
 
     Returns:
@@ -120,6 +122,9 @@ def cache_is_state(
 
     with h5py.File(save_path, "a") as f:
         if save_key in f:
+            if skip_load:  # Skip loading to speed this up for cases where we don't need the return val.
+                return
+
             if verbose:
                 print("Load instance segmentation state from", save_path, ":", save_key)
             g = f[save_key]
@@ -169,6 +174,7 @@ def _precompute_state_for_files(
     predictor: SamPredictor,
     input_files: Union[List[Union[os.PathLike, str]], List[np.ndarray]],
     output_path: Union[os.PathLike, str],
+    key: Optional[str] = None,
     ndim: Optional[int] = None,
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
@@ -185,14 +191,15 @@ def _precompute_state_for_files(
 
         _precompute_state_for_file(
             predictor, file_path, out_path,
-            key=None, ndim=ndim, tile_shape=tile_shape, halo=halo,
+            key=key, ndim=ndim, tile_shape=tile_shape, halo=halo,
             precompute_amg_state=precompute_amg_state, decoder=decoder,
         )
 
 
 def precompute_state(
     input_path: Union[os.PathLike, str],
     output_path: Union[os.PathLike, str],
+    pattern: Optional[str] = None,
     model_type: str = util._DEFAULT_MODEL,
     checkpoint_path: Optional[Union[os.PathLike, str]] = None,
     key: Optional[str] = None,
@@ -209,31 +216,41 @@ def precompute_state(
             In case of a container file the argument `key` must be given. In case of a folder
             it can be given to provide a glob pattern to subselect files from the folder.
         output_path: The output path were the embeddings and other state will be saved.
+        pattern: Glob pattern to select files in a folder. The embeddings will be computed
+            for each of these files. To select all files in a folder pass "*".
         model_type: The SegmentAnything model to use. Will use the standard vit_h model by default.
         checkpoint_path: Path to a checkpoint for a custom model.
         key: The key to the input file. This is needed for contaner files (e.g. hdf5 or zarr)
-            and can be used to provide a glob pattern if the input is a folder with image files.
+            or to load several images as 3d volume. Provide a glob pattern, e.g. "*.tif", for this case.
         ndim: The dimensionality of the data.
         tile_shape: Shape of tiles for tiled prediction. By default prediction is run without tiling.
         halo: Overlap of the tiles for tiled prediction.
         precompute_amg_state: Whether to precompute the state for automatic instance segmentation
             in addition to the image embeddings.
     """
-    predictor = util.get_sam_model(model_type=model_type, checkpoint_path=checkpoint_path)
-    # check if we precompute the state for a single file or for a folder with image files
-    if os.path.isdir(input_path) and Path(input_path).suffix not in (".n5", ".zarr"):
-        pattern = "*" if key is None else key
-        input_files = glob(os.path.join(input_path, pattern))
-        _precompute_state_for_files(
-            predictor, input_files, output_path,
+    predictor, state = util.get_sam_model(
+        model_type=model_type, checkpoint_path=checkpoint_path, return_state=True,
+    )
+    if "decoder_state" in state:
+        decoder = instance_segmentation.get_decoder(predictor.model.image_encoder, state["decoder_state"])
+    else:
+        decoder = None
+
+    # Check if we precompute the state for a single file or for a folder with image files.
+    if pattern is None:
+        _precompute_state_for_file(
+            predictor, input_path, output_path, key,
             ndim=ndim, tile_shape=tile_shape, halo=halo,
             precompute_amg_state=precompute_amg_state,
+            decoder=decoder,
         )
     else:
-        _precompute_state_for_file(
-            predictor, input_path, output_path, key,
+        input_files = glob(os.path.join(input_path, pattern))
+        _precompute_state_for_files(
+            predictor, input_files, output_path, key=key,
             ndim=ndim, tile_shape=tile_shape, halo=halo,
             precompute_amg_state=precompute_amg_state,
+            decoder=decoder,
         )
 
 
@@ -253,11 +270,16 @@ def main():
     parser.add_argument(
         "-e", "--embedding_path", required=True, help="The path where the embeddings will be saved."
     )
+
+    parser.add_argument(
+        "--pattern", help="Pattern / wildcard for selecting files in a folder. To select all files use '*'."
+    )
     parser.add_argument(
         "-k", "--key",
         help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
-        "for a image series it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
+        "for an image stack it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
     )
+
     parser.add_argument(
         "-m", "--model_type", default=util._DEFAULT_MODEL,
         help=f"The segment anything model that will be used, one of {available_models}."
@@ -284,8 +306,10 @@ def main():
 
     args = parser.parse_args()
     precompute_state(
-        args.input_path, args.embedding_path, args.model_type, args.checkpoint,
-        key=args.key, tile_shape=args.tile_shape, halo=args.halo, ndim=args.ndim,
+        args.input_path, args.embedding_path,
+        model_type=args.model_type, checkpoint_path=args.checkpoint,
+        pattern=args.pattern, key=args.key,
+        tile_shape=args.tile_shape, halo=args.halo, ndim=args.ndim,
         precompute_amg_state=args.precompute_amg_state,
     )
 

micro_sam/sam_annotator/_state.py

@@ -126,7 +126,9 @@ def initialize_predictor(
             if save_path is None:
                 raise RuntimeError("Require a save path to precompute the amg state")
 
-            cache_state = cache_amg_state if self.decoder is None else partial(cache_is_state, decoder=self.decoder)
+            cache_state = cache_amg_state if self.decoder is None else partial(
+                cache_is_state, decoder=self.decoder, skip_load=True,
+            )
 
             if ndim == 2:
                 self.amg = cache_state(

micro_sam/sam_annotator/_widgets.py

@@ -891,11 +891,22 @@ def _update_model(self):
             tile_shape=[self.tile_x, self.tile_y],
             halo=[self.halo_x, self.halo_y]
         )
+
+        # Set the default settings for this model in the autosegment widget if it is part of
+        # the currently used plugin.
         if "autosegment" in state.widgets:
             with_decoder = state.decoder is not None
             vutil._sync_autosegment_widget(
                 state.widgets["autosegment"], self.model_type, self.custom_weights, update_decoder=with_decoder
             )
+            # Load the AMG/AIS state if we have a 3d segmentation plugin.
+            if state.widgets["autosegment"].volumetric and with_decoder:
+                state.amg_state = vutil._load_is_state(state.embedding_path)
+            elif state.widgets["autosegment"].volumetric and not with_decoder:
+                state.amg_state = vutil._load_amg_state(state.embedding_path)
+
+        # Set the default settings for this model in the nd-segmentation widget if it is part of
+        # the currently used plugin.
         if "segment_nd" in state.widgets:
             vutil._sync_ndsegment_widget(state.widgets["segment_nd"], self.model_type, self.custom_weights)
 

micro_sam/sam_annotator/annotator_3d.py

@@ -1,58 +1,16 @@
-import os
-import pickle
-from glob import glob
-from pathlib import Path
 from typing import Optional, Tuple, Union
 
-import h5py
 import napari
 import numpy as np
 import torch
 
 from ._annotator import _AnnotatorBase
 from ._state import AnnotatorState
 from . import _widgets as widgets
-from .util import _initialize_parser, _sync_embedding_widget
+from .util import _initialize_parser, _sync_embedding_widget, _load_amg_state, _load_is_state
 from .. import util
 
 
-def _load_amg_state(embedding_path):
-    if embedding_path is None or not os.path.exists(embedding_path):
-        return {"cache_folder": None}
-
-    cache_folder = os.path.join(embedding_path, "amg_state")
-    os.makedirs(cache_folder, exist_ok=True)
-    amg_state = {"cache_folder": cache_folder}
-
-    state_paths = glob(os.path.join(cache_folder, "*.pkl"))
-    for path in state_paths:
-        with open(path, "rb") as f:
-            state = pickle.load(f)
-        i = int(Path(path).stem.split("-")[-1])
-        amg_state[i] = state
-    return amg_state
-
-
-def _load_is_state(embedding_path):
-    if embedding_path is None or not os.path.exists(embedding_path):
-        return {"cache_path": None}
-
-    cache_path = os.path.join(embedding_path, "is_state.h5")
-    is_state = {"cache_path": cache_path}
-
-    with h5py.File(cache_path, "a") as f:
-        for name, g in f.items():
-            i = int(name.split("-")[-1])
-            state = {
-                "foreground": g["foreground"][:],
-                "boundary_distances": g["boundary_distances"][:],
-                "center_distances": g["center_distances"][:],
-            }
-            is_state[i] = state
-
-    return is_state
-
-
 class Annotator3d(_AnnotatorBase):
     def _get_widgets(self):
         autosegment = widgets.AutoSegmentWidget(self._viewer, with_decoder=self._with_decoder, volumetric=True)

micro_sam/sam_annotator/util.py

@@ -1,7 +1,13 @@
 import argparse
+import os
+import pickle
 import warnings
+
+from glob import glob
+from pathlib import Path
 from typing import List, Optional, Tuple
 
+import h5py
 import napari
 import numpy as np
 
@@ -729,3 +735,40 @@ def _sync_ndsegment_widget(widget, model_type, checkpoint_path):
     for param in params:
         if param in settings:
             getattr(widget, f"{param}_param").setValue(settings[param])
+
+
+def _load_amg_state(embedding_path):
+    if embedding_path is None or not os.path.exists(embedding_path):
+        return {"cache_folder": None}
+
+    cache_folder = os.path.join(embedding_path, "amg_state")
+    os.makedirs(cache_folder, exist_ok=True)
+    amg_state = {"cache_folder": cache_folder}
+
+    state_paths = glob(os.path.join(cache_folder, "*.pkl"))
+    for path in state_paths:
+        with open(path, "rb") as f:
+            state = pickle.load(f)
+        i = int(Path(path).stem.split("-")[-1])
+        amg_state[i] = state
+    return amg_state
+
+
+def _load_is_state(embedding_path):
+    if embedding_path is None or not os.path.exists(embedding_path):
+        return {"cache_path": None}
+
+    cache_path = os.path.join(embedding_path, "is_state.h5")
+    is_state = {"cache_path": cache_path}
+
+    with h5py.File(cache_path, "a") as f:
+        for name, g in f.items():
+            i = int(name.split("-")[-1])
+            state = {
+                "foreground": g["foreground"][:],
+                "boundary_distances": g["boundary_distances"][:],
+                "center_distances": g["center_distances"][:],
+            }
+            is_state[i] = state
+
+    return is_state

test/test_sam_annotator/test_cli.py

@@ -1,8 +1,25 @@
+import os
+import platform
 import unittest
-from shutil import which
+from shutil import which, rmtree
+from subprocess import run
+
+import imageio.v3 as imageio
+import micro_sam.util as util
+import zarr
+from skimage.data import binary_blobs
 
 
 class TestCLI(unittest.TestCase):
+    model_type = "vit_t" if util.VIT_T_SUPPORT else "vit_b"
+    tmp_folder = "tmp-files"
+
+    def setUp(self):
+        os.makedirs(self.tmp_folder, exist_ok=True)
+
+    def tearDown(self):
+        rmtree(self.tmp_folder)
+
     def _test_command(self, cmd):
         self.assertTrue(which(cmd) is not None)
 
@@ -21,6 +38,47 @@ def test_image_series_annotator(self):
     def test_precompute_embeddings(self):
         self._test_command("micro_sam.precompute_embeddings")
 
+        # The filepaths can't be found on windows, probably due different filepath conventions.
+        # The actual functionality likely works despite this issue.
+        if platform.system() == "Windows":
+            return
+
+        # Create 3 images as testdata.
+        for i in range(3):
+            im_path = os.path.join(self.tmp_folder, f"image-{i}.tif")
+            image_data = binary_blobs(512).astype("uint8") * 255
+            imageio.imwrite(im_path, image_data)
+
+        # Test precomputation with a single image.
+        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
+        ])
+        self.assertTrue(os.path.exists(emb_path1))
+        with zarr.open(emb_path1, "r") as f:
+            self.assertIn("features", f)
+
+        # 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"
+        ])
+        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)
+
+        # 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"
+        ])
+        for i in range(3):
+            self.assertTrue(os.path.exists(os.path.join(emb_path3, f"image-{i}.zarr")))
+
 
 if __name__ == "__main__":
     unittest.main()