diff --git a/examples/annotator_2d.py b/examples/annotator_2d.py index a90e612f5..e356b317f 100644 --- a/examples/annotator_2d.py +++ b/examples/annotator_2d.py @@ -26,7 +26,10 @@ def livecell_annotator(use_finetuned_model): embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-livecell.zarr") model_type = "vit_b" - annotator(image, embedding_path=embedding_path, model_type=model_type, precompute_amg_state=True) + annotator( + image, embedding_path=embedding_path, model_type=model_type, + precompute_autoseg_state=True, + ) def hela_2d_annotator(use_finetuned_model): diff --git a/examples/annotator_3d.py b/examples/annotator_3d.py index 77f841ea6..195d223a0 100644 --- a/examples/annotator_3d.py +++ b/examples/annotator_3d.py @@ -21,14 +21,17 @@ def em_3d_annotator(use_finetuned_model): if use_finetuned_model: embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-lucchi-vit_b_em_organelles.zarr") model_type = "vit_b_em_organelles" - precompute_amg_state = True + precompute_autoseg_state = True else: embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-lucchi.zarr") model_type = "vit_h" - precompute_amg_state = False + precompute_autoseg_state = False # start the annotator, cache the embeddings - annotator(raw, embedding_path=embedding_path, model_type=model_type, precompute_amg_state=precompute_amg_state) + annotator( + raw, embedding_path=embedding_path, model_type=model_type, + precompute_autoseg_state=precompute_autoseg_state, + ) def main(): diff --git a/examples/image_series_annotator.py b/examples/image_series_annotator.py index d234894e7..1e2b96d69 100644 --- a/examples/image_series_annotator.py +++ b/examples/image_series_annotator.py @@ -27,7 +27,7 @@ def series_annotation(use_finetuned_model): pattern="*.tif", embedding_path=embedding_path, model_type=model_type, - precompute_amg_state=False, + precompute_autoseg_state=False, ) diff --git a/micro_sam/_cli.py b/micro_sam/_cli.py index a8c880b18..8f8ee06b3 100644 --- a/micro_sam/_cli.py +++ b/micro_sam/_cli.py @@ -112,7 +112,8 @@ def _interactive_options(f): help="The number of spatial dimensions (2 or 3). If not given, auto-detected from the image shape." ) @click.option( - "--precompute_amg_state", is_flag=True, default=False, + "--precompute_autoseg_state", + "precompute_autoseg_state", is_flag=True, default=False, help="Whether to precompute the state for automatic instance segmentation (longer start-up, faster first run)." ) @click.option( @@ -122,7 +123,7 @@ def _interactive_options(f): @_interactive_options def annotator_segmentation( input_, key, embedding_path, model_type, checkpoint_path, decoder_path, device, tile_shape, halo, - segmentation_result, segmentation_key, ndim, precompute_amg_state, prefer_decoder, + segmentation_result, segmentation_key, ndim, precompute_autoseg_state, prefer_decoder, ): """Interactively segment a 2d or 3d image.""" from .util import load_image_data @@ -140,7 +141,7 @@ def annotator_segmentation( model_type=model_type or DEFAULT_MODEL, tile_shape=tile_shape or None, halo=halo or None, - precompute_amg_state=precompute_amg_state, + precompute_autoseg_state=precompute_autoseg_state, checkpoint_path=checkpoint_path, decoder_path=decoder_path, device=device, @@ -286,7 +287,12 @@ def annotator_object_classification( ) @click.option("--tile_shape", type=int, nargs=2, default=None, help="The tile shape for tiled prediction.") @click.option("--overlap", "halo", type=int, nargs=2, default=None, help="The tile overlap for tiled prediction.") -@click.option("--precompute_amg_state", is_flag=True, default=False, help="Whether to precompute the AMG state.") +@click.option( + "--precompute_autoseg_state", + "precompute_autoseg_state", is_flag=True, default=False, + help="Whether to precompute the automatic segmentation state (AMG masks, or decoder predictions " + "if the model has a decoder)." +) @click.option( "--prefer_decoder", is_flag=True, default=True, flag_value=False, help="Whether to use decoder based instance segmentation if the model has an additional decoder for that purpose." @@ -297,8 +303,8 @@ def annotator_object_classification( ) def annotator_batch( input_folder, output_folder, task, ndim, pattern, initial_segmentation_folder, initial_segmentation_pattern, - embedding_path, model_type, checkpoint_path, device, tile_shape, halo, precompute_amg_state, prefer_decoder, - skip_segmented, + embedding_path, model_type, checkpoint_path, device, tile_shape, halo, + precompute_autoseg_state, prefer_decoder, skip_segmented, ): """Annotate multiples images within a folder. @@ -322,7 +328,8 @@ def annotator_batch( initial_segmentation_folder=initial_segmentation_folder, initial_segmentation_pattern=initial_segmentation_pattern, embedding_path=embedding_path, model_type=model_type, - tile_shape=tile_shape, halo=halo, precompute_amg_state=precompute_amg_state, + tile_shape=tile_shape, halo=halo, + precompute_autoseg_state=precompute_autoseg_state, checkpoint_path=checkpoint_path, device=device, prefer_decoder=prefer_decoder, skip_segmented=skip_segmented, ) @@ -337,7 +344,8 @@ def annotator_batch( batch_tracking_annotator( images, output_folder, model_type=model_type, embedding_path=embedding_path, tile_shape=tile_shape, halo=halo, checkpoint_path=checkpoint_path, device=device, - precompute_amg_state=precompute_amg_state, skip_done=skip_segmented, + precompute_autoseg_state=precompute_autoseg_state, + skip_done=skip_segmented, ) elif task == "pixel-classification": from .sam_annotator.pixel_classifier import batch_pixel_classifier @@ -756,8 +764,20 @@ def inference_object_classification( "-n", "--ndim", type=int, default=None, help="The number of spatial dimensions. Specify this if your data has a channel dimension." ) -def precompute_embeddings(input_path, embedding_path, pattern, key, model_type, checkpoint_path, ndim): - """Precompute image embeddings.""" +@click.option( + "--precompute_autoseg_state", + "precompute_autoseg_state", is_flag=True, default=False, + help="Whether to also precompute the automatic-segmentation state in the embedding Zarr (SAM2 only)." +) +@click.option( + "--prefer_decoder", is_flag=True, default=True, flag_value=False, + help="Whether to use decoder-based state (AIS) when the model has a decoder, instead of grid-based AMG." +) +def precompute_embeddings( + input_path, embedding_path, pattern, key, model_type, checkpoint_path, ndim, + precompute_autoseg_state, prefer_decoder, +): + """Precompute image embeddings (and optionally the automatic-segmentation state).""" from .precompute_state import precompute_state from .v2.util import _DEFAULT_MODEL @@ -765,6 +785,8 @@ def precompute_embeddings(input_path, embedding_path, pattern, key, model_type, input_path, embedding_path, model_type=model_type or _DEFAULT_MODEL, checkpoint_path=checkpoint_path, pattern=pattern, key=key, ndim=ndim, + precompute_autoseg_state=precompute_autoseg_state, + prefer_decoder=prefer_decoder, ) diff --git a/micro_sam/precompute_state.py b/micro_sam/precompute_state.py index 7a3bd64a2..c4d8aa380 100644 --- a/micro_sam/precompute_state.py +++ b/micro_sam/precompute_state.py @@ -5,14 +5,12 @@ import pickle from glob import glob from pathlib import Path -from functools import partial -from typing import Optional, Tuple, Union, List +from typing import Optional, Tuple, Union import h5py import numpy as np import torch -import torch.nn as nn from segment_anything.predictor import SamPredictor @@ -22,7 +20,6 @@ from tqdm import tqdm from . import util -from .v1.util import precompute_image_embeddings from .v1 import instance_segmentation @@ -34,7 +31,7 @@ def cache_amg_state( verbose: bool = True, i: Optional[int] = None, **kwargs, -) -> instance_segmentation.AMGBase: +) -> instance_segmentation.AutoSegBase: """Compute and cache or load the state for the automatic mask generator. Args: @@ -99,19 +96,19 @@ def cache_is_state( i: Optional[int] = None, skip_load: bool = False, **kwargs, -) -> Optional[instance_segmentation.AMGBase]: - """Compute and cache or load the state for the automatic mask generator. +) -> Optional[instance_segmentation.AutoSegBase]: + """Compute and cache or load the state for the decoder-based instance segmentation. Args: predictor: The Segment Anything predictor. decoder: The instance segmentation decoder. raw: The image data. image_embeddings: The image embeddings. - save_path: The embedding save path. The AMG state will be stored in 'save_path/amg_state.pickle'. + save_path: The embedding save path. The state will be stored in 'save_path/is_state.h5'. verbose: Whether to run the computation verbose. By default, set to 'True'. i: The index for which to cache the state. skip_load: Skip loading the state if it is precomputed. By default, set to 'False'. - kwargs: The keyword arguments for the amg class. + kwargs: The keyword arguments for the instance segmentation class. Returns: The instance segmentation class with the cached state. @@ -157,70 +154,504 @@ def cache_is_state( return amg -def _precompute_state_for_file( - predictor, input_path, output_path, key, ndim, tile_shape, halo, precompute_amg_state, decoder, verbose +AUTOSEG_STATE_GROUP = "autoseg_state" +AUTOSEG_STATE_ATTRIBUTE = "autoseg_state" +AUTOSEG_STATE_VERSION = 1 + + +def _autoseg_state_key(i): + """Return the Zarr key for a whole-image/volume state or a per-slice state.""" + return "state" if i is None else f"state-{i}" + + +def _get_autoseg_state_group(embeddings, mode, create=False): + """Return the mode-specific state group from an open embedding Zarr.""" + if mode not in ("amg", "ais"): + raise ValueError(f"Invalid automatic segmentation mode: {mode}") + + if create: + root = embeddings.require_group(AUTOSEG_STATE_GROUP) + group = root.require_group(mode) + group.attrs["version"] = AUTOSEG_STATE_VERSION + return group + + if AUTOSEG_STATE_GROUP not in embeddings: + return None + root = embeddings[AUTOSEG_STATE_GROUP] + return root.get(mode, None) + + +def _record_autoseg_state(embeddings, mode, group): + """Record which automatic-segmentation states are present in the embedding metadata.""" + modes = embeddings.attrs.get(AUTOSEG_STATE_ATTRIBUTE, []) + if isinstance(modes, str): + modes = [modes] + modes = sorted(set(modes) | {mode}) + embeddings.attrs[AUTOSEG_STATE_ATTRIBUTE] = modes + group.attrs["state_count"] = len(group) + + +# Embedding metadata that defines the identity of the embeddings the state was computed from. If any +# of these change (image, tiling, model, normalization) the cached state is stale and must not be +# reused. 'micro_sam_version' is deliberately excluded (a version bump alone does not invalidate it). +EMBEDDING_SIGNATURE_KEYS = ("data_signature", "tile_shape", "halo", "normalization", "model_name", "model_hash") + + +def _embedding_signature(save_path): + """A stable signature of the embeddings at `save_path`, or None if it cannot be read. + + Stamped into the cached automatic-segmentation state so it is not reused against embeddings that + were recomputed with different settings (e.g. tiling toggled) but the same image data. + """ + if save_path is None: + return None + try: + attrs = dict(util._open_embeddings(save_path, mode="r").attrs) + except Exception: + return None + return "|".join(f"{key}={attrs.get(key)}" for key in EMBEDDING_SIGNATURE_KEYS) + + +def _signature_matches(cached, requested): + """Whether a cached signature is compatible: equal, or unknown on either side (lenient).""" + return cached is None or requested is None or cached == requested + + +def _save_amg_state_v2(segmenter, save_path, key, embedding_signature=None): + """Serialize an AMG state as a byte array inside the embedding Zarr.""" + state = segmenter.get_state() + payload = np.frombuffer(pickle.dumps(state, protocol=pickle.HIGHEST_PROTOCOL), dtype="uint8") + + embeddings = util._open_embeddings(save_path, mode="a") + group = _get_autoseg_state_group(embeddings, "amg", create=True) + if key in group: + del group[key] + chunks = (min(len(payload), 1024 ** 2),) + dataset = util._create_dataset_with_data(group, key, data=payload, chunks=chunks) + if embedding_signature is not None: + dataset.attrs["embedding_signature"] = embedding_signature + _record_autoseg_state(embeddings, "amg", group) + + +def _load_amg_state_v2(save_path, key): + """Load one AMG state from the embedding Zarr, without touching other slice states.""" + embeddings = util._open_embeddings(save_path, mode="r") + group = _get_autoseg_state_group(embeddings, "amg") + if group is None or key not in group: + return None + dataset = group[key] + state = pickle.loads(np.asarray(dataset[:], dtype="uint8").tobytes()) + state["embedding_signature"] = dataset.attrs.get("embedding_signature", None) + return state + + +def _save_ais_state_v2(segmenter, save_path, key, model_type, embedding_signature=None): + """Store decoder predictions as an array inside the embedding Zarr.""" + prediction = segmenter.get_state()["prediction"] + embeddings = util._open_embeddings(save_path, mode="a") + group = _get_autoseg_state_group(embeddings, "ais", create=True) + if key in group: + del group[key] + state_group = group.create_group(key) + + # Keep channel and z chunks independent so writing and later reading volumetric predictions + # does not require materializing one very large compressed chunk. + chunks = list(prediction.shape) + chunks[0] = 1 + if prediction.ndim == 4: + chunks[1] = 1 + chunks[-2] = min(chunks[-2], 512) + chunks[-1] = min(chunks[-1], 512) + util._create_dataset_with_data(state_group, "prediction", data=prediction, chunks=tuple(chunks)) + + # Record which model produced the prediction so it is not reused with a different decoder. + if model_type is not None: + state_group.attrs["model_type"] = model_type + if embedding_signature is not None: + state_group.attrs["embedding_signature"] = embedding_signature + _record_autoseg_state(embeddings, "ais", group) + + +def _load_ais_state_v2(save_path, key): + """Load one AIS prediction from the embedding Zarr on demand.""" + embeddings = util._open_embeddings(save_path, mode="r") + group = _get_autoseg_state_group(embeddings, "ais") + if group is None or key not in group: + return None + state_group = group[key] + return { + "prediction": state_group["prediction"][:], + "model_type": state_group.attrs.get("model_type", None), + "embedding_signature": state_group.attrs.get("embedding_signature", None), + } + + +def _has_autoseg_state(save_path, mode, state_count=1): + """Check cache completeness from Zarr metadata without loading state arrays or bitstreams.""" + if save_path is None or not os.path.exists(save_path): + return False + # Any failure to read the metadata (unreadable / partial zarr, backend error) means we cannot + # confirm a usable cache, so we treat it as absent rather than propagate into the GUI gating. + try: + embeddings = util._open_embeddings(save_path, mode="r") + group = _get_autoseg_state_group(embeddings, mode) + if group is None: + return False + return int(group.attrs.get("state_count", len(group))) >= state_count + except Exception: + return False + + +def _ais_state_matches(state, model_type): + """Whether a cached AIS state may be reused for `model_type`. + + The AIS prediction depends only on the decoder and the embeddings, so the only staleness risk is + reusing it with a different decoder. We reuse the cached state unless both the stored and the + requested `model_type` are known and differ (a state written without a signature is reused). + """ + cached = state.get("model_type") + return cached is None or model_type is None or cached == model_type + + +def cache_autoseg_state(mode, model_or_decoder, raw, image_embeddings, save_path, **kwargs): + """Compute, cache or load the SAM2 automatic-segmentation state for one image / slice / volume. + + A single entry point over the two modes; both reuse a matching cached state instead of recomputing. + - 'amg': grid-based mask generation. `model_or_decoder` is the SAM2 model; extra kwargs are the + AMG parameters (see `_cache_amg_state_v2`). + - 'ais': UniSAM2 decoder prediction. `model_or_decoder` is the decoder; `ndim` is required + (see `_cache_ais_state_v2`). + + Args: + mode: The automatic-segmentation mode, 'amg' or 'ais'. + model_or_decoder: The SAM2 model (AMG) or the UniSAM2 decoder (AIS). + raw: The image data. + image_embeddings: The (optionally precomputed) image embeddings. + save_path: The embedding save path used to store the state, or None to skip caching. + kwargs: Mode-specific keyword arguments forwarded to the underlying cache function. + + Returns: + The segmenter with the (cached or freshly computed) state set. + """ + if mode == "amg": + return _cache_amg_state_v2(model_or_decoder, raw, image_embeddings, save_path, **kwargs) + if mode == "ais": + return _cache_ais_state_v2(model_or_decoder, raw, image_embeddings, save_path, **kwargs) + raise ValueError(f"Invalid automatic-segmentation state mode: {mode!r} (expected 'amg' or 'ais').") + + +def _save_autoseg_state(mode, segmenter, save_path, key, **kwargs): + """Save an autoseg state into the embedding Zarr: 'amg' as pickled masks, 'ais' as a decoder array.""" + if mode == "amg": + return _save_amg_state_v2(segmenter, save_path, key, **kwargs) + if mode == "ais": + return _save_ais_state_v2(segmenter, save_path, key, **kwargs) + raise ValueError(f"Invalid automatic-segmentation state mode: {mode!r} (expected 'amg' or 'ais').") + + +def _load_autoseg_state(mode, save_path, key): + """Load one autoseg state (AMG masks or AIS prediction) from the embedding Zarr, or None if absent.""" + if mode == "amg": + return _load_amg_state_v2(save_path, key) + if mode == "ais": + return _load_ais_state_v2(save_path, key) + raise ValueError(f"Invalid automatic-segmentation state mode: {mode!r} (expected 'amg' or 'ais').") + + +def _cache_amg_state_v2( + model: torch.nn.Module, + raw: np.ndarray, + image_embeddings: Optional[dict], + save_path: Optional[Union[str, os.PathLike]], + model_type: Optional[str] = None, + i: Optional[int] = None, + state_index: Optional[int] = None, + is_tiled: Optional[bool] = None, + tile_shape: Optional[Tuple[int, int]] = None, + halo: Optional[Tuple[int, int]] = None, + points_per_side: int = 32, + pred_iou_thresh: float = 0.8, + stability_score_thresh: float = 0.9, + verbose: bool = True, + pbar_init: Optional[callable] = None, + pbar_update: Optional[callable] = None, + **kwargs, ): - if isinstance(input_path, np.ndarray): - image_data = input_path - else: - image_data = util.load_image_data(input_path, key) + """Compute and cache, or load, the SAM2 grid-based (AMG) automatic-segmentation state. - # Precompute the image embeddings. - output_path = Path(output_path).with_suffix(".zarr") - embeddings = precompute_image_embeddings( - predictor, image_data, output_path, ndim=ndim, tile_shape=tile_shape, halo=halo, verbose=verbose + The SAM2 counterpart of `cache_amg_state`. The state (the predicted masks) is stored in the + embedding Zarr under 'autoseg_state/amg'. A cached state is reused only if it was + computed with the same AMG parameters; otherwise it is recomputed and overwritten. Pass + 'save_path=None' to compute in memory without caching. + + Args: + model: The SAM2 model, loaded via `micro_sam.v2.util.get_sam2_model`. + raw: The image data. + image_embeddings: The (optionally precomputed) image embeddings. When None the segmenter + computes them from `raw`. + save_path: The embedding save path used to store the state, or None to skip caching. + model_type: The SAM2 model type, e.g. 'hvit_t'. Recorded in the cached state. + i: The slice index passed to the segmenter's `initialize` (for reusing volume embeddings). + state_index: The index used for the on-disk state (defaults to `i`). Set this to identify a + slice on disk when `i`/`image_embeddings` do not (e.g. a prebuilt single-slice embedding). + is_tiled: Whether to use the tiled segmenter. By default inferred from the embeddings. + tile_shape: The tile shape for the tiled segmenter. + halo: The tile overlap for the tiled segmenter. + points_per_side: The number of grid points sampled per image side. + pred_iou_thresh: The predicted-IoU filter threshold. + stability_score_thresh: The stability-score filter threshold. + verbose: Whether to run verbose. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + kwargs: Additional keyword arguments for the AMG segmenter. + + Returns: + The AMG segmenter with the (cached or freshly computed) state set. + """ + from .v2.instance_segmentation import get_amg_segmenter + + if is_tiled is None: + is_tiled = image_embeddings is not None and image_embeddings.get("input_size") is None + + segmenter = get_amg_segmenter( + model, is_tiled=is_tiled, model_type=model_type, points_per_side=points_per_side, + pred_iou_thresh=pred_iou_thresh, stability_score_thresh=stability_score_thresh, **kwargs, ) - # Precompute the state for automatic instance segmnetaiton (AMG or AIS). - if precompute_amg_state: - if decoder is None: - 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 - ) + key_index = i if state_index is None else state_index + key, signature = None, None + if save_path is not None: + key = _autoseg_state_key(key_index) + signature = _embedding_signature(save_path) + state = _load_amg_state_v2(save_path, key) + if state is not None: + matches = state.get("params") == segmenter._amg_params + matches = matches and _signature_matches(state.get("embedding_signature"), signature) + if matches: + if verbose: + print("Load the AMG state from", save_path, ":", key) + segmenter.set_state(state) + return segmenter - if ndim is None: - ndim = image_data.ndim + if verbose: + print("Precomputing the state for automatic mask generation.") - if ndim == 2: - cache_function(raw=image_data, verbose=verbose) - else: - n = image_data.shape[0] - for i in tqdm(range(n), total=n, desc="Precompute instance segmentation state", disable=not verbose): - cache_function(raw=image_data, i=i, verbose=False) + init_kwargs = {"tile_shape": tile_shape, "halo": halo} if is_tiled else {} + segmenter.initialize( + raw, image_embeddings=image_embeddings, i=i, verbose=verbose, + pbar_init=pbar_init, pbar_update=pbar_update, **init_kwargs, + ) + if key is not None: + _save_amg_state_v2(segmenter, save_path, key, embedding_signature=signature) + return segmenter -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, +def _cache_amg_slice(segmenter, save_path, i, init_fn, embedding_signature=None): + """Load slice `i`'s AMG state from `save_path` if present and matching, else init and save. + + Used by `micro_sam.v2.instance_segmentation.automatic_3d_segmentation` to cache the per-slice + grid-prediction state of a volume. `init_fn(i)` runs the (expensive) `initialize` for the slice. + """ + key = _autoseg_state_key(i) + state = _load_amg_state_v2(save_path, key) + if state is not None: + matches = state.get("params") == segmenter._amg_params + matches = matches and _signature_matches(state.get("embedding_signature"), embedding_signature) + if matches: + segmenter.set_state(state) + return + init_fn(i) + _save_amg_state_v2(segmenter, save_path, key, embedding_signature=embedding_signature) + + +def _cache_amg_volume_state( + model: torch.nn.Module, + get_slice: callable, + n_slices: int, + image_embeddings: Optional[dict], + save_path: Union[str, os.PathLike], + model_type: Optional[str] = None, + is_tiled: Optional[bool] = None, + tile_shape: Optional[Tuple[int, int]] = None, + halo: Optional[Tuple[int, int]] = None, + verbose: bool = True, + pbar_init: Optional[callable] = None, + pbar_update: Optional[callable] = None, + **amg_kwargs, +): + """Cache the per-slice AMG grid-prediction state for a whole volume with one shared segmenter. + + Builds the AMG segmenter and reads the embedding signature once, then caches each slice via + `_cache_amg_slice`, reusing the shared 3d embeddings. This avoids rebuilding the segmenter and + re-reading the signature on every slice, which a per-slice `cache_autoseg_state` loop would do. + + Args: + model: The SAM2 model. + get_slice: Callable mapping a slice index to the 2d image for that slice. The image is only + used if `image_embeddings` is None; with embeddings the slice is taken from them. + n_slices: The number of slices in the volume. + image_embeddings: The precomputed 3d (video-style) embeddings for the volume. + save_path: The embedding save path used to store the per-slice state. + model_type: The SAM2 model type, recorded in the cached state. + is_tiled: Whether to use the tiled segmenter. By default inferred from the embeddings. + tile_shape: The tile shape for the tiled segmenter. + halo: The tile overlap for the tiled segmenter. + verbose: Whether to report progress. By default, set to 'True'. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + amg_kwargs: Additional keyword arguments for the AMG segmenter. + + Returns: + The AMG segmenter, with the last slice's state set. + """ + from .v2.instance_segmentation import get_amg_segmenter + + if is_tiled is None: + is_tiled = image_embeddings is not None and image_embeddings.get("input_size") is None + + segmenter = get_amg_segmenter(model, is_tiled=is_tiled, model_type=model_type, **amg_kwargs) + signature = _embedding_signature(save_path) + init_kwargs = {"tile_shape": tile_shape, "halo": halo} if is_tiled else {} + + def init_slice(i): + segmenter.initialize(get_slice(i), image_embeddings=image_embeddings, i=i, verbose=False, **init_kwargs) + + _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update) + pbar_init(n_slices, "Precompute automatic segmentation state") + for i in range(n_slices): + _cache_amg_slice(segmenter, save_path, i, init_slice, embedding_signature=signature) + pbar_update(1) + pbar_close() + return segmenter + + +def _cache_ais_state_v2( + decoder: torch.nn.Module, + raw: np.ndarray, + image_embeddings: Optional[dict], + save_path: Optional[Union[str, os.PathLike]], + ndim: int, + model_type: Optional[str] = None, + i: Optional[int] = None, + state_index: Optional[int] = None, + is_tiled: Optional[bool] = None, tile_shape: Optional[Tuple[int, int]] = None, halo: Optional[Tuple[int, int]] = None, - precompute_amg_state: bool = False, - decoder: Optional["nn.Module"] = None, + device: Optional[str] = None, + z_block: Optional[int] = None, + z_halo: Optional[int] = None, + verbose: bool = True, + pbar_init: Optional[callable] = None, + pbar_update: Optional[callable] = None, ): - os.makedirs(output_path, exist_ok=True) - idx = 0 - for file_path in tqdm(input_files, total=len(input_files), desc="Precompute state for files"): - - if isinstance(file_path, np.ndarray): - out_path = os.path.join(output_path, f"embedding_{idx:05}.tif") - else: - out_path = os.path.join(output_path, os.path.basename(file_path)) - - _precompute_state_for_file( - predictor, file_path, out_path, - key=key, ndim=ndim, tile_shape=tile_shape, halo=halo, - precompute_amg_state=precompute_amg_state, decoder=decoder, - verbose=False, + """Compute and cache, or load, the SAM2 decoder-based (AIS) automatic-segmentation state. + + The SAM2 counterpart of `cache_is_state`, using the UniSAM2 decoder. The state (the foreground + and directed-distance predictions) is stored in the embedding Zarr under + 'autoseg_state/ais', using the key 'state' (whole image / volume) or 'state-{i}' + (a slice). It is independent of the post-processing parameters, so it is always reusable. Pass + 'save_path=None' to skip caching. + + Args: + decoder: The UniSAM2 model, loaded via `micro_sam.v2.automatic_segmentation.get_unisam2_model`. + raw: The image data. + image_embeddings: The (optionally precomputed) image embeddings. When given only the decoder + is run on them (no encoder pass). + save_path: The embedding save path used to store the state, or None to skip caching. + ndim: The number of spatial dimensions (2 or 3). + model_type: The SAM2 model type. Recorded in the cached state so it is not reused with a + different decoder. + i: The slice index passed to the segmenter's `initialize`. + state_index: The index used for the on-disk state (defaults to `i`). + is_tiled: Whether to use the tiled segmenter. By default inferred from the embeddings. + tile_shape: The tile shape for the tiled segmenter. + halo: The tile overlap for the tiled segmenter. + device: The device to run inference on. + z_block: Number of slices to decode per z block for volumes. + z_halo: Number of overlapping slices between z blocks for volumes. + verbose: Whether to run verbose. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + + Returns: + The AIS segmenter with the (cached or freshly computed) state set. + """ + from .v2.automatic_segmentation import get_unisam2_segmentation_generator + + if is_tiled is None: + is_tiled = image_embeddings is not None and image_embeddings.get("input_size") is None + + segmenter = get_unisam2_segmentation_generator(decoder, is_tiled=is_tiled, device=device) + + key_index = i if state_index is None else state_index + key, signature = None, None + if save_path is not None: + key = _autoseg_state_key(key_index) + signature = _embedding_signature(save_path) + state = _load_ais_state_v2(save_path, key) + matches = state is not None and _ais_state_matches(state, model_type) + matches = matches and _signature_matches(state.get("embedding_signature"), signature) + if matches: + if verbose: + print("Load instance segmentation state from", save_path, ":", key) + segmenter.set_state(state) + return segmenter + + if verbose: + print("Precomputing the state for automatic instance segmentation.") + + segmenter.initialize( + raw, ndim, image_embeddings=image_embeddings, i=i, tile_shape=tile_shape, halo=halo, + z_block=z_block, z_halo=z_halo, pbar_init=pbar_init, pbar_update=pbar_update, + ) + if key is not None: + _save_ais_state_v2(segmenter, save_path, key, model_type, embedding_signature=signature) + return segmenter + + +def _resolve_unisam2_decoder(model_type, checkpoint_path, device): + """Return a UniSAM2 decoder for the SAM2 model if one is available, else None (fall back to AMG). + + Mirrors `micro_sam.v2.automatic_segmentation.get_segmenter`: a decoder from a custom + `checkpoint_path`, or the registered decoder of a finetuned model (e.g. 'hvit_t_cells'). Any + failure (e.g. an interactive-only checkpoint without a decoder) returns None. + """ + from .v2.util import FINETUNED_MODELS, has_registered_decoder, _download_finetuned_sam2_model + from .v2.automatic_segmentation import get_unisam2_model + + encoder = model_type[:6] + if checkpoint_path is not None: + decoder_source = checkpoint_path + elif model_type in FINETUNED_MODELS and has_registered_decoder(model_type): + _, _, decoder_source = _download_finetuned_sam2_model(model_type) + else: + return None + try: + return get_unisam2_model(decoder_source, device=util.get_device(device), encoder=encoder) + except Exception as e: + print(f"Could not load a UniSAM2 decoder from '{decoder_source}': {e}") + return None + + +def _cache_autoseg_state_for_file( + predictor, decoder, model_type, image_data, embeddings, save_path, ndim, verbose, +): + """Cache the SAM2 automatic-segmentation state for one file: AIS if a decoder is given, else AMG.""" + if decoder is not None: # AIS segments the whole image / volume in one pass. + device = next(decoder.parameters()).device + cache_autoseg_state( + "ais", decoder, image_data, embeddings, save_path, ndim=ndim, model_type=model_type, + device=device, verbose=verbose, + ) + elif ndim == 2: # AMG on a single 2d image. + model = getattr(predictor, "model", predictor) + cache_autoseg_state("amg", model, image_data, embeddings, save_path, model_type=model_type, verbose=verbose) + else: # AMG on a volume: cache the per-slice grid state, reusing the 3d embeddings and one segmenter. + model = getattr(predictor, "model", predictor) + _cache_amg_volume_state( + model, lambda i: image_data[i], image_data.shape[0], embeddings, save_path, + model_type=model_type, verbose=verbose, ) - idx += 1 def precompute_state( @@ -231,8 +662,10 @@ def precompute_state( checkpoint_path: Optional[Union[os.PathLike, str]] = None, key: Optional[str] = None, ndim: Optional[int] = None, + precompute_autoseg_state: bool = False, + prefer_decoder: bool = True, ) -> None: - """Precompute and cache the image embeddings for the input image(s). + """Precompute and cache the image embeddings (and, optionally, the automatic-segmentation state). The embeddings are saved in the same zarr format the annotators use, so the output can be loaded directly by the `micro_sam.annotator` CLI and the napari GUI by passing the same path as the @@ -253,14 +686,32 @@ def precompute_state( key: The key to the input file. This is needed for container files (e.g. hdf5 or zarr) or to load several images as 3d volume. Provide a glob pattern, e.g. "*.tif", for this case. ndim: The dimensionality of the data. By default, computed from the input data. + precompute_autoseg_state: Whether to also precompute the automatic-segmentation state in the + embeddings (a longer start-up in exchange for a faster first automatic segmentation). + Supported for SAM2 ('hvit_*') models. + prefer_decoder: Whether to use the decoder-based state (AIS) when the SAM2 model has a decoder, + instead of grid-based mask generation (AMG). """ # Imported lazily to avoid a circular import ('_state' imports from this module). from micro_sam.sam_annotator._state import _get_sam_model + is_sam2 = model_type.startswith("hvit") + if precompute_autoseg_state and not is_sam2: + raise ValueError( + "Precomputing the automatic-segmentation state via 'precompute_state' is only supported for " + "SAM2 ('hvit_*') models. For SAM1 use the annotator command with " + "'--precompute_autoseg_state'." + ) + # Dispatch to the embedding function for the model family (SAM1 / SAM2 / VFM). All share the same # interface, so the per-family predictor from '_get_sam_model' feeds directly into it. compute_embeddings = util.get_embedding_function(model_type) + # Resolve the UniSAM2 decoder once (AIS); when none is available the state is cached with AMG. + decoder = None + if precompute_autoseg_state and prefer_decoder: + decoder = _resolve_unisam2_decoder(model_type, checkpoint_path, device=None) + # Determine the input files and matching output embedding paths. single = pattern is None if single: @@ -289,6 +740,11 @@ def precompute_state( current_ndim = file_ndim save_path = str(Path(out_path).with_suffix(".zarr")) - compute_embeddings( + embeddings = compute_embeddings( predictor=predictor, input_=image_data, save_path=save_path, ndim=file_ndim, verbose=single ) + + if precompute_autoseg_state: + _cache_autoseg_state_for_file( + predictor, decoder, model_type, image_data, embeddings, save_path, file_ndim, verbose=single, + ) diff --git a/micro_sam/sam_annotator/_batch_classification.py b/micro_sam/sam_annotator/_batch_classification.py index bcee0c7ce..acde27f2b 100644 --- a/micro_sam/sam_annotator/_batch_classification.py +++ b/micro_sam/sam_annotator/_batch_classification.py @@ -68,7 +68,7 @@ def _init_predictor(self, viewer, image, embedding_path, reuse): kwargs = {"predictor": state.predictor} if (reuse and state.predictor is not None) else {} state.initialize_predictor( image, model_type=self.model_type, save_path=embedding_path, halo=self.halo, - tile_shape=self.tile_shape, precompute_amg_state=False, ndim=self.ndim, + tile_shape=self.tile_shape, precompute_autoseg_state=False, ndim=self.ndim, checkpoint_path=self.checkpoint_path, device=self.device, skip_load=False, use_cli=True, **kwargs, ) state.image_shape = image.shape if image.ndim == self.ndim else image.shape[:-1] diff --git a/micro_sam/sam_annotator/_state.py b/micro_sam/sam_annotator/_state.py index 0cfc649fb..860203d1e 100644 --- a/micro_sam/sam_annotator/_state.py +++ b/micro_sam/sam_annotator/_state.py @@ -17,8 +17,10 @@ import micro_sam import micro_sam.util as util from micro_sam.v1.util import get_sam_model -from micro_sam.v1.instance_segmentation import AMGBase, get_decoder -from micro_sam.precompute_state import cache_amg_state, cache_is_state +from micro_sam.v1.instance_segmentation import AutoSegBase, get_decoder +from micro_sam.precompute_state import ( + cache_amg_state, cache_is_state, cache_autoseg_state, _cache_amg_volume_state +) from segment_anything import SamPredictor @@ -45,7 +47,7 @@ def _get_sam_model(model_type, ndim, device, checkpoint_path, decoder_path, use_ encoder = get_vfm_model(model_type, device=device, checkpoint_path=checkpoint_path) return encoder, {} - if model_type.startswith("h"): # i.e. SAM2 models. + if model_type.startswith("hvit"): # i.e. SAM2 models. from micro_sam.v2.util import get_sam2_image_predictor, get_sam2_model # 'device=None' lets 'get_sam2_model' auto-detect the best device (cuda > mps > cpu); @@ -98,11 +100,11 @@ class AnnotatorState(metaclass=Singleton): # Whether the one-time CPU info popup has been shown this session (not reset on recompute). cpu_info_shown: Optional[bool] = None - # amg: needs to be initialized for the automatic segmentation functionality. - # amg_state: for storing the instance segmentation state for the 3d segmentation tool. + # The segmenter and its cached state for automatic segmentation. The state contains grid masks + # for AMG or decoder predictions for AIS and is loaded on demand for SAM2. # decoder: for direct prediction of instance segmentation - amg: Optional[AMGBase] = None - amg_state: Optional[Dict] = None + automatic_segmenter: Optional[AutoSegBase] = None + autoseg_state: Optional[Dict] = None decoder: Optional[nn.Module] = None # current_track_id, lineage, committed_lineages: @@ -158,7 +160,7 @@ def initialize_predictor( decoder_path=None, tile_shape=None, halo=None, - precompute_amg_state=False, + precompute_autoseg_state=False, prefer_decoder=True, pbar_init=None, pbar_update=None, @@ -177,7 +179,7 @@ def initialize_predictor( decoder_path = None from micro_sam.models.vfm import is_vfm_model - self.is_sam2 = model_type.startswith("h") + self.is_sam2 = model_type.startswith("hvit") self.is_vfm = is_vfm_model(model_type) # Initialize the model if necessary. @@ -297,32 +299,74 @@ def initialize_predictor( else: self.data_signature = util._compute_data_signature(image_data) - # Precompute the amg state (if specified). - if precompute_amg_state: + # Precompute the automatic-segmentation state (if specified). Decoder present -> AIS + # (decoder predictions), otherwise -> AMG (grid masks); this mirrors the v1 dispatch. + if precompute_autoseg_state: if save_path is None: - raise RuntimeError("Require a save path to precompute the amg state") + raise RuntimeError("Require a save path to precompute the automatic segmentation state") - cache_state = cache_amg_state if self.decoder is None else partial( - cache_is_state, decoder=self.decoder, skip_load=skip_load, - ) - - if ndim == 2: - self.amg = cache_state( - predictor=self.predictor, - raw=image_data, - image_embeddings=self.image_embeddings, - save_path=save_path + if self.is_sam2: + self._precompute_autoseg_state_sam2( + image_data, ndim, save_path, model_type, pbar_init=pbar_init, pbar_update=pbar_update, ) else: - n_slices = image_data.shape[0] if image_data.ndim == 3 else image_data.shape[1] - for i in tqdm(range(n_slices), desc="Precompute amg state"): - slice_ = np.s_[i] if image_data.ndim == 3 else np.s_[:, i] - cache_state( - predictor=self.predictor, - raw=image_data[slice_], - image_embeddings=self.image_embeddings, - save_path=save_path, i=i, verbose=False, - ) + self._precompute_autoseg_state_sam1(image_data, ndim, save_path, skip_load) + + def _precompute_autoseg_state_sam1(self, image_data, ndim, save_path, skip_load): + cache_state = cache_amg_state if self.decoder is None else partial( + cache_is_state, decoder=self.decoder, skip_load=skip_load, + ) + if ndim == 2: + self.automatic_segmenter = cache_state( + predictor=self.predictor, raw=image_data, + image_embeddings=self.image_embeddings, save_path=save_path, + ) + else: + n_slices = image_data.shape[0] if image_data.ndim == 3 else image_data.shape[1] + for i in tqdm(range(n_slices), desc="Precompute automatic segmentation state"): + slice_ = np.s_[i] if image_data.ndim == 3 else np.s_[:, i] + cache_state( + predictor=self.predictor, raw=image_data[slice_], + image_embeddings=self.image_embeddings, save_path=save_path, i=i, verbose=False, + ) + + def _precompute_autoseg_state_sam2( + self, image_data, ndim, save_path, model_type, pbar_init=None, pbar_update=None, + ): + model = getattr(self.predictor, "model", self.predictor) + resolved_model_type = getattr(self.predictor, "model_type", model_type) + + # The decoder / AMG functions drive the bar per tile / slice / z-block, but label it as if it + # were the actual run. Relabel it with a single clear description while keeping the real unit + # total, so the precompute phase (after the embeddings) reads meaningfully. + def relabel_pbar_init(total, _description): + pbar_init(total, "Precompute automatic segmentation state") + init_cb = relabel_pbar_init if pbar_init is not None else None + + if self.decoder is not None: # AIS: decoder over the whole image / volume (per tile / z-block). + device = next(self.decoder.parameters()).device + cache_autoseg_state( + "ais", self.decoder, image_data, self.image_embeddings, save_path, ndim=ndim, + model_type=resolved_model_type, device=device, pbar_init=init_cb, pbar_update=pbar_update, + ) + elif ndim == 2: # AMG on a single 2d image. + if pbar_init is not None: + pbar_init(1, "Precompute automatic segmentation state") + cache_autoseg_state( + "amg", model, image_data, self.image_embeddings, save_path, model_type=resolved_model_type, + ) + if pbar_update is not None: + pbar_update(1) + else: # AMG on a volume: cache the grid-prediction state per slice, reusing the 3d embeddings. + n_slices = image_data.shape[0] if image_data.ndim == 3 else image_data.shape[1] + + def get_slice(i): + return image_data[np.s_[i] if image_data.ndim == 3 else np.s_[:, i]] + + _cache_amg_volume_state( + model, get_slice, n_slices, self.image_embeddings, save_path, + model_type=resolved_model_type, pbar_init=pbar_init, pbar_update=pbar_update, + ) # Get the name of the image layer used to compute the embeddings. # If the 'image_name' attribute exists we can just use it. @@ -402,8 +446,8 @@ def reset_state(self): self.ndim = None self.image_name = None self.embedding_path = None - self.amg = None - self.amg_state = None + self.automatic_segmenter = None + self.autoseg_state = None self.decoder = None self.current_track_id = None self.lineage = None diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py index 30bef73a7..21998ed0a 100644 --- a/micro_sam/sam_annotator/_tooltips.py +++ b/micro_sam/sam_annotator/_tooltips.py @@ -2,6 +2,7 @@ tooltips = { "embedding": { + "cache_state": "Cache the automatic segmentation state to disk for faster (re)runs.", "custom_weights": "Select custom model weights. For example for a model you have finetuned", "device": "Select the computational device to use for processing.", "embeddings_save_path": "Select path to save or load the computed image embeddings.", diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index 0dad49f39..ab61c674b 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -1556,6 +1556,10 @@ class EmbeddingWidget(_WidgetBase): # A tiled 2D image only gets slow on the CPU once it has many tiles; warn from this many on. cpu_warn_tiles = 64 + # Whether to offer the 'cache automatic segmentation state' option (segmentation / tracking only, + # not the classification tools which have no automatic segmentation). + supports_state_caching = True + def __init__(self, parent=None, sam2_only=False, ndim_choice=False, is_timeseries=False): super().__init__(parent=parent) self.sam2_only = sam2_only @@ -1676,11 +1680,17 @@ def _update_model(self, state): 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) + # Load the AMG/AIS state cache. For SAM2 the state cache (grid masks or decoder + # predictions) is recorded via '_autoseg_state_descriptor' and the widget reads/writes it + # on demand; SAM1 preloads the per-slice 3d state as before. + if state.is_sam2: + state.autoseg_state = vutil._autoseg_state_descriptor( + state.embedding_path, "ais" if with_decoder else "amg", + ) + elif state.widgets["autosegment"].volumetric and with_decoder: + state.autoseg_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) + state.autoseg_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. @@ -1789,6 +1799,18 @@ def _create_settings_widget(self): ) setting_values.layout().addLayout(save_layout) + # Opt-in disk caching of the automatic-segmentation state (off by default). When on, the state + # is precomputed to disk (next to the embeddings) while the embeddings are computed and reused + # across runs / sessions. The automatic segmentation widget reads this flag from here. Not + # offered by the classification tools (no automatic segmentation). + self.cache_state = False + if self.supports_state_caching: + self.cache_state_checkbox = self._add_boolean_param( + "cache_state", self.cache_state, title="cache automatic segmentation state", + tooltip=get_tooltip("embedding", "cache_state"), + ) + setting_values.layout().addWidget(self.cache_state_checkbox) + # Create UI for the custom weights. self.custom_weights = None self.custom_weights_param, weights_layout = self._add_path_param( @@ -2145,6 +2167,17 @@ def __call__(self, skip_validate=False): # Warn CPU users once per session that processing can be slow. self._maybe_warn_cpu(ndim, tile_shape, state.image_shape) + # Eager caching of the automatic-segmentation state: if the 'cache automatic segmentation state' + # option (in these embedding settings) is on, precompute the state to disk while computing the + # embeddings. It persists wherever the embeddings live on disk - the given save path, or the + # ephemeral zarr the eager setup creates for SAM2 volumes / tiled images (removed on reset). + # Plain in-memory 2d has no disk location, so it needs a save path to persist. + is_sam2 = self.model_type.startswith("hvit") + embeddings_on_disk = save_path is not None or (is_sam2 and (ndim == 3 or tile_shape is not None)) + precompute_autoseg_state = self.cache_state and embeddings_on_disk + if self.cache_state and not embeddings_on_disk: + show_info("Set an embeddings save path to cache the automatic segmentation state.") + # Set up progress bar and signals for using it within a threadworker. pbar, pbar_signals = _create_pbar_for_threadworker() @@ -2157,6 +2190,10 @@ def compute_image_embedding(): def pbar_init(total, description): if self.is_timeseries: # A timeseries goes through the 3D compute path; relabel it. description = description.replace("3D", "Timeseries") + # Reset the counter to 0 so each phase starts fresh: the embeddings, then (when caching + # is on) the automatic-segmentation state precompute reuse the same bar, and without a + # reset the second phase inherits the first's completed count and sits stuck at full. + pbar_signals.pbar_reset.emit() pbar_signals.pbar_total.emit(total) pbar_signals.pbar_description.emit(description) QtWidgets.QApplication.processEvents() @@ -2175,6 +2212,7 @@ def pbar_update(update): tile_shape=tile_shape, halo=halo, prefer_decoder=True, + precompute_autoseg_state=precompute_autoseg_state, pbar_init=pbar_init, pbar_update=pbar_update, ) @@ -2200,6 +2238,9 @@ class ClassificationEmbeddingWidget(EmbeddingWidget): the SAM1 families and the VFM (DINO / UNI) families (`_advanced_family_suffixes` and `_dino_families`). """ + # The classification tools have no automatic segmentation, so the state-caching option is hidden. + supports_state_caching = False + size_order = ["tiny", "small", "base", "large", "huge", "giant"] # The classifiers only run the image encoder and a lightweight classifier on top, so they stay @@ -3169,11 +3210,11 @@ def update_segmentation(results): # -# Messy amg state handling, would be good to refactor this properly at some point. -def _handle_amg_state(state, i, pbar_init, pbar_update): - if state.amg is None: +# Messy automatic-segmentation state handling, would be good to refactor this properly at some point. +def _handle_autoseg_state(state, i, pbar_init, pbar_update): + if state.automatic_segmenter is None: is_tiled = state.image_embeddings["input_size"] is None - state.amg = instance_segmentation.get_instance_segmentation_generator( + state.automatic_segmenter = instance_segmentation.get_instance_segmentation_generator( state.predictor, is_tiled=is_tiled, decoder=state.decoder ) @@ -3181,15 +3222,15 @@ def _handle_amg_state(state, i, pbar_init, pbar_update): # Further optimization: refactor parts of this so that we can also use it in the automatic 3d segmentation fucnction # For 3D we store the amg state in a dict and check if it is computed already. - if state.amg_state is not None: + if state.autoseg_state is not None: assert i is not None - if i in state.amg_state: - amg_state_i = state.amg_state[i] - state.amg.set_state(amg_state_i) + if i in state.autoseg_state: + segmentation_state_i = state.autoseg_state[i] + state.automatic_segmenter.set_state(segmentation_state_i) else: dummy_image = np.zeros(shape[-2:], dtype="uint8") - state.amg.initialize( + state.automatic_segmenter.initialize( dummy_image, image_embeddings=state.image_embeddings, i=i, @@ -3197,43 +3238,43 @@ def _handle_amg_state(state, i, pbar_init, pbar_update): pbar_init=pbar_init, pbar_update=pbar_update, ) - amg_state_i = state.amg.get_state() - state.amg_state[i] = amg_state_i + segmentation_state_i = state.automatic_segmenter.get_state() + state.autoseg_state[i] = segmentation_state_i - cache_folder = state.amg_state.get("cache_folder", None) + cache_folder = state.autoseg_state.get("cache_folder", None) if cache_folder is not None: cache_path = os.path.join(cache_folder, f"state-{i}.pkl") with open(cache_path, "wb") as f: - pickle.dump(amg_state_i, f) + pickle.dump(segmentation_state_i, f) - cache_path = state.amg_state.get("cache_path", None) + cache_path = state.autoseg_state.get("cache_path", None) if cache_path is not None: save_key = f"state-{i}" with h5py.File(cache_path, "a") as f: g = f.create_group(save_key) g.create_dataset( "foreground", - data=amg_state_i["foreground"], + data=segmentation_state_i["foreground"], compression="gzip", ) g.create_dataset( "boundary_distances", - data=amg_state_i["boundary_distances"], + data=segmentation_state_i["boundary_distances"], compression="gzip", ) g.create_dataset( "center_distances", - data=amg_state_i["center_distances"], + data=segmentation_state_i["center_distances"], compression="gzip", ) # Otherwise (2d segmentation) we just check if the amg is initialized or not. - elif not state.amg.is_initialized: + elif not state.automatic_segmenter.is_initialized: assert i is None # We don't need to pass the actual image data here, since the embeddings are passed. # (The image data is only used by the amg to compute image embeddings, so not needed here.) dummy_image = np.zeros(shape, dtype="uint8") - state.amg.initialize( + state.automatic_segmenter.initialize( dummy_image, image_embeddings=state.image_embeddings, verbose=pbar_init is not None, @@ -3246,8 +3287,8 @@ def _instance_segmentation_impl( min_object_size, i=None, pbar_init=None, pbar_update=None, **kwargs ): state = AnnotatorState() - _handle_amg_state(state, i, pbar_init, pbar_update) - seg = state.amg.generate(**kwargs) + _handle_autoseg_state(state, i, pbar_init, pbar_update) + seg = state.automatic_segmenter.generate(**kwargs) assert isinstance(seg, np.ndarray) return seg @@ -3814,9 +3855,15 @@ def _allow_segment_3d(self): predictor = state.predictor if str(predictor.device) == "cpu" or str(predictor.device) == "mps": n_slices = self._viewer.layers["auto_segmentation"].data.shape[0] - embeddings_are_precomputed = (state.amg_state is not None) and ( - len(state.amg_state) > n_slices - ) + if state.is_sam2: + from micro_sam.precompute_state import _has_autoseg_state + embeddings_are_precomputed = _has_autoseg_state( + state.embedding_path, "amg", state_count=n_slices, + ) + else: + embeddings_are_precomputed = (state.autoseg_state is not None) and ( + len(state.autoseg_state) > n_slices + ) if not embeddings_are_precomputed: return False return True @@ -3928,6 +3975,10 @@ class AutoSegmentWidget(_WidgetBase): offered as a fallback when no decoder is available. A mode dropdown sits next to the 'Apply to Volume' switch and the post-processing parameters refresh on mode change. + Disk-backed caching of the state is opted into via the 'cache automatic segmentation state' + checkbox in the embedding settings (read here through the embedding widget); when off, the state + is kept in memory only. + Args: viewer: The napari viewer. with_decoder: Whether the loaded model has a UniSAM2 decoder for automatic segmentation. @@ -4213,10 +4264,17 @@ def _z_tiling(self, n_slices): z_halo = self.halo_z if z_block < n_slices else 0 return z_block, z_halo + def _state_save_path(self, state): + # The state cache is opted into via the embedding settings' 'cache automatic segmentation state' + # checkbox; when on it persists next to the embeddings, else in-memory only ('_segmenter' cache). + embed_widget = state.widgets.get("embeddings") + return state.embedding_path if getattr(embed_widget, "cache_state", False) else None + def _run_unisam2(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None): - from micro_sam.v2.automatic_segmentation import get_unisam2_segmentation_generator + from micro_sam.precompute_state import cache_autoseg_state device = next(state.decoder.parameters()).device + save_path = self._state_save_path(state) # All decoder auto-seg cases reuse the precomputed embeddings and run the decoder on them (no # encoder re-run). The tiling is taken from the embeddings (tiled embeddings have a top-level @@ -4243,14 +4301,21 @@ def _run_unisam2(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None else: image_embeddings, is_tiled = emb3d, True - # The cache avoids re-running the model when only the post-processing parameters change. - cache_key = (state.data_signature, "unisam2", ndim, z, tile_shape, halo, z_block, z_halo, - image_embeddings is not None) + # The in-memory cache avoids re-running the model when only the post-processing parameters + # change; 'cache_autoseg_state' additionally persists the decoder predictions in the + # embedding Zarr so a later run / session reuses them. The whole volume is + # cached under one key ('state'); a single segmented slice under 'state-{z}'. + cache_key = ( + state.data_signature, "unisam2", ndim, z, tile_shape, halo, z_block, z_halo, + image_embeddings is not None, save_path, + ) if self._segmenter is None or self._segmenter_key != cache_key: - self._segmenter = get_unisam2_segmentation_generator(state.decoder, is_tiled=is_tiled, device=device) - self._segmenter.initialize( - run_raw, ndim, image_embeddings=image_embeddings, tile_shape=tile_shape, halo=halo, i=z, - pbar_init=pbar_init, pbar_update=pbar_update, z_block=z_block, z_halo=z_halo, + self._segmenter = cache_autoseg_state( + "ais", state.decoder, run_raw, image_embeddings, save_path, ndim=ndim, + model_type=getattr(state.predictor, "model_type", None), + i=z, state_index=(None if ndim == 3 else z), is_tiled=is_tiled, + tile_shape=tile_shape, halo=halo, device=device, z_block=z_block, z_halo=z_halo, + pbar_init=pbar_init, pbar_update=pbar_update, verbose=False, ) self._segmenter_key = cache_key @@ -4258,27 +4323,28 @@ def _run_unisam2(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None def _run_amg(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None): from micro_sam.v2.instance_segmentation import get_amg_segmenter, automatic_3d_segmentation + from micro_sam.precompute_state import cache_autoseg_state # The SAM2 model: 'state.predictor' is the image predictor (2d) wrapping the model, or the # video predictor itself (3d); both can drive the grid-based mask generator. model = getattr(state.predictor, "model", state.predictor) model_type = getattr(state.predictor, "model_type", None) + save_path = self._state_save_path(state) generate_kwargs = dict(min_object_size=self.min_object_size, with_background=True) - - def _build(is_tiled): - return get_amg_segmenter( - model, is_tiled=is_tiled, model_type=model_type, - points_per_side=self.points_per_side, pred_iou_thresh=self.pred_iou_thresh, - stability_score_thresh=self.stability_score_thresh, - ) + amg_params = dict( + points_per_side=self.points_per_side, pred_iou_thresh=self.pred_iou_thresh, + stability_score_thresh=self.stability_score_thresh, + ) if ndim == 3: # Segment slice-by-slice and stitch across z. Tiling is in-plane (None if off). tile_shape, halo = self._get_tiling() - # Reuse the precomputed 3d embeddings per slice (tiled or not) so AMG does not re-encode. + segmenter = get_amg_segmenter(model, is_tiled=tile_shape is not None, model_type=model_type, **amg_params) + # Reuse the precomputed 3d embeddings per slice (tiled or not) so AMG does not re-encode, + # and cache each slice's grid-prediction state in the embedding Zarr. return automatic_3d_segmentation( - run_raw, _build(tile_shape is not None), tile_shape=tile_shape, halo=halo, - image_embeddings=state.image_embeddings, + run_raw, segmenter, tile_shape=tile_shape, halo=halo, + image_embeddings=state.image_embeddings, state_save_path=save_path, pbar_init=pbar_init, pbar_update=pbar_update, **generate_kwargs, ) @@ -4292,20 +4358,26 @@ def _build(is_tiled): tile_shape, halo, image_embeddings = None, None, state.image_embeddings is_tiled = image_embeddings["input_size"] is None - # The cache lets changing the post-processing parameters re-run only the cheap 'generate'. - cache_key = (state.data_signature, "amg", z, tile_shape, halo, image_embeddings is not None, - self.points_per_side, self.pred_iou_thresh, self.stability_score_thresh) + # The in-memory cache lets changing the post-processing parameters re-run only the cheap + # 'generate'; the on-disk cache (via 'cache_autoseg_state') persists the state across sessions. + cache_key = ( + state.data_signature, "amg", z, tile_shape, halo, image_embeddings is not None, + self.points_per_side, self.pred_iou_thresh, self.stability_score_thresh, save_path, + ) if self._segmenter is None or self._segmenter_key != cache_key: - self._segmenter = _build(is_tiled) if is_tiled: # The tiled segmenter reports per-tile progress. - self._segmenter.initialize( - run_raw, tile_shape=tile_shape, halo=halo, image_embeddings=image_embeddings, - pbar_init=pbar_init, pbar_update=pbar_update, + self._segmenter = cache_autoseg_state( + "amg", model, run_raw, image_embeddings, save_path, model_type=model_type, + state_index=z, is_tiled=True, tile_shape=tile_shape, halo=halo, + pbar_init=pbar_init, pbar_update=pbar_update, verbose=False, **amg_params, ) else: # A single 2d image is one step. if pbar_init is not None: pbar_init(1, "Automatic segmentation") - self._segmenter.initialize(run_raw, tile_shape=tile_shape, halo=halo, image_embeddings=image_embeddings) + self._segmenter = cache_autoseg_state( + "amg", model, run_raw, image_embeddings, save_path, model_type=model_type, + state_index=z, is_tiled=False, verbose=False, **amg_params, + ) if pbar_update is not None: pbar_update(1) self._segmenter_key = cache_key diff --git a/micro_sam/sam_annotator/annotator.py b/micro_sam/sam_annotator/annotator.py index cac0bdc2b..c473d00a4 100644 --- a/micro_sam/sam_annotator/annotator.py +++ b/micro_sam/sam_annotator/annotator.py @@ -296,9 +296,9 @@ def _update_image(self, segmentation_result=None): if self._ndim == 3: state = AnnotatorState() if state.decoder is not None: - state.amg_state = _load_is_state(state.embedding_path) + state.autoseg_state = _load_is_state(state.embedding_path) else: - state.amg_state = _load_amg_state(state.embedding_path) + state.autoseg_state = _load_amg_state(state.embedding_path) def annotator( @@ -312,7 +312,7 @@ def annotator( halo: Optional[Tuple[int, int]] = None, return_viewer: bool = False, viewer: Optional["napari.viewer.Viewer"] = None, - precompute_amg_state: bool = False, + precompute_autoseg_state: bool = False, checkpoint_path: Optional[str] = None, decoder_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, @@ -337,7 +337,8 @@ def annotator( By default, does not return the napari viewer. viewer: The viewer to which the Segment Anything functionality should be added. This enables using a pre-initialized viewer. - precompute_amg_state: Whether to precompute the state for automatic mask generation. + precompute_autoseg_state: Whether to precompute the automatic segmentation state (AMG masks, or + decoder predictions if the model has a decoder). Requires an embedding path. This will take more time when precomputing embeddings, but will then make automatic mask generation much faster. By default, set to 'False'. checkpoint_path: Path to a custom checkpoint from which to load the SAM model. @@ -371,7 +372,7 @@ def annotator( save_path=embedding_path, halo=halo, tile_shape=tile_shape, - precompute_amg_state=precompute_amg_state, + precompute_autoseg_state=precompute_autoseg_state, ndim=ndim, checkpoint_path=checkpoint_path, decoder_path=decoder_path, diff --git a/micro_sam/sam_annotator/annotator_tracking.py b/micro_sam/sam_annotator/annotator_tracking.py index 2f78b5575..5eabe1c0d 100644 --- a/micro_sam/sam_annotator/annotator_tracking.py +++ b/micro_sam/sam_annotator/annotator_tracking.py @@ -382,9 +382,9 @@ def _update_image(self): self._init_track_state() state = AnnotatorState() if self._with_decoder: - state.amg_state = vutil._load_is_state(state.embedding_path) + state.autoseg_state = vutil._load_is_state(state.embedding_path) else: - state.amg_state = vutil._load_amg_state(state.embedding_path) + state.autoseg_state = vutil._load_amg_state(state.embedding_path) def annotator_tracking( @@ -396,7 +396,7 @@ def annotator_tracking( halo: Optional[Tuple[int, int]] = None, return_viewer: bool = False, viewer: Optional["napari.viewer.Viewer"] = None, - precompute_amg_state: bool = False, + precompute_autoseg_state: bool = False, checkpoint_path: Optional[str] = None, decoder_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, @@ -415,7 +415,8 @@ def annotator_tracking( By default, does not return the napari viewer. viewer: The viewer to which the Segment Anything functionality should be added. This enables using a pre-initialized viewer. - precompute_amg_state: Whether to precompute the state for automatic mask generation. + precompute_autoseg_state: Whether to precompute the automatic segmentation state (AMG masks, or + decoder predictions if the model has a decoder). Requires an embedding path. This will take more time when precomputing embeddings, but will then make automatic mask generation much faster. By default, set to 'False'. checkpoint_path: Path to a custom checkpoint from which to load the SAM model. @@ -426,7 +427,6 @@ def annotator_tracking( Returns: The napari viewer, only returned if `return_viewer=True`. """ - _validate_tracking_model_type(model_type) # Initialize the predictor state. @@ -442,7 +442,7 @@ def annotator_tracking( checkpoint_path=checkpoint_path, decoder_path=decoder_path, device=device, - precompute_amg_state=precompute_amg_state, + precompute_autoseg_state=precompute_autoseg_state, use_cli=True, ) state.image_shape = image.shape[:-1] if image.ndim == 4 else image.shape @@ -485,7 +485,8 @@ class TrackingBatchTask(BatchAnnotatorTask): def __init__( self, *, model_type, embedding_path=None, tile_shape=None, halo=None, - checkpoint_path=None, decoder_path=None, device=None, precompute_amg_state=False, + checkpoint_path=None, decoder_path=None, device=None, + precompute_autoseg_state=False, ): _validate_tracking_model_type(model_type) self.model_type = model_type @@ -495,7 +496,7 @@ def __init__( self.checkpoint_path = checkpoint_path self.decoder_path = decoder_path self.device = device - self.precompute_amg_state = precompute_amg_state + self.precompute_autoseg_state = precompute_autoseg_state def result_filename(self, entry, index): if self.have_inputs_as_arrays: @@ -526,7 +527,8 @@ def _init_predictor(self, image, embedding_path, reuse): image, model_type=self.model_type, save_path=embedding_path, halo=self.halo, tile_shape=self.tile_shape, ndim=3, checkpoint_path=self.checkpoint_path, decoder_path=self.decoder_path, device=self.device, - precompute_amg_state=self.precompute_amg_state, use_cli=True, **kwargs, + precompute_autoseg_state=self.precompute_autoseg_state, + use_cli=True, **kwargs, ) state.image_shape = image.shape[:-1] if image.ndim == 4 else image.shape @@ -573,7 +575,7 @@ def batch_tracking_annotator( checkpoint_path: Optional[str] = None, decoder_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, - precompute_amg_state: bool = False, + precompute_autoseg_state: bool = False, viewer: Optional["napari.viewer.Viewer"] = None, return_viewer: bool = False, skip_done: bool = True, @@ -592,7 +594,8 @@ def batch_tracking_annotator( decoder_path: Path to a custom decoder checkpoint from which to load the `micro-sam` decoder. device: The computational device to use for the SAM model. By default, automatically chooses the best available device. - precompute_amg_state: Whether to precompute the state for automatic mask generation. + precompute_autoseg_state: Whether to precompute the automatic segmentation state (AMG masks, or + decoder predictions if the model has a decoder). Requires an embedding path. viewer: The viewer to which the functionality should be added. return_viewer: Whether to return the napari viewer instead of starting the event loop. skip_done: Whether to skip videos whose tracking result already exists in `output_folder`. @@ -604,7 +607,7 @@ def batch_tracking_annotator( task = TrackingBatchTask( model_type=model_type, embedding_path=embedding_path, tile_shape=tile_shape, halo=halo, checkpoint_path=checkpoint_path, decoder_path=decoder_path, device=device, - precompute_amg_state=precompute_amg_state, + precompute_autoseg_state=precompute_autoseg_state, ) return run_batch( images, output_folder, task, have_inputs_as_arrays=have_inputs_as_arrays, diff --git a/micro_sam/sam_annotator/batch_annotator.py b/micro_sam/sam_annotator/batch_annotator.py index 282f02ec8..8f90b5a8b 100644 --- a/micro_sam/sam_annotator/batch_annotator.py +++ b/micro_sam/sam_annotator/batch_annotator.py @@ -45,14 +45,15 @@ class SegmentationBatchTask(BatchAnnotatorTask): def __init__( self, *, ndim, model_type, embedding_path, tile_shape, halo, - precompute_amg_state, checkpoint_path, device, prefer_decoder, initial_segmentations=None, + precompute_autoseg_state, checkpoint_path, device, prefer_decoder, + initial_segmentations=None, ): self.ndim = ndim self.model_type = model_type self.embedding_path = embedding_path self.tile_shape = tile_shape self.halo = halo - self.precompute_amg_state = precompute_amg_state + self.precompute_autoseg_state = precompute_autoseg_state self.checkpoint_path = checkpoint_path self.device = device self.prefer_decoder = prefer_decoder @@ -98,7 +99,9 @@ def _init_predictor(self, viewer, image, embedding_path): kwargs = dict(prefer_decoder=self.prefer_decoder) state.initialize_predictor( image, model_type=self.model_type, save_path=embedding_path, halo=self.halo, tile_shape=self.tile_shape, - ndim=self.ndim, precompute_amg_state=self.precompute_amg_state, checkpoint_path=self.checkpoint_path, + ndim=self.ndim, + precompute_autoseg_state=self.precompute_autoseg_state, + checkpoint_path=self.checkpoint_path, device=self.device, skip_load=False, use_cli=True, **kwargs, ) # Capture the loaded model so subsequent items reuse it instead of reloading. @@ -131,8 +134,8 @@ def advance(self, viewer, annotator, entry, image, embedding_path, index): viewer.layers["committed_objects"].data = np.zeros_like(viewer.layers["committed_objects"].data) segmentation_result = self._resolve_initial_result(entry, index) viewer.layers["image"].data = image - if state.amg is not None: - state.amg.clear_state() + if state.automatic_segmenter is not None: + state.automatic_segmenter.clear_state() self._init_predictor(viewer, image, embedding_path) annotator._update_image(segmentation_result=segmentation_result) @@ -156,7 +159,7 @@ def batch_annotator( halo: Optional[Tuple[int, int]] = None, viewer: Optional["napari.viewer.Viewer"] = None, return_viewer: bool = False, - precompute_amg_state: bool = False, + precompute_autoseg_state: bool = False, checkpoint_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, prefer_decoder: bool = True, @@ -181,7 +184,8 @@ def batch_annotator( This enables using a pre-initialized viewer. return_viewer: Whether to return the napari viewer to further modify it before starting the tool. By default, does not return the napari viewer. - precompute_amg_state: Whether to precompute the state for automatic mask generation. + precompute_autoseg_state: Whether to precompute the automatic segmentation state (AMG masks, or + decoder predictions if the model has a decoder). Requires an embedding path. This will take more time when precomputing embeddings, but will then make automatic mask generation much faster. By default, set to 'False'. checkpoint_path: Path to a custom checkpoint from which to load the SAM model. @@ -210,7 +214,8 @@ def batch_annotator( task = SegmentationBatchTask( ndim=ndim, model_type=model_type, embedding_path=embedding_path, - tile_shape=tile_shape, halo=halo, precompute_amg_state=precompute_amg_state, + tile_shape=tile_shape, halo=halo, + precompute_autoseg_state=precompute_autoseg_state, checkpoint_path=checkpoint_path, device=device, prefer_decoder=prefer_decoder, initial_segmentations=initial_segmentations, ) diff --git a/micro_sam/sam_annotator/object_classifier.py b/micro_sam/sam_annotator/object_classifier.py index 576bce164..a32d0f376 100644 --- a/micro_sam/sam_annotator/object_classifier.py +++ b/micro_sam/sam_annotator/object_classifier.py @@ -245,7 +245,7 @@ def object_classifier( state.initialize_predictor( image, model_type=model_type, save_path=embedding_path, - halo=halo, tile_shape=tile_shape, precompute_amg_state=False, + halo=halo, tile_shape=tile_shape, precompute_autoseg_state=False, ndim=ndim, checkpoint_path=checkpoint_path, device=device, skip_load=False, use_cli=True, ) diff --git a/micro_sam/sam_annotator/pixel_classifier.py b/micro_sam/sam_annotator/pixel_classifier.py index adf987458..af1f257cf 100644 --- a/micro_sam/sam_annotator/pixel_classifier.py +++ b/micro_sam/sam_annotator/pixel_classifier.py @@ -104,7 +104,7 @@ def pixel_classifier( state.initialize_predictor( image, model_type=model_type, save_path=embedding_path, - halo=halo, tile_shape=tile_shape, precompute_amg_state=False, + halo=halo, tile_shape=tile_shape, precompute_autoseg_state=False, ndim=ndim, checkpoint_path=checkpoint_path, device=device, skip_load=False, use_cli=True, ) diff --git a/micro_sam/sam_annotator/util.py b/micro_sam/sam_annotator/util.py index 826c0de51..cc307c3d2 100644 --- a/micro_sam/sam_annotator/util.py +++ b/micro_sam/sam_annotator/util.py @@ -1182,7 +1182,7 @@ def _sync_embedding_widget(widget, model_type, save_path, checkpoint_path, devic # Update the index for model size, eg. 'base', 'tiny', etc. size_map = {"t": "tiny", "s": "small", "b": "base", "l": "large", "h": "huge"} - size_idx = 5 if model_type.startswith("h") else 4 + size_idx = 5 if model_type.startswith("hvit") else 4 model_size = size_map.get(model_type[size_idx]) if model_size is not None: @@ -1284,3 +1284,15 @@ def _load_is_state(embedding_path): is_state[i] = state return is_state + + +def _autoseg_state_descriptor(embedding_path, mode): + """Descriptor of the SAM2 automatic-segmentation state cache in the embedding Zarr. + + Returns the embedding path and mode ('amg' or 'ais'); the state itself is loaded on demand by + `micro_sam.precompute_state.cache_autoseg_state`. The SAM2 automatic segmentation widget + reads/writes the cache directly, so this only records where it lives. + """ + if embedding_path is None or not os.path.exists(embedding_path): + return {"embedding_path": None, "mode": mode} + return {"embedding_path": embedding_path, "mode": mode} diff --git a/micro_sam/v1/automatic_segmentation.py b/micro_sam/v1/automatic_segmentation.py index 3e57961c5..4d9243c6d 100644 --- a/micro_sam/v1/automatic_segmentation.py +++ b/micro_sam/v1/automatic_segmentation.py @@ -14,7 +14,7 @@ from .. import util from .util import get_sam_model, precompute_image_embeddings, get_model_names from .instance_segmentation import ( - get_instance_segmentation_generator, get_decoder, AMGBase, + get_instance_segmentation_generator, get_decoder, AutoSegBase, AutomaticMaskGenerator, TiledAutomaticMaskGenerator, AutomaticPromptGenerator, TiledAutomaticPromptGenerator, InstanceSegmentationWithDecoder, TiledInstanceSegmentationWithDecoder, @@ -32,7 +32,7 @@ def get_predictor_and_segmenter( predictor=None, state=None, **kwargs, -) -> Tuple[util.SamPredictor, Union[AMGBase, InstanceSegmentationWithDecoder]]: +) -> Tuple[util.SamPredictor, AutoSegBase]: f"""Get the Segment Anything model and class for automatic instance segmentation. Args: @@ -88,7 +88,7 @@ def _add_suffix_to_output_path(output_path: Union[str, os.PathLike], suffix: str def automatic_tracking( predictor: util.SamPredictor, - segmenter: Union[AMGBase, InstanceSegmentationWithDecoder], + segmenter: AutoSegBase, input_path: Union[Union[os.PathLike, str], np.ndarray], output_path: Optional[Union[os.PathLike, str]] = None, embedding_path: Optional[Union[os.PathLike, str]] = None, @@ -171,7 +171,7 @@ def automatic_tracking( def automatic_instance_segmentation( predictor: util.SamPredictor, - segmenter: Union[AMGBase, InstanceSegmentationWithDecoder], + segmenter: AutoSegBase, input_path: Union[Union[os.PathLike, str], np.ndarray], output_path: Optional[Union[os.PathLike, str]] = None, embedding_path: Optional[Union[os.PathLike, str]] = None, diff --git a/micro_sam/v1/evaluation/instance_segmentation.py b/micro_sam/v1/evaluation/instance_segmentation.py index 90a2e5381..270bfae6b 100644 --- a/micro_sam/v1/evaluation/instance_segmentation.py +++ b/micro_sam/v1/evaluation/instance_segmentation.py @@ -16,7 +16,7 @@ from elf.evaluation import mean_segmentation_accuracy, matching from ..util import precompute_image_embeddings -from ..instance_segmentation import AMGBase, InstanceSegmentationWithDecoder +from ..instance_segmentation import AutoSegBase def _get_range_of_search_values(input_vals, step): @@ -169,7 +169,7 @@ def default_grid_search_values_apg( def _grid_search_iteration( - segmenter: Union[AMGBase, InstanceSegmentationWithDecoder], + segmenter: AutoSegBase, gs_combinations: List[Dict], gt: np.ndarray, image_name: str, @@ -216,7 +216,7 @@ def _load_image(path, key, roi): def run_instance_segmentation_grid_search( - segmenter: Union[AMGBase, InstanceSegmentationWithDecoder], + segmenter: AutoSegBase, grid_search_values: Dict[str, List], image_paths: List[Union[str, os.PathLike]], gt_paths: List[Union[str, os.PathLike]], @@ -322,7 +322,7 @@ def run_instance_segmentation_grid_search( def run_instance_segmentation_inference( - segmenter: Union[AMGBase, InstanceSegmentationWithDecoder], + segmenter: AutoSegBase, image_paths: List[Union[str, os.PathLike]], embedding_dir: Optional[Union[str, os.PathLike]], prediction_dir: Union[str, os.PathLike], @@ -429,7 +429,7 @@ def save_grid_search_best_params(best_kwargs, best_msa, grid_search_result_dir=N def run_instance_segmentation_grid_search_and_inference( - segmenter: Union[AMGBase, InstanceSegmentationWithDecoder], + segmenter: AutoSegBase, grid_search_values: Dict[str, List], val_image_paths: List[Union[str, os.PathLike]], val_gt_paths: List[Union[str, os.PathLike]], diff --git a/micro_sam/v1/instance_segmentation.py b/micro_sam/v1/instance_segmentation.py index 8707925fb..54d945249 100644 --- a/micro_sam/v1/instance_segmentation.py +++ b/micro_sam/v1/instance_segmentation.py @@ -7,7 +7,7 @@ import shutil import tempfile import warnings -from abc import ABC +from abc import ABC, abstractmethod from contextlib import contextmanager from copy import deepcopy from collections import OrderedDict @@ -63,9 +63,43 @@ def __getitem__(self, index): # -class AMGBase(ABC): - """Base class for the automatic mask generators. +class AutoSegBase(ABC): + """Common interface for automatic-segmentation generators. + + Unifies the grid-based mask generators (`AMGBase` and its subclasses) and the decoder-based + instance segmentation (`InstanceSegmentationWithDecoder` and its subclasses): both are + initialized on an image, produce an instance segmentation via `generate`, and cache / restore + their expensive intermediate state via `get_state` / `set_state`. """ + + @property + def is_initialized(self) -> bool: + """Whether `initialize` has been run and the state is available.""" + return self._is_initialized + + @abstractmethod + def initialize(self, *args, **kwargs) -> None: + """Compute and store the (expensive) state needed by `generate`.""" + + @abstractmethod + def generate(self, *args, **kwargs): + """Produce the instance segmentation from the initialized state.""" + + @abstractmethod + def get_state(self) -> Dict[str, Any]: + """Return the cached state so it can be serialized and later restored.""" + + @abstractmethod + def set_state(self, state: Dict[str, Any]) -> None: + """Restore a state produced by `get_state`.""" + + @abstractmethod + def clear_state(self) -> None: + """Clear the cached state.""" + + +class AMGBase(AutoSegBase): + """Base class for the grid-based (AMG) automatic mask generators.""" def __init__(self): # the state that has to be computed by the 'initialize' method of the child classes self._is_initialized = False @@ -73,12 +107,6 @@ def __init__(self): self._crop_boxes = None self._original_size = None - @property - def is_initialized(self): - """Whether the mask generator has already been initialized. - """ - return self._is_initialized - @property def crop_list(self): """The list of mask data after initialization. @@ -951,10 +979,11 @@ def _apply_smoothing(foreground, foreground_smoothing, tile_shape, n_threads): return foreground -class InstanceSegmentationWithDecoder: +class InstanceSegmentationWithDecoder(AutoSegBase): """Generates an instance segmentation without prompts, using a decoder. - Implements the same interface as `AutomaticMaskGenerator`. + A concrete `AutoSegBase` (like `AutomaticMaskGenerator`), but predicts the segmentation with a + decoder instead of grid prompts. Use this class as follows: ```python @@ -978,12 +1007,6 @@ def __init__(self, predictor: SamPredictor, decoder: torch.nn.Module) -> None: self._is_initialized = False - @property - def is_initialized(self): - """Whether the mask generator has already been initialized. - """ - return self._is_initialized - @torch.no_grad() def initialize( self, @@ -1637,7 +1660,7 @@ def get_instance_segmentation_generator( decoder: Optional[torch.nn.Module] = None, segmentation_mode: Optional[Literal["amg", "ais", "apg"]] = None, **kwargs, -) -> Union[AMGBase, InstanceSegmentationWithDecoder]: +) -> AutoSegBase: f"""Get the automatic mask generator. Args: diff --git a/micro_sam/v1/multi_dimensional_segmentation.py b/micro_sam/v1/multi_dimensional_segmentation.py index adc56250f..eccad58cc 100644 --- a/micro_sam/v1/multi_dimensional_segmentation.py +++ b/micro_sam/v1/multi_dimensional_segmentation.py @@ -44,7 +44,7 @@ from .. import util from .util import precompute_image_embeddings from .prompt_based_segmentation import segment_from_mask -from .instance_segmentation import AMGBase +from .instance_segmentation import AutoSegBase PROJECTION_MODES = ("box", "mask", "points", "points_and_mask", "single_point") @@ -424,7 +424,7 @@ def _segment_slices( def automatic_3d_segmentation( volume: np.ndarray, predictor: SamPredictor, - segmentor: AMGBase, + segmentor: AutoSegBase, embedding_path: Optional[Union[str, os.PathLike]] = None, with_background: bool = True, gap_closing: Optional[int] = None, @@ -686,7 +686,7 @@ def track_across_frames( def automatic_tracking_implementation( timeseries: np.ndarray, predictor: SamPredictor, - segmentor: AMGBase, + segmentor: AutoSegBase, embedding_path: Optional[Union[str, os.PathLike]] = None, gap_closing: Optional[int] = None, min_time_extent: Optional[int] = None, diff --git a/micro_sam/v2/automatic_segmentation.py b/micro_sam/v2/automatic_segmentation.py index 4dd15da1f..af5eafa29 100644 --- a/micro_sam/v2/automatic_segmentation.py +++ b/micro_sam/v2/automatic_segmentation.py @@ -791,6 +791,22 @@ def generate(self, mode: str = "sparse", **kwargs) -> np.ndarray: raise RuntimeError("The segmenter has not been initialized. Call 'initialize' first.") return segment_from_predictions(self._prediction, mode=mode, **kwargs) + def get_state(self) -> dict: + """Return the cached decoder predictions so they can be serialized and later restored. + + The state holds the (4, *spatial) foreground + directed-distance predictions. Restore it + with `set_state` to skip the expensive decoder inference in `initialize`. It is independent + of the post-processing parameters (those are applied in `generate`), so it is always reusable. + """ + if not self._is_initialized: + raise RuntimeError("Cannot get the state before the segmenter has been initialized.") + return {"prediction": self._prediction} + + def set_state(self, state: dict) -> None: + """Restore the state produced by `get_state`, marking the segmenter initialized.""" + self._prediction = np.asarray(state["prediction"]) + self._is_initialized = True + class TiledUniSAM2InstanceSegmentation(UniSAM2InstanceSegmentation): """Generates a tiled instance segmentation with the UniSAM2 model. diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index 19a7725ad..6df634779 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -159,6 +159,14 @@ def __init__( predictor = self._mask_generator.predictor predictor.model_type = model_type or getattr(model, "model_type", None) or "hvit" predictor.model_name = model_type or getattr(model, "model_name", None) or predictor.model_type + # The parameters that are baked into the predicted masks during 'initialize'. They are stored + # in the cached state so a reused state can be validated against the requested parameters. + self._amg_params = { + "points_per_side": points_per_side, + "pred_iou_thresh": pred_iou_thresh, + "stability_score_thresh": stability_score_thresh, + "model_type": predictor.model_type, + } self._masks = None self._original_size = None self._is_initialized = False @@ -276,6 +284,30 @@ def generate( with_background=with_background, ) + def get_state(self) -> Dict[str, Any]: + """Return the cached mask-generation state so it can be serialized and later restored. + + The state holds the predicted masks (as compact RLE dicts), the image size and the + parameters the masks were generated with (used to validate a reused state). Restore it + with `set_state` to skip the expensive grid prediction in `initialize`. + """ + if not self._is_initialized: + raise RuntimeError("Cannot get the state before the segmenter has been initialized.") + return { + "masks": [dict(mask) for mask in self._masks], + "original_size": self._original_size, + "params": dict(self._amg_params), + } + + def set_state(self, state: Dict[str, Any]) -> None: + """Restore the state produced by `get_state`, marking the segmenter initialized. + + The masks are re-wrapped in `_LazyRLEMask` so `generate` decodes them lazily as before. + """ + self._masks = [_LazyRLEMask(mask) for mask in state["masks"]] + self._original_size = tuple(int(s) for s in state["original_size"]) + self._is_initialized = True + class TiledAutomaticMaskGenerationSegmenter(AutomaticMaskGenerationSegmenter): """Generates a tiled instance segmentation for the SAM2 model using grid-based prompting (AMG). @@ -344,6 +376,7 @@ def initialize( tile_shape = tuple(int(s) for s in feats.attrs["tile_shape"]) halo = tuple(int(s) for s in feats.attrs["halo"]) self._original_size = tuple(int(s) for s in feats.attrs["shape"]) + self._tile_shape = tile_shape self._tiling = Blocking([0, 0], list(self._original_size), list(tile_shape)) self._halo = tuple(halo) @@ -376,6 +409,7 @@ def _initialize_slice_from_3d_embeddings(self, image_embeddings, i): tile_shape = tuple(int(s) for s in feats.attrs["tile_shape"]) halo = tuple(int(s) for s in feats.attrs["halo"]) self._original_size = full_shape[1:] # in-plane (Y, X); z is not tiled + self._tile_shape = tile_shape self._tiling = Blocking([0, 0], list(self._original_size), list(tile_shape)) self._halo = halo @@ -450,6 +484,27 @@ def generate( return segmentation + def get_state(self) -> Dict[str, Any]: + """Return the cached per-tile mask state, plus the tiling needed to restore it.""" + if not self._is_initialized: + raise RuntimeError("Cannot get the state before the segmenter has been initialized.") + return { + "masks": [[dict(mask) for mask in tile_masks] for tile_masks in self._masks], + "original_size": self._original_size, + "tile_shape": self._tile_shape, + "halo": self._halo, + "params": dict(self._amg_params), + } + + def set_state(self, state: Dict[str, Any]) -> None: + """Restore the per-tile masks and rebuild the tiling from `get_state`.""" + self._masks = [[_LazyRLEMask(mask) for mask in tile_masks] for tile_masks in state["masks"]] + self._original_size = tuple(int(s) for s in state["original_size"]) + self._tile_shape = tuple(int(s) for s in state["tile_shape"]) + self._halo = tuple(int(s) for s in state["halo"]) + self._tiling = Blocking([0, 0], list(self._original_size), list(self._tile_shape)) + self._is_initialized = True + def get_amg_segmenter( model: torch.nn.Module, is_tiled: bool = False, **kwargs @@ -479,6 +534,7 @@ def automatic_3d_segmentation( tile_shape: Optional[Tuple[int, int]] = None, halo: Optional[Tuple[int, int]] = None, image_embeddings: Optional[dict] = None, + state_save_path: Optional[str] = None, verbose: bool = True, pbar_init: Optional[callable] = None, pbar_update: Optional[callable] = None, @@ -504,6 +560,9 @@ def automatic_3d_segmentation( halo: The overlap between the tiles, (y, x). By default 'None'. image_embeddings: Optional precomputed 3d (video-style) embeddings for the volume. When given (and not tiled), each slice's AMG reuses the precomputed features instead of re-encoding. + state_save_path: Optional path to the embedding Zarr in which to cache the per-slice + grid-prediction state. When set, a slice reuses its cached state instead of re-running + the grid prediction; freshly computed slices are written back. verbose: Verbosity flag. By default 'True'. pbar_init: Callback to initialize an external progress bar, called with the number of slices. pbar_update: Callback to update an external progress bar, called once per segmented slice. @@ -526,13 +585,23 @@ def automatic_3d_segmentation( _, pbar_init, pbar_update, pbar_close = handle_pbar(verbose, pbar_init, pbar_update) pbar_init(volume.shape[0], "Automatic segmentation (slices)") - segmentation = np.zeros(volume.shape, dtype="uint32") - offset = 0 - for i in range(volume.shape[0]): + def init_slice(i): if reuse_embeddings: segmenter.initialize(volume[i], image_embeddings=image_embeddings, i=i, verbose=False, **init_kwargs) else: segmenter.initialize(volume[i], verbose=False, **init_kwargs) + + if state_save_path is not None: + from micro_sam.precompute_state import _cache_amg_slice, _embedding_signature + state_signature = _embedding_signature(state_save_path) + + segmentation = np.zeros(volume.shape, dtype="uint32") + offset = 0 + for i in range(volume.shape[0]): + if state_save_path is not None: + _cache_amg_slice(segmenter, state_save_path, i, init_slice, embedding_signature=state_signature) + else: + init_slice(i) seg = segmenter.generate(**kwargs) # Offset the per-slice instance ids so that they are unique across the whole volume. diff --git a/test/test_auto_state.py b/test/test_auto_state.py new file mode 100644 index 000000000..78a797163 --- /dev/null +++ b/test/test_auto_state.py @@ -0,0 +1,365 @@ +"""Tests for the SAM2 automatic-segmentation state caching ('autoseg_state'). + +These are model-free: the segmenters are constructed directly (bypassing the model download and +encoder pass) and populated with hand-built state, so the tests exercise `get_state`/`set_state`, +the serialization helpers and the staleness guards without loading a SAM2 checkpoint. +""" + +import numpy as np +import pytest +import torch +from bioimage_cpp.utils import Blocking +from sam2.utils.amg import mask_to_rle_pytorch + +from micro_sam.precompute_state import ( + AUTOSEG_STATE_ATTRIBUTE, + AUTOSEG_STATE_GROUP, + _ais_state_matches, + _autoseg_state_key, + _has_autoseg_state, + _load_ais_state_v2, + _load_amg_state_v2, + _load_autoseg_state, + _save_ais_state_v2, + _save_amg_state_v2, + _save_autoseg_state, + _signature_matches, + cache_autoseg_state, +) +from micro_sam.util import _open_embeddings, _create_dataset_with_data +from micro_sam.v2.automatic_segmentation import UniSAM2InstanceSegmentation +from micro_sam.v2.instance_segmentation import ( + AutomaticMaskGenerationSegmenter, + TiledAutomaticMaskGenerationSegmenter, + _LazyRLEMask, +) + +DEFAULT_AMG_PARAMS = { + "points_per_side": 32, "pred_iou_thresh": 0.8, "stability_score_thresh": 0.9, "model_type": "hvit_t", +} + + +def _rle_mask(binary): + """A single AMG mask dict with an uncompressed RLE 'segmentation', as SAM2's AMG produces.""" + rle = mask_to_rle_pytorch(torch.from_numpy(binary[None]).to(torch.bool))[0] + return {"segmentation": rle, "area": int(binary.sum())} + + +def _make_amg_segmenter(masks, original_size, params=None): + segmenter = object.__new__(AutomaticMaskGenerationSegmenter) + segmenter._masks = [_LazyRLEMask(mask) for mask in masks] + segmenter._original_size = original_size + segmenter._amg_params = dict(params or DEFAULT_AMG_PARAMS) + segmenter._is_initialized = True + return segmenter + + +def _make_tiled_amg_segmenter(masks_per_tile, original_size, tile_shape, halo, params=None): + segmenter = object.__new__(TiledAutomaticMaskGenerationSegmenter) + segmenter._masks = [[_LazyRLEMask(mask) for mask in tile] for tile in masks_per_tile] + segmenter._original_size = original_size + segmenter._tile_shape = tile_shape + segmenter._halo = halo + segmenter._amg_params = dict(params or DEFAULT_AMG_PARAMS) + segmenter._tiling = Blocking([0, 0], list(original_size), list(tile_shape)) + segmenter._is_initialized = True + return segmenter + + +def test_autoseg_state_key(): + assert _autoseg_state_key(None) == "state" + assert _autoseg_state_key(3) == "state-3" + + +def test_amg_get_set_state_roundtrip(): + m1 = np.zeros((16, 16), bool) + m1[2:6, 2:6] = True + m2 = np.zeros((16, 16), bool) + m2[9:13, 9:13] = True + segmenter = _make_amg_segmenter([_rle_mask(m1), _rle_mask(m2)], (16, 16)) + + state = segmenter.get_state() + # The masks are stored as compact RLE dicts (not decoded arrays) and the params are recorded. + assert isinstance(state["masks"][0], dict) and isinstance(state["masks"][0]["segmentation"], dict) + assert state["params"] == DEFAULT_AMG_PARAMS + assert state["original_size"] == (16, 16) + + restored = object.__new__(AutomaticMaskGenerationSegmenter) + restored.set_state(state) + assert restored._is_initialized + assert isinstance(restored._masks[0], _LazyRLEMask) + + out_original = segmenter.generate(min_object_size=0, with_background=False) + out_restored = restored.generate(min_object_size=0, with_background=False) + assert np.array_equal(out_original, out_restored) + assert out_restored.max() == 2 # both objects survive the round-trip + + +def test_amg_empty_state_roundtrip(): + segmenter = _make_amg_segmenter([], (12, 12)) + restored = object.__new__(AutomaticMaskGenerationSegmenter) + restored.set_state(segmenter.get_state()) + out = restored.generate() + assert out.shape == (12, 12) and out.max() == 0 + + +def test_tiled_amg_get_set_state_roundtrip(): + # Two tiles side by side: full (16, 32), tile (16, 16), no halo. + tile1 = np.zeros((16, 16), bool) + tile1[2:6, 2:6] = True + tile2 = np.zeros((16, 16), bool) + tile2[8:12, 8:12] = True + segmenter = _make_tiled_amg_segmenter( + [[_rle_mask(tile1)], [_rle_mask(tile2)]], (16, 32), (16, 16), (0, 0), + ) + + state = segmenter.get_state() + assert state["tile_shape"] == (16, 16) and state["halo"] == (0, 0) + + restored = object.__new__(TiledAutomaticMaskGenerationSegmenter) + restored.set_state(state) + assert restored._tiling.number_of_blocks == segmenter._tiling.number_of_blocks + assert restored._tile_shape == (16, 16) and restored._halo == (0, 0) + + out_original = segmenter.generate(min_object_size=0, with_background=False) + out_restored = restored.generate(min_object_size=0, with_background=False) + assert np.array_equal(out_original, out_restored) + + +def test_ais_get_set_state_roundtrip(): + prediction = np.random.RandomState(0).rand(4, 16, 16).astype("float32") + segmenter = UniSAM2InstanceSegmentation(model=None) + segmenter._prediction = prediction + segmenter._is_initialized = True + + restored = UniSAM2InstanceSegmentation(model=None) + restored.set_state(segmenter.get_state()) + assert restored._is_initialized + assert np.array_equal(restored._prediction, prediction) + + +def test_amg_serialization_and_param_match(tmp_path): + m = np.zeros((16, 16), bool) + m[3:7, 3:7] = True + segmenter = _make_amg_segmenter([_rle_mask(m)], (16, 16)) + + save_path = str(tmp_path / "embeddings.zarr") + key = _autoseg_state_key(None) + _save_amg_state_v2(segmenter, save_path, key) + loaded = _load_amg_state_v2(save_path, key) + + assert loaded["params"] == segmenter._amg_params and len(loaded["masks"]) == 1 + # The param dict is what 'cache_autoseg_state' compares to decide whether to reuse the state. + assert loaded.get("params") == segmenter._amg_params + changed = _make_amg_segmenter([], (16, 16), params={**DEFAULT_AMG_PARAMS, "points_per_side": 64}) + assert loaded.get("params") != changed._amg_params + + +def test_ais_serialization_and_staleness_guard(tmp_path): + prediction = np.arange(4 * 4 * 4, dtype="float32").reshape(4, 4, 4) + segmenter = UniSAM2InstanceSegmentation(model=None) + segmenter._prediction = prediction + segmenter._is_initialized = True + + save_path = str(tmp_path / "embeddings.zarr") + key = _autoseg_state_key(None) + _save_ais_state_v2(segmenter, save_path, key, "hvit_t_cells") + loaded = _load_ais_state_v2(save_path, key) + + assert np.array_equal(loaded["prediction"], prediction) + assert loaded["model_type"] == "hvit_t_cells" + assert _ais_state_matches(loaded, "hvit_t_cells") # same model -> reuse + assert not _ais_state_matches(loaded, "hvit_t_other") # different model -> recompute + assert _ais_state_matches(loaded, None) # unknown request -> reuse + assert _ais_state_matches({"prediction": prediction}, "hvit_t_cells") # legacy (no signature) -> reuse + + +def test_signature_matches(): + assert _signature_matches("a", "a") + assert not _signature_matches("a", "b") # both known and different -> stale + assert _signature_matches(None, "a") # cached unknown (legacy) -> reuse + assert _signature_matches("a", None) # requested unknown -> reuse + assert _signature_matches(None, None) + + +def test_amg_embedding_signature_roundtrip(tmp_path): + m = np.zeros((16, 16), bool) + m[3:7, 3:7] = True + segmenter = _make_amg_segmenter([_rle_mask(m)], (16, 16)) + save_path = str(tmp_path / "embeddings.zarr") + key = _autoseg_state_key(None) + _save_amg_state_v2(segmenter, save_path, key, embedding_signature="sig-A") + loaded = _load_amg_state_v2(save_path, key) + assert loaded["embedding_signature"] == "sig-A" + assert _signature_matches(loaded.get("embedding_signature"), "sig-A") + assert not _signature_matches(loaded.get("embedding_signature"), "sig-B") # embeddings changed -> stale + + +def test_ais_embedding_signature_roundtrip(tmp_path): + segmenter = UniSAM2InstanceSegmentation(model=None) + segmenter._prediction = np.zeros((4, 4, 4), dtype="float32") + segmenter._is_initialized = True + save_path = str(tmp_path / "embeddings.zarr") + key = _autoseg_state_key(None) + _save_ais_state_v2(segmenter, save_path, key, "hvit_t_cells", embedding_signature="sig-A") + loaded = _load_ais_state_v2(save_path, key) + assert loaded["embedding_signature"] == "sig-A" + assert not _signature_matches(loaded.get("embedding_signature"), "sig-B") # embeddings changed -> stale + + +def test_states_share_embedding_zarr_and_record_metadata(tmp_path): + save_path = str(tmp_path / "embeddings.zarr") + amg = _make_amg_segmenter([], (16, 16)) + ais = UniSAM2InstanceSegmentation(model=None) + ais._prediction = np.zeros((4, 16, 16), dtype="float32") + ais._is_initialized = True + + _save_amg_state_v2(amg, save_path, "state-0") + _save_amg_state_v2(amg, save_path, "state-1") + _save_ais_state_v2(ais, save_path, "state", "hvit_t_cells") + + embeddings = _open_embeddings(save_path, mode="r") + assert set(embeddings.attrs[AUTOSEG_STATE_ATTRIBUTE]) == {"amg", "ais"} + state_root = embeddings[AUTOSEG_STATE_GROUP] + assert state_root["amg"].attrs["state_count"] == 2 + assert state_root["ais"].attrs["state_count"] == 1 + assert "state-0" in state_root["amg"] and "state" in state_root["ais"] + + # Cache completeness is checked from group metadata, without deserializing masks or predictions. + assert _has_autoseg_state(save_path, "amg", state_count=2) + assert not _has_autoseg_state(save_path, "amg", state_count=3) + assert _has_autoseg_state(save_path, "ais") + + +def test_state_stored_inside_embedding_zarr_representations(tmp_path): + """The state lives in the SAME zarr container as the embeddings (no sidecar file). AMG is saved as + a pickle bitstream (a 1-D uint8 dataset); AIS as an individual float32 array. This is the storage + contract behind the h5 -> zarr migration.""" + save_path = str(tmp_path / "embeddings.zarr") + + # Mimic a precomputed embedding container: a 'features' dataset plus an identifying attr. + embeddings = _open_embeddings(save_path, mode="a") + _create_dataset_with_data(embeddings, "features", data=np.zeros((1, 256, 8, 8), dtype="float32")) + embeddings.attrs["original_size"] = [[64, 64]] + + m = np.zeros((16, 16), bool) + m[3:7, 3:7] = True + _save_amg_state_v2(_make_amg_segmenter([_rle_mask(m)], (16, 16)), save_path, "state") + + ais = UniSAM2InstanceSegmentation(model=None) + ais._prediction = np.zeros((4, 16, 16), dtype="float32") + ais._is_initialized = True + _save_ais_state_v2(ais, save_path, "state", "hvit_t_cells") + + reopened = _open_embeddings(save_path, mode="r") + # One container holds the embeddings and both state modes (same filepath, no sidecar). + assert "features" in reopened + assert "original_size" in reopened.attrs # embeddings stay identifiable after writing the state + root = reopened[AUTOSEG_STATE_GROUP] + + amg_ds = root["amg"]["state"] + assert amg_ds.dtype == np.uint8 and len(amg_ds.shape) == 1 # a pickle bitstream, not decoded arrays + + ais_ds = root["ais"]["state"]["prediction"] + assert ais_ds.dtype == np.float32 and tuple(ais_ds.shape) == (4, 16, 16) # an individual array + assert ais_ds.chunks[0] == 1 # per-channel chunks so a read does not inflate one big chunk + + +def test_cache_autoseg_state_ais_is_on_demand(tmp_path, monkeypatch): + """The lazy contract via the mode dispatcher: a matching cached AIS state is loaded on demand (no + decoder rerun), and a stale one is recomputed. Mirrors how image embeddings are reused when cached.""" + prediction = np.arange(4 * 8 * 8, dtype="float32").reshape(4, 8, 8) + seed = UniSAM2InstanceSegmentation(model=None) + seed._prediction = prediction + seed._is_initialized = True + + save_path = str(tmp_path / "embeddings.zarr") + key = _autoseg_state_key(None) + _save_ais_state_v2(seed, save_path, key, "hvit_t_cells") + + # Track whether the (expensive) decoder pass runs; it must not run on a cache hit. + initialize_calls = [] + + def fake_initialize(self, *args, **kwargs): + initialize_calls.append(True) + self._prediction = np.zeros((4, 8, 8), dtype="float32") + self._is_initialized = True + + monkeypatch.setattr(UniSAM2InstanceSegmentation, "initialize", fake_initialize) + raw = np.zeros((8, 8), dtype="uint8") + + # Cache hit: the state is loaded from disk and no decoder is run ('decoder=None' would crash the + # real 'initialize', so a clean return with the cached prediction proves it was reused as-is). + segmenter = cache_autoseg_state( + "ais", None, raw, None, save_path, ndim=2, model_type="hvit_t_cells", verbose=False, + ) + assert initialize_calls == [] + assert np.array_equal(segmenter.get_state()["prediction"], prediction) + + # Stale cache (different model): recomputed on demand instead of silently reusing the wrong state. + cache_autoseg_state( + "ais", None, raw, None, save_path, ndim=2, model_type="hvit_t_other", verbose=False, + ) + assert initialize_calls == [True] + + +def test_amg_state_loads_per_slice_on_demand(tmp_path): + """Each slice's AMG state is a separate on-disk key, loaded one at a time. A volume's state is + streamed per slice (like the lazy per-slice embeddings), not materialized as one whole array.""" + m0 = np.zeros((16, 16), bool) + m0[1:5, 1:5] = True + m1 = np.zeros((16, 16), bool) + m1[10:14, 10:14] = True + _save_amg_state_v2(_make_amg_segmenter([_rle_mask(m0)], (16, 16)), str(tmp_path / "e.zarr"), "state-0") + _save_amg_state_v2(_make_amg_segmenter([_rle_mask(m1)], (16, 16)), str(tmp_path / "e.zarr"), "state-1") + save_path = str(tmp_path / "e.zarr") + + # Loading one slice returns only that slice's state, without touching the sibling slices. + state0 = _load_amg_state_v2(save_path, _autoseg_state_key(0)) + restored = object.__new__(AutomaticMaskGenerationSegmenter) + restored.set_state(state0) + out = restored.generate(min_object_size=0, with_background=False) + assert out[3, 3] == 1 and out[12, 12] == 0 # slice-0 object present, slice-1 object not loaded + + # A slice that was never cached loads as None (so it is computed on demand), not an error. + assert _load_amg_state_v2(save_path, _autoseg_state_key(5)) is None + + +def test_cache_autoseg_state_routes_by_mode(monkeypatch): + """The 'cache_autoseg_state' dispatcher forwards to the AMG / AIS implementation by mode, and + rejects unknown modes.""" + import micro_sam.precompute_state as ps + + calls = [] + monkeypatch.setattr(ps, "_cache_amg_state_v2", lambda *a, **k: calls.append(("amg", a, k)) or "amg-seg") + monkeypatch.setattr(ps, "_cache_ais_state_v2", lambda *a, **k: calls.append(("ais", a, k)) or "ais-seg") + + assert ps.cache_autoseg_state("amg", "MODEL", "RAW", None, "sp", points_per_side=8) == "amg-seg" + assert ps.cache_autoseg_state("ais", "DECODER", "RAW", None, "sp", ndim=3) == "ais-seg" + assert calls[0] == ("amg", ("MODEL", "RAW", None, "sp"), {"points_per_side": 8}) + assert calls[1] == ("ais", ("DECODER", "RAW", None, "sp"), {"ndim": 3}) + + with pytest.raises(ValueError, match="Invalid automatic-segmentation state mode"): + ps.cache_autoseg_state("bogus", None, None, None, None) + + +def test_save_load_autoseg_state_dispatch(tmp_path): + """The save/load dispatchers route to the AMG (pickle bitstream) and AIS (array) implementations.""" + save_path = str(tmp_path / "e.zarr") + m = np.zeros((16, 16), bool) + m[3:7, 3:7] = True + _save_autoseg_state("amg", _make_amg_segmenter([_rle_mask(m)], (16, 16)), save_path, "state") + + ais = UniSAM2InstanceSegmentation(model=None) + ais._prediction = np.zeros((4, 16, 16), dtype="float32") + ais._is_initialized = True + _save_autoseg_state("ais", ais, save_path, "state", model_type="hvit_t_cells") + + amg_loaded = _load_autoseg_state("amg", save_path, "state") + ais_loaded = _load_autoseg_state("ais", save_path, "state") + assert len(amg_loaded["masks"]) == 1 # AMG masks came back + assert tuple(ais_loaded["prediction"].shape) == (4, 16, 16) # AIS array came back + assert _load_autoseg_state("amg", save_path, "state-99") is None # missing key -> None + + with pytest.raises(ValueError, match="Invalid automatic-segmentation state mode"): + _load_autoseg_state("bogus", save_path, "state") diff --git a/test/test_instance_segmentation.py b/test/test_instance_segmentation.py index 75e91e09b..6d6190807 100644 --- a/test/test_instance_segmentation.py +++ b/test/test_instance_segmentation.py @@ -1,6 +1,7 @@ import unittest from copy import deepcopy +import pytest import micro_sam.util as util from micro_sam.v1.util import get_sam_model, precompute_image_embeddings import numpy as np @@ -150,5 +151,43 @@ def test_tiled_automatic_prompt_generator(self): ) +def test_autoseg_base_hierarchy_and_contract(): + """AutoSegBase is the common ABC for both the grid (AMG) and decoder-based families; every leaf + generator is a subclass exposing the shared contract, and the abstract bases cannot be built.""" + from micro_sam.v1.instance_segmentation import ( + AutoSegBase, AMGBase, AutomaticMaskGenerator, TiledAutomaticMaskGenerator, + InstanceSegmentationWithDecoder, TiledInstanceSegmentationWithDecoder, + AutomaticPromptGenerator, TiledAutomaticPromptGenerator, + ) + + # Both family bases descend from the common interface. + assert issubclass(AMGBase, AutoSegBase) + assert issubclass(InstanceSegmentationWithDecoder, AutoSegBase) + + # Every concrete generator is an AutoSegBase and exposes the full contract. + contract = ("is_initialized", "initialize", "generate", "get_state", "set_state", "clear_state") + leaves = [ + AutomaticMaskGenerator, TiledAutomaticMaskGenerator, + InstanceSegmentationWithDecoder, TiledInstanceSegmentationWithDecoder, + AutomaticPromptGenerator, TiledAutomaticPromptGenerator, + ] + import inspect + for cls in leaves: + assert issubclass(cls, AutoSegBase) + assert all(hasattr(cls, member) for member in contract) + assert not inspect.isabstract(cls) # every leaf implements the full contract -> instantiable + + # The abstract bases cannot be instantiated (initialize / generate stay abstract). + assert inspect.isabstract(AutoSegBase) and inspect.isabstract(AMGBase) + with pytest.raises(TypeError): + AutoSegBase() + with pytest.raises(TypeError): + AMGBase() + + # A concrete decoder segmenter constructs without model work and is an AutoSegBase. + seg = InstanceSegmentationWithDecoder(predictor=None, decoder=None) + assert isinstance(seg, AutoSegBase) and seg.is_initialized is False + + if __name__ == "__main__": unittest.main() diff --git a/test/test_sam_annotator/test_annotator.py b/test/test_sam_annotator/test_annotator.py index 5a4dab848..edd55803d 100644 --- a/test/test_sam_annotator/test_annotator.py +++ b/test/test_sam_annotator/test_annotator.py @@ -443,6 +443,132 @@ def test_z_tiling_hidden_for_2d_autoseg(self, make_napari_viewer_proxy): viewer.close() +class TestAutoSegVolumeDispatch: + """'Apply to volume' decides the run dimensionality of automatic segmentation on a 3d volume: + off -> only the current slice (2d); on -> the whole volume (3d), segmented slice by slice. + This is what makes the state caching happen per-slice, on demand.""" + + def _dispatch(self, monkeypatch, *, apply_to_volume, current_slice=2): + from types import SimpleNamespace + from micro_sam.sam_annotator import _widgets + from micro_sam.sam_annotator._widgets import AutoSegmentWidget + + volume = np.zeros((5, 16, 16), dtype="float32") + auto_layer = SimpleNamespace(data=np.zeros((5, 16, 16), dtype="uint32"), refresh=lambda: None) + viewer = SimpleNamespace( + layers={"image": SimpleNamespace(data=volume), "auto_segmentation": auto_layer}, + dims=SimpleNamespace(point=(current_slice, 0, 0)), + ) + + calls = {} + + def fake_run_amg(state, run_raw, ndim, z, pbar_init=None, pbar_update=None): + calls.update(run_raw_shape=tuple(run_raw.shape), ndim=ndim, z=z) + return np.zeros(run_raw.shape if ndim == 3 else run_raw.shape[-2:], dtype="uint32") + + # Duck-typed stand-in so we exercise the '__call__' dispatch without instantiating a QWidget. + widget = SimpleNamespace( + _viewer=viewer, mode="amg", volumetric=True, apply_to_volume=apply_to_volume, + _run_amg=fake_run_amg, + ) + + def fake_pbar(): + signal = lambda: SimpleNamespace(emit=lambda *a, **k: None) # noqa + signals = SimpleNamespace( + pbar_total=signal(), pbar_description=signal(), pbar_update=signal(), pbar_stop=signal(), + ) + return SimpleNamespace(), signals + + monkeypatch.setattr(_widgets, "AnnotatorState", lambda: SimpleNamespace(get_image_name=lambda v: "image")) + monkeypatch.setattr(_widgets, "_validate_layers", lambda *a, **k: False) + monkeypatch.setattr(_widgets, "_validate_embeddings", lambda *a, **k: False) + monkeypatch.setattr(_widgets, "_select_layer", lambda *a, **k: None) + monkeypatch.setattr(_widgets, "_create_pbar_for_threadworker", fake_pbar) + monkeypatch.setattr( + _widgets, "QtWidgets", + SimpleNamespace(QApplication=SimpleNamespace(processEvents=lambda *a, **k: None)), + ) + + AutoSegmentWidget.__call__(widget) + return calls + + def test_apply_to_volume_off_runs_current_slice_only(self, monkeypatch): + calls = self._dispatch(monkeypatch, apply_to_volume=False, current_slice=2) + assert calls["ndim"] == 2 # a single 2d slice, not the whole volume + assert calls["z"] == 2 # the currently viewed slice + assert calls["run_raw_shape"] == (16, 16) + + def test_apply_to_volume_on_runs_whole_volume(self, monkeypatch): + calls = self._dispatch(monkeypatch, apply_to_volume=True) + assert calls["ndim"] == 3 # the whole volume, segmented slice by slice + assert calls["z"] is None + assert calls["run_raw_shape"] == (5, 16, 16) + + +class TestAutoSegStatePersistence: + """By default (caching off) auto-seg persists no state, so the embedding zarr gets no + 'autoseg_state' group; it is written only when the user opts in. '_state_save_path' + is the gate: None means in-memory only, a path means persist into the embedding zarr.""" + + def _state_save_path(self, cache_state, embedding_path="/tmp/e.zarr", with_widget=True): + from types import SimpleNamespace + from micro_sam.sam_annotator._widgets import AutoSegmentWidget + + widgets = {"embeddings": SimpleNamespace(cache_state=cache_state)} if with_widget else {} + state = SimpleNamespace(widgets=widgets, embedding_path=embedding_path) + return AutoSegmentWidget._state_save_path(SimpleNamespace(), state) + + def test_no_persist_by_default(self): + assert self._state_save_path(cache_state=False) is None # default: in-memory only, no zarr group + + def test_persist_when_opted_in(self): + assert self._state_save_path(cache_state=True) == "/tmp/e.zarr" + + def test_no_persist_without_embedding_widget(self): + assert self._state_save_path(cache_state=True, with_widget=False) is None + + def test_enabling_persistence_invalidates_the_in_memory_cache(self, monkeypatch): + """Turning disk caching on after an in-memory run must route through the cache helper again.""" + from types import MethodType, SimpleNamespace + + import micro_sam.precompute_state as precompute_state + from micro_sam.sam_annotator._widgets import AutoSegmentWidget + + calls = [] + + class _Segmenter: + def generate(self, **kwargs): + return np.zeros((8, 8), dtype="uint32") + + def fake_cache_autoseg_state(*args, **kwargs): + calls.append(args[4]) # save_path + return _Segmenter() + + monkeypatch.setattr(precompute_state, "cache_autoseg_state", fake_cache_autoseg_state) + + embedding_widget = SimpleNamespace(cache_state=False) + state = SimpleNamespace( + predictor=SimpleNamespace(model=object(), model_type="hvit_t"), + image_embeddings={"input_size": (8, 8)}, + embedding_path="/tmp/embeddings.zarr", + data_signature="data", + widgets={"embeddings": embedding_widget}, + ) + widget = SimpleNamespace( + volumetric=False, min_object_size=0, points_per_side=32, + pred_iou_thresh=0.8, stability_score_thresh=0.9, + _segmenter=None, _segmenter_key=None, + ) + widget._state_save_path = MethodType(AutoSegmentWidget._state_save_path, widget) + + raw = np.zeros((8, 8), dtype="uint8") + AutoSegmentWidget._run_amg(widget, state, raw, ndim=2, z=None) + embedding_widget.cache_state = True + AutoSegmentWidget._run_amg(widget, state, raw, ndim=2, z=None) + + assert calls == [None, "/tmp/embeddings.zarr"] + + @pytest.mark.gui @pytest.mark.skipif(platform.system() in ("Windows",), reason="Gui test is not working on windows.") class TestAutoSegDefaultMode: diff --git a/test/test_sam_annotator/test_state.py b/test/test_sam_annotator/test_state.py index fc4c46b30..ce5193894 100644 --- a/test/test_sam_annotator/test_state.py +++ b/test/test_sam_annotator/test_state.py @@ -29,6 +29,18 @@ def test_state_for_tracking(self): self.assertTrue(state.initialized_for_tracking()) +def test_autoseg_names_have_no_legacy_aliases(): + """Only the canonical automatic segmenter and cached-state names are exposed.""" + from micro_sam.sam_annotator._state import AnnotatorState + + state = AnnotatorState() + assert hasattr(state, "automatic_segmenter") + assert hasattr(state, "autoseg_state") # canonical name + assert not hasattr(state, "amg") + assert not hasattr(state, "amg_state") # old aliases gone + assert not hasattr(state, "auto_state") + + def test_blank_model_paths_are_normalized(monkeypatch): """Blank API paths must fall back to the registered model instead of being loaded as files.""" import micro_sam.sam_annotator._state as state_module diff --git a/test/test_util.py b/test/test_util.py index 486f5b152..f11ffc51f 100644 --- a/test/test_util.py +++ b/test/test_util.py @@ -497,6 +497,7 @@ def _check_predictor_initialization_2d(self, predictor, embeddings): def test_precompute_image_embeddings_2d(self): from micro_sam.v2.normalization import RAW_NORMALIZATION from micro_sam.v2.util import precompute_image_embeddings + from micro_sam.v2.normalization import RAW_NORMALIZATION predictor = self._get_predictor(ndim=2) input_ = np.random.rand(512, 512).astype("float32")