From 7b936a3485353f66ab3c128a2caead34b4ebae0a Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Tue, 21 Jul 2026 18:47:11 +0200 Subject: [PATCH 01/11] Add batched and pipelined multi-GPU inference for SAM2 --- environment.yaml | 2 +- micro_sam/v2/automatic_segmentation.py | 33 +- micro_sam/v2/batched_inference.py | 1099 +++++++++++++++++++++ micro_sam/v2/instance_segmentation.py | 412 ++++---- micro_sam/v2/prompt_based_segmentation.py | 147 ++- micro_sam/v2/util.py | 364 +------ setup.cfg | 2 +- test/test_v2_batched_inference.py | 170 ++++ 8 files changed, 1659 insertions(+), 570 deletions(-) create mode 100644 micro_sam/v2/batched_inference.py create mode 100644 test/test_v2_batched_inference.py diff --git a/environment.yaml b/environment.yaml index 9e5553dc8..b5d52e370 100644 --- a/environment.yaml +++ b/environment.yaml @@ -25,7 +25,7 @@ dependencies: - pytorch >=2.5 - segment-anything - torchvision - - torch_em >=0.9 + - torch_em >=0.10.1 - tqdm - timm - transformers >=4.56 diff --git a/micro_sam/v2/automatic_segmentation.py b/micro_sam/v2/automatic_segmentation.py index 9c21a049b..d3947942f 100644 --- a/micro_sam/v2/automatic_segmentation.py +++ b/micro_sam/v2/automatic_segmentation.py @@ -15,7 +15,7 @@ import numpy as np import torch -from .util import DEFAULT_MODEL +from .util import DEFAULT_MODEL, Devices def get_predictor_and_segmenter( @@ -97,6 +97,10 @@ def automatic_instance_segmentation( mode: str = "sparse", device: Optional[Union[str, torch.device]] = None, verbose: bool = True, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, + num_write_workers: int = 1, **generate_kwargs, ) -> np.ndarray: """Run automatic instance segmentation for a single input and save the result. @@ -118,6 +122,10 @@ def automatic_instance_segmentation( mode: The AIS post-processing mode, 'sparse' (flow) or 'dense' (multicut). Ignored for AMG. device: The device to run inference on. verbose: Whether to print progress. + batch_size: Explicit tile or slice batch size, or None for automatic capacity probing. + devices: Inference device or devices. None uses all visible GPUs when the model is on CUDA. + num_prefetch_workers: Number of input reading and preprocessing threads. + num_write_workers: Number of output writing threads for full tiled inference. generate_kwargs: Additional post-processing parameters forwarded to the segmenter's `generate`. Returns: @@ -147,11 +155,28 @@ def automatic_instance_segmentation( else: emb_predictor = predictor image_embeddings = precompute_image_embeddings( - emb_predictor, raw, save_path=embedding_path, ndim=ndim, - tile_shape=tile_shape, halo=halo, verbose=verbose, lazy_loading=(ndim == 3), + emb_predictor, + raw, + save_path=embedding_path, + ndim=ndim, + tile_shape=tile_shape, + halo=halo, + verbose=verbose, + lazy_loading=(ndim == 3), + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, ) segmenter.initialize( - raw, ndim=ndim, image_embeddings=image_embeddings, tile_shape=tile_shape, halo=halo, + raw, + ndim=ndim, + image_embeddings=image_embeddings, + tile_shape=tile_shape, + halo=halo, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) segmentation = segmenter.generate(mode=mode, **generate_kwargs) diff --git a/micro_sam/v2/batched_inference.py b/micro_sam/v2/batched_inference.py new file mode 100644 index 000000000..d698dd9f3 --- /dev/null +++ b/micro_sam/v2/batched_inference.py @@ -0,0 +1,1099 @@ +"""Batched, pipelined, multi-GPU SAM2 inference: scheduling engine, encoder embeddings, and decoder passes.""" + +import gc +import os +import queue +import threading +import warnings +from collections import defaultdict +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union + +import numpy as np +import torch + +from micro_sam.util import _create_dataset_with_data, _create_dataset_without_data +from .normalization import to_image +from .util import Device, Devices + + +STOP = object() + + +class PipelineAborted(Exception): + """Raised inside workers when another pipeline worker has failed.""" + + +class AtomicCounter: + """Lock-guarded counter used to send completion sentinels exactly once.""" + + def __init__(self, value: int) -> None: + self.value = value + self.lock = threading.Lock() + + def decrement(self) -> int: + with self.lock: + self.value -= 1 + return self.value + + +@dataclass +class PipelineJob: + """A work item moving from the input loader to inference and output writing.""" + + spec: Any + data: Any + + +def _normalize_device(device: Union[str, torch.device]) -> torch.device: + device = torch.device(device) + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _model_device(model: torch.nn.Module) -> torch.device: + try: + return _normalize_device(next(model.parameters()).device) + except (AttributeError, StopIteration): + return torch.device("cpu") + + +def resolve_devices(model: torch.nn.Module, devices: Devices = None) -> List[torch.device]: + """Resolve inference devices, using every visible CUDA device by default. + + Automatic multi-GPU execution is enabled only when the supplied model already lives on CUDA. + This preserves an explicitly CPU- or MPS-loaded model. Pass a scalar device to force one device, + or a sequence to select an explicit set of devices. + """ + if devices is None: + device = _model_device(model) + if device.type == "cuda" and torch.cuda.device_count() > 1: + resolved = [torch.device("cuda", index) for index in range(torch.cuda.device_count())] + else: + resolved = [device] + elif isinstance(devices, (str, torch.device)): + resolved = [_normalize_device(devices)] + else: + resolved = [_normalize_device(device) for device in devices] + + if len(resolved) == 0: + raise ValueError("At least one inference device is required.") + if len(set(resolved)) != len(resolved): + raise ValueError(f"Inference devices must be unique, got {resolved}.") + if any(device.type == "cuda" for device in resolved) and not torch.cuda.is_available(): + raise RuntimeError("CUDA inference was requested, but PyTorch CUDA support is unavailable.") + return resolved + + +def prepare_models( + model: torch.nn.Module, devices: Sequence[torch.device], +) -> List[Tuple[torch.nn.Module, torch.device]]: + """Create one eval-mode model replica per device, reusing the original where possible.""" + source_device = _model_device(model) + models = [] + for device in devices: + replica = model if device == source_device else deepcopy(model).to(device) + if hasattr(replica, "eval"): + replica.eval() + models.append((replica, device)) + return models + + +def release_model_replicas(model_devices: List[Tuple[torch.nn.Module, torch.device]]) -> None: + """Drop the replicas built by :func:`prepare_models` and free their GPU memory. + + Clearing the list releases only the per-device copies; the caller's own model stays alive + through its other references. + """ + model_devices.clear() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def compute_auto_batch_sizes( + model_devices: Sequence[Tuple[torch.nn.Module, torch.device]], + n_jobs: int, + patch_shape: Tuple[int, ...], + in_channels: int, + prediction_function: Callable[[torch.nn.Module, torch.Tensor], Any], +) -> List[int]: + """Probe the largest safe batch size independently on every CUDA device. + + The probe delegates its exponential and binary OOM search to + :func:`torch_em.util.compute_max_batch_size`. CPU and MPS use a conservative batch size of one, + because the torch-em probe intentionally supports CUDA OOM detection only. + """ + if n_jobs < 1: + return [1] * len(model_devices) + + # Cap the probe at the per-device job count and at 256 (no run needs a larger batch). + jobs_per_device = (int(n_jobs) + len(model_devices) - 1) // len(model_devices) + upper_bound = min(256, jobs_per_device) + batch_sizes = [] + from torch_em.util import compute_max_batch_size + + for model, device in model_devices: + if device.type != "cuda": + batch_sizes.append(1) + continue + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", message="The batch size search reached the upper bound.*", category=UserWarning, + ) + batch_size = compute_max_batch_size( + model=model, + patch_shape=patch_shape, + in_channels=in_channels, + device=device, + dtype=torch.float32, + safety_factor=0.8, # leave GPU-memory headroom below the probed maximum + max_batch_size=upper_bound, + prediction_function=prediction_function, + ) + batch_sizes.append(batch_size) + return batch_sizes + + +def _safe_get(input_queue: queue.Queue, stop_event: threading.Event, timeout: float = 0.2) -> Any: + while not stop_event.is_set(): + try: + return input_queue.get(timeout=timeout) + except queue.Empty: + continue + raise PipelineAborted() + + +def _safe_put(output_queue: queue.Queue, item: Any, stop_event: threading.Event, timeout: float = 0.2) -> None: + while not stop_event.is_set(): + try: + output_queue.put(item, timeout=timeout) + return + except queue.Full: + continue + raise PipelineAborted() + + +def run_batched_pipeline( + jobs: Iterable[Any], + model_devices: Sequence[Tuple[torch.nn.Module, torch.device]], + batch_sizes: Sequence[int], + load_fn: Callable[[Any], Any], + predict_fn: Callable[[torch.nn.Module, List[Any], torch.device], List[Any]], + write_fn: Callable[[Any, Any], None], + num_prefetch_workers: int = 4, + update_progress: Optional[Callable[[int], None]] = None, +) -> None: + """Run load/preprocess, batched inference, and writing as a bounded threaded pipeline. + + There is one inference consumer per device and one writer. Keeping output writes on one thread + is safe for zarr, N5, HDF5, and in-memory outputs, while still overlapping I/O with GPU work. + """ + jobs = list(jobs) + if len(jobs) == 0: + return + if len(model_devices) != len(batch_sizes): + raise ValueError("Expected one batch size for every model/device pair.") + if any(int(batch_size) < 1 for batch_size in batch_sizes): + raise ValueError(f"Batch sizes must be positive, got {batch_sizes}.") + + num_prefetch_workers = max(1, min(int(num_prefetch_workers), len(jobs))) + batch_sizes = [int(batch_size) for batch_size in batch_sizes] + + job_queue = queue.Queue() + for job in jobs: + job_queue.put(job) + for _ in range(num_prefetch_workers): + job_queue.put(STOP) + + input_queue = queue.Queue(maxsize=max(2 * sum(batch_sizes), 2)) + output_queue = queue.Queue(maxsize=max(2 * len(model_devices), 2)) + stop_event = threading.Event() + error_box = [] + error_lock = threading.Lock() + remaining_producers = AtomicCounter(num_prefetch_workers) + remaining_consumers = AtomicCounter(len(model_devices)) + + def record_error(exc): + with error_lock: + if not error_box: + error_box.append(exc) + stop_event.set() + + def producer(): + try: + while True: + spec = job_queue.get() + if spec is STOP or stop_event.is_set(): + break + _safe_put(input_queue, PipelineJob(spec, load_fn(spec)), stop_event) + except PipelineAborted: + pass + except Exception as exc: # noqa + record_error(exc) + finally: + if remaining_producers.decrement() == 0 and not stop_event.is_set(): + for _ in model_devices: + input_queue.put(STOP) + + def consumer(worker_id): + model, device = model_devices[worker_id] + batch_size = batch_sizes[worker_id] + try: + while True: + batch = [] + got_stop = False + while len(batch) < batch_size: + item = _safe_get(input_queue, stop_event) + if item is STOP: + got_stop = True + break + batch.append(item) + + if batch: + with torch.no_grad(): + predictions = predict_fn(model, [item.data for item in batch], device) + if len(predictions) != len(batch): + raise RuntimeError( + f"The batch predictor returned {len(predictions)} outputs for {len(batch)} inputs." + ) + for item, prediction in zip(batch, predictions): + item.data = prediction + _safe_put(output_queue, item, stop_event) + + if got_stop: + break + except PipelineAborted: + pass + except Exception as exc: # noqa + record_error(exc) + finally: + if remaining_consumers.decrement() == 0 and not stop_event.is_set(): + output_queue.put(STOP) + + def writer(): + try: + while True: + item = _safe_get(output_queue, stop_event) + if item is STOP: + break + write_fn(item.spec, item.data) + if update_progress is not None: + update_progress(1) + except PipelineAborted: + pass + except Exception as exc: # noqa + record_error(exc) + + writer_thread = threading.Thread(target=writer, name="sam2-writer") + consumer_threads = [ + threading.Thread(target=consumer, args=(worker_id,), name=f"sam2-consumer-{worker_id}") + for worker_id in range(len(model_devices)) + ] + producer_threads = [ + threading.Thread(target=producer, name=f"sam2-producer-{worker_id}") + for worker_id in range(num_prefetch_workers) + ] + threads = [writer_thread, *consumer_threads, *producer_threads] + + try: + for thread in threads: + thread.start() + for thread in producer_threads: + thread.join() + for thread in consumer_threads: + thread.join() + writer_thread.join() + finally: + stop_event.set() + for thread in threads: + thread.join() + + if error_box: + raise error_box[0] + + +IMAGE_MEAN = (0.485, 0.456, 0.406) +IMAGE_STD = (0.229, 0.224, 0.225) + + +def _clear_group(group) -> None: + for key in list(group.keys()): + del group[key] + + +def _embedding_cache_complete(features, root) -> bool: + return bool(features.attrs.get("complete", "input_size" in root.attrs)) + + +def _prepare_encoder_pipeline( + predictor, n_jobs: int, batch_size: Optional[int], devices: Devices, +) -> Tuple[torch.nn.Module, List[Tuple[torch.nn.Module, torch.device]], List[int]]: + model = getattr(predictor, "model", predictor) + resolved_devices = resolve_devices(model, devices) + model_devices = prepare_models(model, resolved_devices) + if batch_size is None: + image_size = int(getattr(model, "image_size", getattr(predictor, "image_size", 1024))) + batch_sizes = compute_auto_batch_sizes( + model_devices=model_devices, + n_jobs=n_jobs, + patch_shape=(image_size, image_size), + in_channels=3, + prediction_function=lambda this_model, inputs: this_model.forward_image(inputs), + ) + else: + if int(batch_size) < 1: + raise ValueError(f"batch_size must be positive or None, got {batch_size}.") + batch_sizes = [int(batch_size)] * len(model_devices) + return model, model_devices, batch_sizes + + +def _forward_image_batch( + model: torch.nn.Module, items: List[Dict], device: torch.device, feature_sizes: Sequence, +) -> List[Dict]: + batch = torch.stack([item["tensor"] for item in items]).to(device, non_blocking=True) + backbone_out = model.forward_image(batch) + _, vision_feats, _, _ = model._prepare_backbone_features(backbone_out) + if model.directly_add_no_mem_embed: + vision_feats[-1] = vision_feats[-1] + model.no_mem_embed + + batch_size = len(items) + features = [ + feat.permute(1, 2, 0).reshape(batch_size, -1, *feat_size) + for feat, feat_size in zip(vision_feats[::-1], feature_sizes[::-1]) + ][::-1] + features = [feature.detach().cpu().numpy() for feature in features] + return [ + { + "features": features[-1][index:index + 1], + "high_res_feats": [feature[index:index + 1] for feature in features[:-1]], + "original_size": items[index]["original_size"], + } + for index in range(batch_size) + ] + + +def compute_tiled_2d( + input_: np.ndarray, + predictor, + tile_shape: Tuple[int, int], + halo: Tuple[int, int], + root, + save_path: Optional[Union[str, os.PathLike]], + pbar_init: Callable, + pbar_update: Callable, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> Dict: + """Compute 2D tile embeddings with batched encoders and queued input/output I/O.""" + from bioimage_cpp.utils import Blocking + from micro_sam.v2.util import _write_embedding_signature + + features = root.require_group("features") + high_res_group = root.require_group("high_res_feats") + if save_path is not None and "shape" in features.attrs and _embedding_cache_complete(features, root): + return {"features": features, "high_res_feats": high_res_group, "input_size": None, "original_size": None} + + _clear_group(features) + _clear_group(high_res_group) + + tiling = Blocking([0, 0], list(input_.shape[:2]), list(tile_shape)) + n_tiles = tiling.number_of_blocks + features.attrs["shape"] = list(input_.shape[:2]) + features.attrs["tile_shape"] = list(tile_shape) + features.attrs["halo"] = list(halo) + features.attrs["complete"] = False + + pbar_init(n_tiles, "Compute Image Embeddings 2D tiled") + model, model_devices, batch_sizes = _prepare_encoder_pipeline(predictor, n_tiles, batch_size, devices) + feature_sizes = predictor._bb_feat_sizes + + def load_tile(tile_id): + block = tiling.get_block_with_halo(tile_id, list(halo)).outer_block + bb = tuple(slice(begin, end) for begin, end in zip(block.begin, block.end)) + image = to_image(np.asarray(input_[bb])) + return {"tensor": predictor._transforms(image), "original_size": tuple(image.shape[:2])} + + def predict_tiles(this_model, items, device): + return _forward_image_batch(this_model, items, device, feature_sizes) + + def write_tile(tile_id, result): + name = str(tile_id) + dataset = _create_dataset_with_data(features, name, data=result["features"]) + dataset.attrs["input_size"] = model.image_size + dataset.attrs["original_size"] = [list(result["original_size"])] + + tile_high_res = high_res_group.require_group(name) + for level, feature in enumerate(result["high_res_feats"]): + _create_dataset_with_data(tile_high_res, str(level), data=feature) + + try: + run_batched_pipeline( + jobs=range(n_tiles), + model_devices=model_devices, + batch_sizes=batch_sizes, + load_fn=load_tile, + predict_fn=predict_tiles, + write_fn=write_tile, + num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, + ) + finally: + release_model_replicas(model_devices) + + if save_path is not None: + _write_embedding_signature( + root, input_, predictor, tile_shape=tile_shape, halo=halo, input_size=None, original_size=None, + ) + features.attrs["complete"] = True + return {"features": features, "high_res_feats": high_res_group, "input_size": None, "original_size": None} + + +def _prepare_video_frame(raw: np.ndarray, image_size: int) -> torch.Tensor: + from micro_sam.v2.models._video_predictor import _load_img_as_tensor + + image, _, _ = _load_img_as_tensor(np.asarray(raw), image_size) + mean = torch.tensor(IMAGE_MEAN, dtype=torch.float32)[:, None, None] + std = torch.tensor(IMAGE_STD, dtype=torch.float32)[:, None, None] + return (image - mean) / std + + +def _forward_video_batch(model: torch.nn.Module, items: List[Dict], device: torch.device) -> List[Dict]: + batch = torch.stack([item["tensor"] for item in items]).to(device, non_blocking=True) + backbone_out = model.forward_image(batch) + + vision_features = backbone_out["vision_features"].detach().cpu().numpy() + pos_enc = [value.detach().cpu().numpy() for value in backbone_out["vision_pos_enc"]] + fpn = [value.detach().cpu().numpy() for value in backbone_out["backbone_fpn"]] + return [ + { + "features": vision_features[index:index + 1], + "pos_enc": [value[index:index + 1] for value in pos_enc], + "fpn": [value[index:index + 1] for value in fpn], + "original_size": items[index]["original_size"], + } + for index in range(len(items)) + ] + + +def _create_feature_dataset(group, name: str, n_slices: int, tensor: np.ndarray): + shape = (n_slices,) + tuple(tensor.shape) + chunks = (1,) + tuple(tensor.shape) + return _create_dataset_without_data(group, name, shape=shape, dtype="float32", chunks=chunks) + + +def _create_feature_levels(group, n_slices: int, tensors: Sequence[np.ndarray]) -> List: + return [ + _create_feature_dataset(group, str(level), n_slices, tensor) + for level, tensor in enumerate(tensors) + ] + + +def _load_feature_levels(group, lazy_loading: bool) -> List: + values = [] + level = 0 + while str(level) in group: + dataset = group[str(level)] + values.append(dataset if lazy_loading else dataset[:]) + level += 1 + return values + + +def compute_3d( + input_: np.ndarray, + predictor, + root, + save_path: Optional[Union[str, os.PathLike]], + lazy_loading: bool, + pbar_init: Callable, + pbar_update: Callable, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> Dict: + """Compute volume embeddings by batching slices and overlapping preprocessing and zarr writes.""" + from micro_sam.v2.util import _write_embedding_signature + + if save_path is not None and "original_size" in root.attrs: + features = root["features"] if lazy_loading else root["features"][:] + return { + "features": features, + "pos_enc": _load_feature_levels(root["pos_enc"], lazy_loading), + "fpn": _load_feature_levels(root["fpn"], lazy_loading), + "input_size": root.attrs["input_size"], + "original_size": root.attrs["original_size"], + } + + for key in ("features", "pos_enc", "fpn"): + if key in root: + del root[key] + + n_slices = int(input_.shape[0]) + image_size = int(predictor.image_size) + pbar_init(n_slices, "Compute Image Embeddings 3D") + model, model_devices, batch_sizes = _prepare_encoder_pipeline(predictor, n_slices, batch_size, devices) + + if save_path is None: + feature_values = [None] * n_slices + pos_values = [None] * n_slices + fpn_values = [None] * n_slices + feature_dataset = None + pos_datasets = None + fpn_datasets = None + else: + feature_values = pos_values = fpn_values = None + feature_dataset = None + pos_datasets = None + fpn_datasets = None + + def load_slice(z): + raw = np.asarray(input_[z]) + return { + "tensor": _prepare_video_frame(raw, image_size), + "original_size": tuple(int(value) for value in raw.shape[:2]), + } + + def write_slice(z, result): + nonlocal feature_dataset, pos_datasets, fpn_datasets + if save_path is None: + feature_values[z] = torch.from_numpy(result["features"]) + pos_values[z] = [torch.from_numpy(value) for value in result["pos_enc"]] + fpn_values[z] = [torch.from_numpy(value) for value in result["fpn"]] + return + + if feature_dataset is None: + feature_dataset = _create_feature_dataset(root, "features", n_slices, result["features"]) + pos_datasets = _create_feature_levels(root.require_group("pos_enc"), n_slices, result["pos_enc"]) + fpn_datasets = _create_feature_levels(root.require_group("fpn"), n_slices, result["fpn"]) + feature_dataset[z] = result["features"] + for dataset, value in zip(pos_datasets, result["pos_enc"]): + dataset[z] = value + for dataset, value in zip(fpn_datasets, result["fpn"]): + dataset[z] = value + + try: + run_batched_pipeline( + jobs=range(n_slices), + model_devices=model_devices, + batch_sizes=batch_sizes, + load_fn=load_slice, + predict_fn=_forward_video_batch, + write_fn=write_slice, + num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, + ) + finally: + release_model_replicas(model_devices) + + original_size = tuple(int(value) for value in input_.shape[-2:]) + if save_path is None: + features = torch.cat(feature_values).numpy() + n_levels = len(pos_values[0]) + pos_enc = [torch.stack([value[level] for value in pos_values]) for level in range(n_levels)] + fpn = [torch.stack([value[level] for value in fpn_values]) for level in range(n_levels)] + else: + _write_embedding_signature( + root, input_, predictor, tile_shape=None, halo=None, + input_size=image_size, original_size=original_size, + ) + features = feature_dataset if lazy_loading else feature_dataset[:] + pos_enc = _load_feature_levels(root["pos_enc"], lazy_loading) + fpn = _load_feature_levels(root["fpn"], lazy_loading) + + return { + "features": features, + "pos_enc": pos_enc, + "fpn": fpn, + "input_size": image_size, + "original_size": original_size, + } + + +def compute_tiled_3d( + input_: np.ndarray, + predictor, + tile_shape: Tuple[int, int], + halo: Tuple[int, int], + root, + save_path: Optional[Union[str, os.PathLike]], + pbar_init: Callable, + pbar_update: Callable, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> Dict: + """Compute tile/slice embeddings as one pipelined job stream across all available GPUs.""" + from bioimage_cpp.utils import Blocking + from micro_sam.v2.util import _write_embedding_signature + + features = root.require_group("features") + pos_enc_group = root.require_group("pos_enc") + fpn_group = root.require_group("fpn") + if save_path is not None and "shape" in features.attrs and _embedding_cache_complete(features, root): + return { + "features": features, + "pos_enc": pos_enc_group, + "fpn": fpn_group, + "input_size": None, + "original_size": None, + } + + _clear_group(features) + _clear_group(pos_enc_group) + _clear_group(fpn_group) + + tiling = Blocking([0, 0], list(input_.shape[1:]), list(tile_shape)) + n_tiles = tiling.number_of_blocks + n_slices = int(input_.shape[0]) + image_size = int(predictor.image_size) + jobs = [(tile_id, z) for tile_id in range(n_tiles) for z in range(n_slices)] + + features.attrs["shape"] = list(input_.shape) + features.attrs["tile_shape"] = list(tile_shape) + features.attrs["halo"] = list(halo) + features.attrs["complete"] = False + pbar_init(len(jobs), "Compute Image Embeddings 3D tiled") + + bounds = {} + for tile_id in range(n_tiles): + block = tiling.get_block_with_halo(tile_id, list(halo)).outer_block + bounds[tile_id] = tuple(slice(begin, end) for begin, end in zip(block.begin, block.end)) + + model, model_devices, batch_sizes = _prepare_encoder_pipeline(predictor, len(jobs), batch_size, devices) + tile_datasets = {} + + def load_tile_slice(job): + tile_id, z = job + bb = bounds[tile_id] + raw = np.asarray(input_[z, bb[0], bb[1]]) + return { + "tensor": _prepare_video_frame(raw, image_size), + "original_size": tuple(int(value) for value in raw.shape[:2]), + } + + def write_tile_slice(job, result): + tile_id, z = job + if tile_id not in tile_datasets: + name = str(tile_id) + feature_dataset = _create_feature_dataset(features, name, n_slices, result["features"]) + feature_dataset.attrs["input_size"] = image_size + feature_dataset.attrs["original_size"] = list(result["original_size"]) + pos_datasets = _create_feature_levels( + pos_enc_group.require_group(name), n_slices, result["pos_enc"], + ) + fpn_datasets = _create_feature_levels( + fpn_group.require_group(name), n_slices, result["fpn"], + ) + tile_datasets[tile_id] = feature_dataset, pos_datasets, fpn_datasets + + feature_dataset, pos_datasets, fpn_datasets = tile_datasets[tile_id] + feature_dataset[z] = result["features"] + for dataset, value in zip(pos_datasets, result["pos_enc"]): + dataset[z] = value + for dataset, value in zip(fpn_datasets, result["fpn"]): + dataset[z] = value + + try: + run_batched_pipeline( + jobs=jobs, + model_devices=model_devices, + batch_sizes=batch_sizes, + load_fn=load_tile_slice, + predict_fn=_forward_video_batch, + write_fn=write_tile_slice, + num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, + ) + finally: + release_model_replicas(model_devices) + + if save_path is not None: + _write_embedding_signature( + root, input_, predictor, tile_shape=tile_shape, halo=halo, input_size=None, original_size=None, + ) + features.attrs["complete"] = True + return { + "features": features, + "pos_enc": pos_enc_group, + "fpn": fpn_group, + "input_size": None, + "original_size": None, + } + + +class EmptyEncoder(torch.nn.Module): + """Lightweight placeholder used while replicating decoder-only UniSAM2 models.""" + + def __init__(self, img_size: int) -> None: + super().__init__() + self.img_size = img_size + + def forward(self, x: torch.Tensor): + raise RuntimeError("The decoder-only placeholder must be replaced before inference.") + + +def _normalize_feature_block(block: np.ndarray) -> torch.Tensor: + block = np.asarray(block) + if block.ndim == 5 and block.shape[1] == 1: + block = block[:, 0] + if block.ndim == 3: + block = block[None] + if block.ndim != 4: + raise ValueError(f"Expected a (Z, C, H, W) feature block, got shape {block.shape}.") + return torch.from_numpy(np.asarray(block, dtype="float32")) + + +def _load_job(job: Dict) -> Dict: + source = job["source"] + selection = job.get("selection") + block = np.asarray(source) if selection is None else np.asarray(source[selection]) + return { + "feature": _normalize_feature_block(block), + "original_size": job["original_size"], + } + + +def _predict_jobs(model: torch.nn.Module, items: List[Dict], device: torch.device) -> List[np.ndarray]: + from micro_sam.v2.instance_segmentation import _decode_3d_feature_batch + + original_size = items[0]["original_size"] + if any(item["original_size"] != original_size for item in items): + raise RuntimeError("Decoder jobs with different output sizes cannot share a batch.") + features = torch.stack([item["feature"] for item in items]).to(device, non_blocking=True) + output = _decode_3d_feature_batch(model, features, original_size, device) + return [np.array(value) for value in output] + + +def _feature_shape(job: Dict) -> Tuple[int, ...]: + source_shape = tuple(job["source"].shape) + selection = job.get("selection") + if selection is None: + shape = source_shape + else: + start = 0 if selection.start is None else selection.start + stop = source_shape[0] if selection.stop is None else selection.stop + shape = (stop - start,) + source_shape[1:] + + if len(shape) == 5 and shape[1] == 1: + shape = (shape[0],) + shape[2:] + if len(shape) == 3: + shape = (1,) + shape + if len(shape) != 4: + raise ValueError(f"Expected feature shape (Z, C, H, W), got {shape}.") + return shape + + +def _prepare_decoder_pipeline( + model: torch.nn.Module, jobs: List[Dict], batch_size: Optional[int], devices: Devices, +) -> Tuple[List[Tuple[torch.nn.Module, torch.device]], List[int], torch.nn.Module]: + resolved_devices = resolve_devices(model, devices) + encoder = model.encoder + model.encoder = EmptyEncoder(getattr(encoder, "img_size", 1024)) + try: + model_devices = prepare_models(model, resolved_devices) + except Exception: + model.encoder = encoder + raise + + if batch_size is not None: + if int(batch_size) < 1: + release_model_replicas(model_devices) + model.encoder = encoder + raise ValueError(f"batch_size must be positive or None, got {batch_size}.") + return model_devices, [int(batch_size)] * len(model_devices), encoder + + representative = max( + jobs, + key=lambda job: np.prod(_feature_shape(job)) * np.prod(job["original_size"]), + ) + z, channels, height, width = _feature_shape(representative) + original_size = representative["original_size"] + + def prediction_function(this_model, inputs): + from micro_sam.v2.instance_segmentation import _decode_3d_feature_batch + + features = inputs.permute(0, 2, 1, 3, 4) + return _decode_3d_feature_batch( + this_model, features, original_size, inputs.device, + ) + + try: + batch_sizes = compute_auto_batch_sizes( + model_devices=model_devices, + n_jobs=len(jobs), + patch_shape=(z, height, width), + in_channels=channels, + prediction_function=prediction_function, + ) + except Exception: + release_model_replicas(model_devices) + model.encoder = encoder + raise + return model_devices, batch_sizes, encoder + + +def _run_decoder_jobs( + model: torch.nn.Module, + jobs: List[Dict], + write_fn: Callable[[Dict, np.ndarray], None], + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> None: + if len(jobs) == 0: + return + + model_devices, batch_sizes, encoder = _prepare_decoder_pipeline(model, jobs, batch_size, devices) + groups = defaultdict(list) + for job in jobs: + groups[(_feature_shape(job), job["original_size"])].append(job) + + try: + for group_jobs in groups.values(): + run_batched_pipeline( + jobs=group_jobs, + model_devices=model_devices, + batch_sizes=batch_sizes, + load_fn=_load_job, + predict_fn=_predict_jobs, + write_fn=write_fn, + num_prefetch_workers=num_prefetch_workers, + ) + finally: + release_model_replicas(model_devices) + model.encoder = encoder + + +def decode_volume_embeddings( + model: torch.nn.Module, + image_embeddings: Dict, + device: Device = None, + z_block: Optional[int] = None, + z_halo: Optional[int] = None, + pbar_init: Optional[Callable] = None, + pbar_update: Optional[Callable] = None, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> np.ndarray: + """Decode z blocks from non-tiled volume embeddings with queued reads and batched inference.""" + from micro_sam.v2.util import DEFAULT_HALO_Z, DEFAULT_TILE_Z + + features = image_embeddings["features"] + n_dims = len(features.shape) + if n_dims not in (4, 5): + raise ValueError( + f"Decoder-from-embeddings (3d) requires features with ndim 4 or 5, got {n_dims}." + ) + + n_slices = int(features.shape[0]) + z_block = DEFAULT_TILE_Z if z_block is None else int(z_block) + z_halo = DEFAULT_HALO_Z if z_halo is None else int(z_halo) + if z_block < 1 or z_halo < 0: + raise ValueError(f"z_block must be positive and z_halo non-negative, got {z_block}, {z_halo}.") + + original_size = tuple(int(value) for value in np.asarray(image_embeddings["original_size"]).reshape(-1)[:2]) + output = np.zeros((4, n_slices, *original_size), dtype="float32") + jobs = [] + for z0 in range(0, n_slices, z_block): + z1 = min(z0 + z_block, n_slices) + c0, c1 = max(0, z0 - z_halo), min(n_slices, z1 + z_halo) + jobs.append({ + "source": features, + "selection": slice(c0, c1), + "original_size": original_size, + "z0": z0, + "z1": z1, + "c0": c0, + }) + + if pbar_init is not None: + pbar_init(n_slices, "Automatic segmentation (volume)") + + def write_prediction(job, prediction): + inner = job["z0"] - job["c0"] + count = job["z1"] - job["z0"] + output[:, job["z0"]:job["z1"]] = prediction[:, inner:inner + count] + if pbar_update is not None: + pbar_update(count) + + _run_decoder_jobs( + model, jobs, write_prediction, batch_size=batch_size, devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) + return output + + +def _tiled_metadata(image_embeddings: Dict, is_3d: bool) -> Tuple: + from bioimage_cpp.utils import Blocking + + features = image_embeddings["features"] + shape = tuple(int(value) for value in features.attrs["shape"]) + tile_shape = tuple(int(value) for value in features.attrs["tile_shape"]) + halo = tuple(int(value) for value in features.attrs["halo"]) + spatial_shape = shape[1:] if is_3d else shape + tiling = Blocking([0, 0], list(spatial_shape), list(tile_shape)) + return features, shape, halo, tiling + + +def decode_tiled_2d_embeddings( + model: torch.nn.Module, + image_embeddings: Dict, + device: Device = None, + pbar_init: Optional[Callable] = None, + pbar_update: Optional[Callable] = None, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> np.ndarray: + """Decode and stitch 2D embedding tiles in automatically sized batches.""" + features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=False) + output = np.zeros((4, *shape), dtype="float32") + jobs = [] + for tile_id in range(tiling.number_of_blocks): + tile_features = features[str(tile_id)] + original_size = tuple(int(value) for value in np.asarray(tile_features.attrs["original_size"]).reshape(-1)[:2]) + jobs.append({ + "source": tile_features, + "selection": None, + "original_size": original_size, + "tile_id": tile_id, + }) + + if pbar_init is not None: + pbar_init(len(jobs), "Automatic segmentation (tiles)") + + def write_prediction(job, prediction): + block = tiling.get_block_with_halo(job["tile_id"], halo=list(halo)) + local = tuple(slice(begin, end) for begin, end in zip( + block.inner_block_local.begin, block.inner_block_local.end, + )) + inner = tuple(slice(begin, end) for begin, end in zip( + block.inner_block.begin, block.inner_block.end, + )) + output[(slice(None),) + inner] = prediction[(slice(None), slice(0, 1)) + local][:, 0] + if pbar_update is not None: + pbar_update(1) + + _run_decoder_jobs( + model, jobs, write_prediction, batch_size=batch_size, devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) + return output + + +def _tiled_3d_jobs(features, tiling, n_slices: int, z_block: int, z_halo: int) -> List[Dict]: + jobs = [] + for tile_id in range(tiling.number_of_blocks): + tile_features = features[str(tile_id)] + original_size = tuple(int(value) for value in np.asarray(tile_features.attrs["original_size"]).reshape(-1)[:2]) + for z0 in range(0, n_slices, z_block): + z1 = min(z0 + z_block, n_slices) + c0, c1 = max(0, z0 - z_halo), min(n_slices, z1 + z_halo) + jobs.append({ + "source": tile_features, + "selection": slice(c0, c1), + "original_size": original_size, + "tile_id": tile_id, + "z0": z0, + "z1": z1, + "c0": c0, + }) + return jobs + + +def decode_tiled_3d_embeddings( + model: torch.nn.Module, + image_embeddings: Dict, + device: Device = None, + pbar_init: Optional[Callable] = None, + pbar_update: Optional[Callable] = None, + z_block: Optional[int] = None, + z_halo: Optional[int] = None, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> np.ndarray: + """Batch over both tile columns and z blocks while decoding tiled 3D embeddings.""" + from micro_sam.v2.util import DEFAULT_HALO_Z, DEFAULT_TILE_Z + + features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=True) + n_slices = shape[0] + z_block = DEFAULT_TILE_Z if z_block is None else int(z_block) + z_halo = DEFAULT_HALO_Z if z_halo is None else int(z_halo) + jobs = _tiled_3d_jobs(features, tiling, n_slices, z_block, z_halo) + output = np.zeros((4, *shape), dtype="float32") + + if pbar_init is not None: + pbar_init(tiling.number_of_blocks * n_slices, "Automatic segmentation (tiles)") + + def write_prediction(job, prediction): + block = tiling.get_block_with_halo(job["tile_id"], halo=list(halo)) + local = tuple(slice(begin, end) for begin, end in zip( + block.inner_block_local.begin, block.inner_block_local.end, + )) + inner = tuple(slice(begin, end) for begin, end in zip( + block.inner_block.begin, block.inner_block.end, + )) + z_count = job["z1"] - job["z0"] + local_z = job["z0"] - job["c0"] + prediction = prediction[:, local_z:local_z + z_count] + output[(slice(None), slice(job["z0"], job["z1"])) + inner] = prediction[(slice(None), slice(None)) + local] + if pbar_update is not None: + pbar_update(z_count) + + _run_decoder_jobs( + model, jobs, write_prediction, batch_size=batch_size, devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) + return output + + +def decode_tiled_3d_slice( + model: torch.nn.Module, + image_embeddings: Dict, + index: int, + device: Device = None, + pbar_init: Optional[Callable] = None, + pbar_update: Optional[Callable] = None, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, +) -> np.ndarray: + """Decode one volume slice across all embedding tiles in batches.""" + features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=True) + output = np.zeros((4, *shape[1:]), dtype="float32") + jobs = [] + for tile_id in range(tiling.number_of_blocks): + tile_features = features[str(tile_id)] + original_size = tuple(int(value) for value in np.asarray(tile_features.attrs["original_size"]).reshape(-1)[:2]) + jobs.append({ + "source": tile_features, + "selection": slice(index, index + 1), + "original_size": original_size, + "tile_id": tile_id, + }) + + if pbar_init is not None: + pbar_init(len(jobs), "Automatic segmentation (tiles)") + + def write_prediction(job, prediction): + block = tiling.get_block_with_halo(job["tile_id"], halo=list(halo)) + local = tuple(slice(begin, end) for begin, end in zip( + block.inner_block_local.begin, block.inner_block_local.end, + )) + inner = tuple(slice(begin, end) for begin, end in zip( + block.inner_block.begin, block.inner_block.end, + )) + output[(slice(None),) + inner] = prediction[(slice(None), 0) + local] + if pbar_update is not None: + pbar_update(1) + + _run_decoder_jobs( + model, jobs, write_prediction, batch_size=batch_size, devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) + return output diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index a0537ec84..cb44fccd5 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -36,7 +36,7 @@ from micro_sam.v2.postprocessing import flow_instance_segmentation, run_multicut from micro_sam.v2.util import ( precompute_image_embeddings, set_precomputed, _load_list_datasets, - _DEFAULT_MODEL, DEFAULT_TILE_Z, DEFAULT_HALO_Z, + _DEFAULT_MODEL, DEFAULT_TILE_Z, DEFAULT_HALO_Z, Devices, ) DEFAULT_SEGMENTATION_MODE_WITH_DECODER = "ais" @@ -843,40 +843,44 @@ def _block_shape_and_halo(spatial_shape, ndim, tile_shape, halo): class _StubEncoder(torch.nn.Module): """Encoder replacement that returns precomputed per-slice features in call order. - `UNETR3D.forward` runs ``[self.encoder(x[:, :, i])[0] for i in range(Z)]``, i.e. it calls the - encoder once per z slice in order (once total for a single 2d slice, the Z=1 case) and has no - encoder skip connections. Returning the i-th precomputed slice feature on the i-th call therefore - reproduces the full forward without re-running the SAM2 image encoder, for both 2d and 3d. + For a single feature volume, `features` has shape (Z, C, H, W). Batched decoder inference uses + (B, Z, C, H, W), in which case each encoder call returns the corresponding slice for all batch + elements. This reproduces the encoder outputs expected by `UNETR3D.forward` without re-encoding. """ def __init__(self, features: torch.Tensor, img_size: int = 1024) -> None: super().__init__() - self.features = features # (Z, C, h, w); a single 2d slice is the Z=1 case + self.features = features self.img_size = img_size self._idx = 0 def forward(self, x): # noqa - feature = self.features[self._idx:self._idx + 1] + if self.features.ndim == 5: + feature = self.features[:, self._idx] + else: + feature = self.features[self._idx:self._idx + 1] self._idx += 1 return [feature] -def _decode_3d_feature_block(model, feature, original_size, device): - """Run the UNETR3D decoder on a ``(z, C, h, w)`` feature block via the stub encoder. - - Temporarily swaps the model's encoder for `_StubEncoder` (which returns the precomputed per-slice - features in order) and runs the decoder on an original-size dummy, so the model resizes the - prediction back to ``(z, H, W)`` itself. Returns the prediction for this z block, shape ``(4, z, H, W)``. - """ +def _decode_3d_feature_batch(model, features, original_size, device): + """Decode precomputed feature volumes with shape (B, Z, C, H, W).""" + if features.ndim != 5: + raise ValueError(f"Expected batched features with shape (B, Z, C, H, W), got {features.shape}.") img_size = getattr(model.encoder, "img_size", 1024) real_encoder = model.encoder - model.encoder = _StubEncoder(feature, img_size) + model.encoder = _StubEncoder(features, img_size) try: - dummy = torch.zeros((1, 3, feature.shape[0], *original_size), device=device) - output = model(dummy) # (1, 4, z, H, W) + dummy = torch.zeros((features.shape[0], 3, features.shape[1], *original_size), device=device) + output = model(dummy) finally: model.encoder = real_encoder - return output[0].detach().cpu().numpy() # (4, z, H, W) + return output.detach().cpu().numpy() + + +def _decode_3d_feature_block(model, feature, original_size, device): + """Decode one precomputed feature volume with shape (Z, C, H, W).""" + return _decode_3d_feature_batch(model, feature.unsqueeze(0), original_size, device)[0] def _segment_from_predictions(prediction: np.ndarray, mode: str = "sparse", **kwargs) -> np.ndarray: @@ -935,24 +939,31 @@ def __init__(self, model: torch.nn.Module, device: Optional[Union[str, torch.dev self._prediction = None self._is_initialized = False - def _run_full_inference(self, raw, ndim, tile_shape=None, halo=None, pbar_init=None, pbar_update=None): - """Run the UniSAM2 model (encoder + decoder) to predict foreground and directed distances. + def _run_full_inference( + self, raw, ndim, tile_shape=None, halo=None, pbar_init=None, pbar_update=None, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, num_write_workers=1, + ): + """Run queued, batched UniSAM2 encoder and decoder inference. - Inference is tiled with a halo (3d: fully 3d; 2d: in-plane); with `tile_shape=None` the whole - image is a single block. Each block is percentile-normalized before prediction. Returns the - predictions stacked along the channel axis, shape (4, *spatial). + Tiled reads and preprocessing, GPU inference, and output writes overlap through the torch-em + prediction pipeline. If `batch_size` is None, the largest safe batch is probed independently + on every selected CUDA device. """ - from torch_em.util.prediction import predict_with_halo + from torch_em.util.prediction import predict_with_halo_pipelined from micro_sam.v2.normalization import normalize_raw + from micro_sam.v2.batched_inference import ( + compute_auto_batch_sizes, prepare_models, release_model_replicas, resolve_devices, + ) def _preprocess(crop): return np.concatenate([normalize_raw(crop)] * 3, axis=0) is_3d = ndim == 3 block_shape, block_halo = _block_shape_and_halo(tuple(raw.shape), ndim, tile_shape, halo) + n_blocks = _n_blocks(tuple(raw.shape), ndim, block_shape) if pbar_init is not None: desc = "Automatic segmentation (volume)" if is_3d else "Automatic segmentation" - pbar_init(_n_blocks(tuple(raw.shape), ndim, block_shape), desc) + pbar_init(n_blocks, desc) if is_3d: input_ = raw[np.newaxis].astype("float32") @@ -963,17 +974,39 @@ def _preprocess(crop): img_size = getattr(getattr(self._model, "encoder", None), "img_size", 1024) resize_model = ResizeLongestSideWrapper(self._model, img_size) + resolved_devices = resolve_devices(resize_model, devices) + + if batch_size is None: + model_devices = prepare_models(resize_model, resolved_devices) + patch_shape = tuple(block + 2 * overlap for block, overlap in zip(block_shape, block_halo)) + try: + batch_sizes = compute_auto_batch_sizes( + model_devices=model_devices, + n_jobs=n_blocks, + patch_shape=patch_shape, + in_channels=3, + prediction_function=lambda this_model, inputs: this_model(inputs), + ) + batch_size = min(batch_sizes) + finally: + release_model_replicas(model_devices) + elif int(batch_size) < 1: + raise ValueError(f"batch_size must be positive or None, got {batch_size}.") with _bridge_halo_progress(pbar_update): - output = predict_with_halo( + output = predict_with_halo_pipelined( input_=input_, model=resize_model, block_shape=block_shape, halo=block_halo, preprocess=_preprocess, - gpu_ids=[self._device] if self._device is not None else None, + gpu_ids=resolved_devices, output=output, with_channels=True, + batch_size=int(batch_size), + num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, + disable_tqdm=pbar_update is not None, ) if not is_3d: output = output[:, 0] @@ -1010,50 +1043,25 @@ def _run_decoder_2d(self, image_embeddings): return output[0, :, 0].detach().cpu().numpy() @torch.no_grad() - def _run_decoder_3d(self, image_embeddings, z_block=None, z_halo=None, pbar_init=None, pbar_update=None): - """Run only the UniSAM2 decoder on precomputed 3d embeddings (no encoder pass). - - Reuses the per-slice ``vision_features`` (shape ``(Z, C, h, w)``) produced for the volume - - the same embeddings used for interactive 3d segmentation. The decoder pass is chunked along z - (with a halo for 3d-conv context) so a deep volume does not decode the whole stack at once and - run out of memory. Returns predictions stacked along the channel axis, shape ``(4, Z, H, W)``. - """ - z_block = DEFAULT_TILE_Z if z_block is None else int(z_block) - z_halo = DEFAULT_HALO_Z if z_halo is None else int(z_halo) - - # Keep 'features' as the (possibly lazy, on-disk) handle and read only the current z block inside - # the loop, so peak memory is one z block instead of the whole feature volume. Use 'len(shape)' as - # zarr / z5py datasets and numpy arrays all expose 'shape' (but not always 'ndim'). - features = image_embeddings["features"] - n_dims = len(features.shape) - if n_dims not in (4, 5): - raise ValueError( - f"Decoder-from-embeddings (3d) requires 3d embeddings (features with ndim 4 or 5), got {n_dims}." - ) - # Per-slice features are (Z, C, h, w); the tiled / save-path layout keeps a singleton batch axis - # (Z, 1, C, h, w) - squeeze it per block so the stub returns a (1, C, h, w) feature per slice. - squeeze_batch = n_dims == 5 and features.shape[1] == 1 - n_slices = features.shape[0] - original_size = tuple(int(s) for s in np.array(image_embeddings["original_size"]).reshape(-1)[:2]) - - z_starts = list(range(0, n_slices, z_block)) if n_slices > z_block else [0] - if pbar_init is not None: - pbar_init(n_slices, "Automatic segmentation (volume)") - - output = np.zeros((4, n_slices, *original_size), dtype="float32") - for z0 in z_starts: - z1 = min(z0 + z_block, n_slices) - c0, c1 = max(0, z0 - z_halo), min(n_slices, z1 + z_halo) - block = np.asarray(features[c0:c1]) - if squeeze_batch: - block = block[:, 0] - feature_block = torch.as_tensor(block, device=self._device).float() - pred = _decode_3d_feature_block(self._model, feature_block, original_size, self._device) - inner = z0 - c0 - output[:, z0:z1] = pred[:, inner:inner + (z1 - z0)] - if pbar_update is not None: - pbar_update(z1 - z0) # advance by the number of slices in this block - return output + def _run_decoder_3d( + self, image_embeddings, z_block=None, z_halo=None, pbar_init=None, pbar_update=None, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, + ): + """Decode queued z blocks from precomputed 3d embeddings.""" + from micro_sam.v2.batched_inference import decode_volume_embeddings + + return decode_volume_embeddings( + model=self._model, + image_embeddings=image_embeddings, + device=self._device, + z_block=z_block, + z_halo=z_halo, + pbar_init=pbar_init, + pbar_update=pbar_update, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) @torch.no_grad() def initialize( @@ -1068,42 +1076,56 @@ def initialize( pbar_update: Optional[callable] = None, z_block: Optional[int] = None, z_halo: Optional[int] = None, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> None: - """Run the UniSAM2 inference and store the foreground and distance predictions. + """Run the UniSAM2 inference and store foreground and distance predictions. Args: image: The input image, shape (Y, X) for 2d or (Z, Y, X) for 3d. ndim: The number of spatial dimensions (2 or 3). - image_embeddings: Optional precomputed image embeddings. If given, only the decoder is run - on them (no encoder pass), reusing the embeddings shared with interactive / AMG - for - both 2d and 3d. See `precompute_image_embeddings`. + image_embeddings: Optional precomputed image embeddings. If given, only the decoder runs. i: Index for the image data. Unused here, kept for interface compatibility. - tile_shape: Unused for the non-tiled segmenter (no tiling); kept so the interface matches - the tiled segmenter. - halo: Unused for the non-tiled segmenter; kept for interface compatibility. - pbar_init: Callback to initialize an external progress bar. The 2d decoder-on-embeddings - path is a single step; the 3d path reports per z block, and the full-inference path - reports per block. + tile_shape: Unused for the non-tiled segmenter. + halo: Unused for the non-tiled segmenter. + pbar_init: Callback to initialize an external progress bar. pbar_update: Callback to update an external progress bar. - z_block: Number of slices per z block for the 3d decoder pass (defaults to `DEFAULT_TILE_Z`). - z_halo: Overlapping slices between z blocks (defaults to `DEFAULT_HALO_Z`). + z_block: Number of slices per decoder z block. + z_halo: Overlapping decoder slices used as context. + batch_size: Explicit tile or z-block batch size, or None for automatic capacity probing. + devices: Inference device or devices. None uses all visible GPUs when the model is on CUDA. + num_prefetch_workers: Number of input reading and preprocessing threads. + num_write_workers: Number of output writing threads for full tiled inference. """ if image_embeddings is not None and ndim == 3: - # Decoder-only on precomputed 3d embeddings: progress advances per z block. self._prediction = self._run_decoder_3d( - image_embeddings, z_block=z_block, z_halo=z_halo, pbar_init=pbar_init, pbar_update=pbar_update, + image_embeddings, + z_block=z_block, + z_halo=z_halo, + pbar_init=pbar_init, + pbar_update=pbar_update, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, ) elif image_embeddings is not None: - # Decoder-only on precomputed 2d embeddings is a single step. if pbar_init is not None: pbar_init(1, "Automatic segmentation") self._prediction = self._run_decoder_2d(image_embeddings) if pbar_update is not None: pbar_update(1) else: - # Full inference reports per block (per z chunk for 3d). self._prediction = self._run_full_inference( - image, ndim, pbar_init=pbar_init, pbar_update=pbar_update, + image, + ndim, + pbar_init=pbar_init, + pbar_update=pbar_update, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) self._is_initialized = True @@ -1154,116 +1176,64 @@ class TiledUniSAM2InstanceSegmentation(UniSAM2InstanceSegmentation): """ @torch.no_grad() - def _run_decoder_tiled_2d(self, image_embeddings, pbar_init=None, pbar_update=None): - """Run the UniSAM2 decoder on precomputed tiled 2d embeddings and stitch the tiles. - - For each tile the decoder is run on the precomputed per-tile features (no encoder pass), and - the inner block of the per-tile prediction (the halo is used only as context) is written into - the full output - the same halo stitching as the micro-sam v1 tiled decoder. Returns the - predictions stacked along the channel axis, shape (4, Y, X). - """ - feats_group = image_embeddings["features"] - shape = tuple(int(s) for s in feats_group.attrs["shape"]) - tile_shape = tuple(int(s) for s in feats_group.attrs["tile_shape"]) - halo = tuple(int(s) for s in feats_group.attrs["halo"]) - tiling = Blocking([0, 0], list(shape), list(tile_shape)) - - if pbar_init is not None: - pbar_init(tiling.number_of_blocks, "Automatic segmentation (tiles)") - - output = np.zeros((4, *shape), dtype="float32") - for tile_id in range(tiling.number_of_blocks): - tile_features = feats_group[str(tile_id)] - # The UNETR decoder only needs the vision features, so we pass them as a single-image embedding. - tile_embeddings = { - "features": np.asarray(tile_features), "original_size": tile_features.attrs["original_size"], - } - tile_prediction = self._run_decoder_2d(tile_embeddings) - - block = tiling.get_block_with_halo(tile_id, halo=list(halo)) - local_bb = tuple(slice(b, e) for b, e in zip(block.inner_block_local.begin, block.inner_block_local.end)) - inner_bb = tuple(slice(b, e) for b, e in zip(block.inner_block.begin, block.inner_block.end)) - output[(slice(None),) + inner_bb] = tile_prediction[(slice(None),) + local_bb] - if pbar_update is not None: - pbar_update(1) - - return output + def _run_decoder_tiled_2d( + self, image_embeddings, pbar_init=None, pbar_update=None, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, + ): + """Decode and stitch queued 2d embedding tiles.""" + from micro_sam.v2.batched_inference import decode_tiled_2d_embeddings + + return decode_tiled_2d_embeddings( + model=self._model, + image_embeddings=image_embeddings, + device=self._device, + pbar_init=pbar_init, + pbar_update=pbar_update, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) @torch.no_grad() - def _run_decoder_tiled_3d(self, image_embeddings, pbar_init=None, pbar_update=None, z_block=None, z_halo=None): - """Run the UniSAM2 decoder on precomputed tiled 3d embeddings and stitch the tiles in-plane. - - Like `_run_decoder_tiled_2d`, but each tile holds the per-slice features for a full - ``(Z, tile_y, tile_x)`` sub-volume, so the decoder is run in 3d per tile (itself z-chunked) - and the inner block is stitched in-plane. Returns predictions of shape ``(4, Z, Y, X)``. - """ - feats_group = image_embeddings["features"] - shape = tuple(int(s) for s in feats_group.attrs["shape"]) # (Z, Y, X) - tile_shape = tuple(int(s) for s in feats_group.attrs["tile_shape"]) # (y, x) - halo = tuple(int(s) for s in feats_group.attrs["halo"]) # (y, x) - tiling = Blocking([0, 0], list(shape[1:]), list(tile_shape)) # in-plane only; z is full per tile - - # Progress is reported in (tiles x slices) units: each tile decodes the full z stack (in z blocks), - # so the inner per-tile decode advances the shared bar per block (no inner 'pbar_init' that would - # reset the total). This matches the granularity of the 3d-tiled embedding bar. - if pbar_init is not None: - pbar_init(tiling.number_of_blocks * shape[0], "Automatic segmentation (tiles)") - - output = np.zeros((4, *shape), dtype="float32") - for tile_id in range(tiling.number_of_blocks): - # Keep the tile-column features lazy; '_run_decoder_3d' reads them one z block at a time from - # disk rather than materialising the whole column. - tile_features = feats_group[str(tile_id)] # (Z, C, h, w) - tile_embeddings = { - "features": tile_features, "original_size": tile_features.attrs["original_size"], - } - tile_prediction = self._run_decoder_3d( - tile_embeddings, z_block=z_block, z_halo=z_halo, pbar_update=pbar_update, - ) # (4, Z, ty, tx) - - block = tiling.get_block_with_halo(tile_id, halo=list(halo)) - local_bb = tuple(slice(b, e) for b, e in zip(block.inner_block_local.begin, block.inner_block_local.end)) - inner_bb = tuple(slice(b, e) for b, e in zip(block.inner_block.begin, block.inner_block.end)) - # Full channel + full z, inner block in-plane. - output[(slice(None), slice(None)) + inner_bb] = tile_prediction[(slice(None), slice(None)) + local_bb] - - return output + def _run_decoder_tiled_3d( + self, image_embeddings, pbar_init=None, pbar_update=None, z_block=None, z_halo=None, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, + ): + """Decode batches spanning tile columns and z blocks, then stitch them.""" + from micro_sam.v2.batched_inference import decode_tiled_3d_embeddings + + return decode_tiled_3d_embeddings( + model=self._model, + image_embeddings=image_embeddings, + device=self._device, + pbar_init=pbar_init, + pbar_update=pbar_update, + z_block=z_block, + z_halo=z_halo, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) @torch.no_grad() - def _run_decoder_tiled_3d_slice(self, image_embeddings, i, pbar_init=None, pbar_update=None): - """Run the UniSAM2 decoder on a single slice of precomputed tiled 3d embeddings, stitched in-plane. - - For segmenting one slice of a volume without re-encoding: per tile, slice ``i``'s precomputed - features are decoded as a single-slice (Z=1) volume via `_run_decoder_3d`, and the inner block - is stitched in-plane. Returns the slice predictions of shape ``(4, Y, X)``. - """ - feats_group = image_embeddings["features"] - shape = tuple(int(s) for s in feats_group.attrs["shape"]) # (Z, Y, X) - tile_shape = tuple(int(s) for s in feats_group.attrs["tile_shape"]) # (y, x) - halo = tuple(int(s) for s in feats_group.attrs["halo"]) # (y, x) - tiling = Blocking([0, 0], list(shape[1:]), list(tile_shape)) - - if pbar_init is not None: - pbar_init(tiling.number_of_blocks, "Automatic segmentation (tiles)") - - output = np.zeros((4, *shape[1:]), dtype="float32") # (4, Y, X) - for tile_id in range(tiling.number_of_blocks): - tile_features = feats_group[str(tile_id)] # (Z, 1, C, h, w) - # Slice 'i' as a single-slice (Z=1) feature volume for this tile. - slice_embeddings = { - "features": np.asarray(tile_features[i]), "original_size": tile_features.attrs["original_size"], - } - tile_prediction = self._run_decoder_3d(slice_embeddings) # (4, 1, ty, tx) - tile_prediction = tile_prediction[:, 0] # (4, ty, tx) - - block = tiling.get_block_with_halo(tile_id, halo=list(halo)) - local_bb = tuple(slice(b, e) for b, e in zip(block.inner_block_local.begin, block.inner_block_local.end)) - inner_bb = tuple(slice(b, e) for b, e in zip(block.inner_block.begin, block.inner_block.end)) - output[(slice(None),) + inner_bb] = tile_prediction[(slice(None),) + local_bb] - if pbar_update is not None: - pbar_update(1) - - return output + def _run_decoder_tiled_3d_slice( + self, image_embeddings, i, pbar_init=None, pbar_update=None, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, + ): + """Decode one volume slice across all embedding tiles in batches.""" + from micro_sam.v2.batched_inference import decode_tiled_3d_slice + + return decode_tiled_3d_slice( + model=self._model, + image_embeddings=image_embeddings, + index=i, + device=self._device, + pbar_init=pbar_init, + pbar_update=pbar_update, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + ) @torch.no_grad() def initialize( @@ -1278,44 +1248,46 @@ def initialize( pbar_update: Optional[callable] = None, z_block: Optional[int] = None, z_halo: Optional[int] = None, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> None: - """Run the tiled UniSAM2 inference and store the foreground and distance predictions. + """Run tiled UniSAM2 inference and store foreground and distance predictions. - Args: - image: The input image, shape (Y, X) for 2d or (Z, Y, X) for 3d. - ndim: The number of spatial dimensions (2 or 3). - tile_shape: The tile shape for tiled prediction - (y, x) for 2d, (z, y, x) for 3d. - halo: The halo for the overlap between tiles - (y, x) for 2d, (z, y, x) for 3d. - image_embeddings: Optional precomputed tiled image embeddings. If given, the decoder is run - per tile on them and stitched (no encoder pass), reusing the embeddings shared with - interactive / AMG - for both 2d and 3d. See `precompute_image_embeddings`. - i: Slice index, for segmenting a single slice of a volume from tiled 3d embeddings - (ndim 2); unused otherwise. - pbar_init: Callback to initialize an external progress bar, called with the number of - tiles (2d) or the number of blocks / z chunks (3d). - pbar_update: Callback to update an external progress bar, called once per tile / block. - z_block: Number of slices per z block for the per-tile 3d decoder (defaults to `DEFAULT_TILE_Z`). - z_halo: Overlapping slices between z blocks (defaults to `DEFAULT_HALO_Z`). + `batch_size=None` probes the largest safe batch independently on each selected CUDA device. + Reads and preprocessing are queued through `num_prefetch_workers`; full inference also overlaps + output writes through `num_write_workers`. """ - # A single slice of a volume from tiled 3d (video-style) embeddings: reuse slice 'i's features - # per tile (no re-encode). 3d tiled embeddings carry an 'fpn' group; 2d tiled ones do not. + decoder_kwargs = { + "pbar_init": pbar_init, + "pbar_update": pbar_update, + "batch_size": batch_size, + "devices": devices, + "num_prefetch_workers": num_prefetch_workers, + } if image_embeddings is not None and ndim == 2 and "fpn" in image_embeddings and i is not None: self._prediction = self._run_decoder_tiled_3d_slice( - image_embeddings, i, pbar_init=pbar_init, pbar_update=pbar_update, + image_embeddings, i, **decoder_kwargs, ) elif image_embeddings is not None and ndim == 3: self._prediction = self._run_decoder_tiled_3d( - image_embeddings, pbar_init=pbar_init, pbar_update=pbar_update, z_block=z_block, z_halo=z_halo, + image_embeddings, z_block=z_block, z_halo=z_halo, **decoder_kwargs, ) elif image_embeddings is not None and ndim == 2: - self._prediction = self._run_decoder_tiled_2d( - image_embeddings, pbar_init=pbar_init, pbar_update=pbar_update, - ) + self._prediction = self._run_decoder_tiled_2d(image_embeddings, **decoder_kwargs) else: - # 3d tiled inference runs through 'predict_with_halo'; its per-block tqdm is bridged to the - # external progress bar inside '_run_full_inference' so progress is reported per block. self._prediction = self._run_full_inference( - image, ndim, tile_shape=tile_shape, halo=halo, pbar_init=pbar_init, pbar_update=pbar_update, + image, + ndim, + tile_shape=tile_shape, + halo=halo, + pbar_init=pbar_init, + pbar_update=pbar_update, + batch_size=batch_size, + devices=devices, + num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) self._is_initialized = True diff --git a/micro_sam/v2/prompt_based_segmentation.py b/micro_sam/v2/prompt_based_segmentation.py index 5564d9cfa..6be14f2cc 100644 --- a/micro_sam/v2/prompt_based_segmentation.py +++ b/micro_sam/v2/prompt_based_segmentation.py @@ -1,12 +1,16 @@ import gc import ctypes import platform +import threading -from typing import Optional, Union, List +from concurrent import futures +from copy import copy +from typing import Callable, List, Optional, Tuple, Union import numpy as np import torch +from micro_sam.v2.util import Devices from micro_sam.v1.prompt_based_segmentation import _process_box, _compute_logits_from_mask from micro_sam.v2.transforms.resize import resize_longest_side_and_pad_spatial_numpy, ResizeLongestSideTransforms @@ -102,6 +106,35 @@ def _crop_mask_to_tile(tiling, halo, tile_id, mask): return np.asarray(mask)[y0:y1, x0:x1] +def _clone_image_predictor(predictor, model: torch.nn.Module): + """Copy an image predictor onto a replica model, with per-image state cleared.""" + from micro_sam.v2.util import configure_image_predictor + + replica = copy(predictor) + replica.__dict__.pop("_micro_sam_predictor_devices", None) # do not share the source cache + replica.model = model + replica.reset_predictor() + return configure_image_predictor(replica) + + +def _get_image_predictor_devices(predictor, devices: Devices) -> List[Tuple]: + """Create and cache one image-predictor replica per selected device.""" + from micro_sam.v2.batched_inference import prepare_models, resolve_devices + + resolved_devices = resolve_devices(predictor.model, devices) + cache_key = tuple(str(device) for device in resolved_devices) + cached = getattr(predictor, "_micro_sam_predictor_devices", None) + if cached is not None and cached[0] == cache_key: + return cached[1] + + predictor_devices = [ + (predictor if model is predictor.model else _clone_image_predictor(predictor, model), device) + for model, device in prepare_models(predictor.model, resolved_devices) + ] + predictor._micro_sam_predictor_devices = (cache_key, predictor_devices) + return predictor_devices + + def promptable_segmentation_2d( predictor, image: Optional[np.ndarray] = None, @@ -295,6 +328,7 @@ def tiled_promptable_segmentation_2d( boxes: Optional[np.ndarray] = None, masks: Optional[np.ndarray] = None, batched: Optional[bool] = None, + devices: Devices = None, ): """Tiled 2d promptable segmentation for the SAM2 image predictor. @@ -302,6 +336,9 @@ def tiled_promptable_segmentation_2d( the predictor (via `set_precomputed`), runs `promptable_segmentation_2d` on the tile, and stitches the per-tile mask into the full image. Points are in (y, x) order, as passed by the annotator. Same return convention as `promptable_segmentation_2d`. + + Independent active tiles run concurrently on persistent predictor replicas. By default all + visible CUDA devices are used; pass `devices` to select or restrict them. """ from bioimage_cpp.utils import Blocking from micro_sam.v2.util import set_precomputed @@ -336,8 +373,7 @@ def tiled_promptable_segmentation_2d( None if box_mask is None else _crop_mask_to_tile(tiling, halo, tid, box_mask) ) - out = np.zeros(shape, dtype="uint32") - found = False + tile_jobs = {} for tile_id in sorted(set(tile_points) | set(tile_boxes)): tpoints = np.asarray(tile_points.get(tile_id, [])).reshape(-1, 2) tlabels = np.asarray(tile_labels.get(tile_id, []), dtype=int) @@ -351,14 +387,40 @@ def tiled_promptable_segmentation_2d( y0, x0 = int(outer.begin[0]), int(outer.begin[1]) local_points = (tpoints - np.array([y0, x0])) if len(tpoints) else tpoints local_boxes = [b - np.array([y0, x0, y0, x0]) for b in tboxes] if tboxes else None - local_masks = tile_masks.get(tile_id) if tboxes else None + tile_jobs[tile_id] = local_points, tlabels, local_boxes, local_masks + + predictor_devices = _get_image_predictor_devices(predictor, devices) + groups = {} + for tile_id in tile_jobs: + worker_id = tile_id % len(predictor_devices) + groups.setdefault(worker_id, []).append(tile_id) + + def run_group(worker_id, tile_ids): + local_predictor, _ = predictor_devices[worker_id] + results = [] + for tile_id in tile_ids: + local_points, local_labels, local_boxes, local_masks = tile_jobs[tile_id] + set_precomputed(local_predictor, image_embeddings, tile_id=tile_id) + tile_seg = promptable_segmentation_2d( + local_predictor, image=None, points=local_points, labels=local_labels, + boxes=local_boxes, masks=local_masks, batched=batched, + ) + results.append((tile_id, tile_seg)) + return results - set_precomputed(predictor, image_embeddings, tile_id=tile_id) - tile_seg = promptable_segmentation_2d( - predictor, image=None, points=local_points, labels=tlabels, boxes=local_boxes, - masks=local_masks, batched=batched, - ) + if len(groups) < 2: + results = run_group(*next(iter(groups.items()))) if groups else [] + else: + results = [] + with futures.ThreadPoolExecutor(max_workers=len(groups)) as pool: + tasks = [pool.submit(run_group, worker_id, tile_ids) for worker_id, tile_ids in groups.items()] + for task in tasks: + results.extend(task.result()) + + out = np.zeros(shape, dtype="uint32") + found = False + for tile_id, tile_seg in sorted(results): if tile_seg is None: continue @@ -944,10 +1006,16 @@ class TiledPromptableSegmentation3D: volume_embeddings: The precomputed tiled 3d embeddings (with per-tile `features`/`pos_enc`/ `fpn` groups and `shape`/`tile_shape`/`halo` attrs). See `precompute_image_embeddings`. device: The device to run inference on. + devices: Devices used for independent tile columns. By default all visible CUDA devices + are used when the predictor is on CUDA. """ - def __init__(self, predictor, volume, volume_embeddings, device=None, **kwargs): + def __init__(self, predictor, volume, volume_embeddings, device=None, devices: Devices = None, **kwargs): from bioimage_cpp.utils import Blocking + from micro_sam.v2.batched_inference import prepare_models, resolve_devices + + resolved_devices = resolve_devices(predictor, devices) + self._predictor_devices = prepare_models(predictor, resolved_devices) self.predictor = predictor self.volume = volume @@ -1013,8 +1081,9 @@ def _get_segmenter(self, tile_id): "input_size": tile_dataset.attrs["input_size"], "original_size": tile_dataset.attrs["original_size"], } + predictor, device = self._predictor_devices[tile_id % len(self._predictor_devices)] self._segmenters[tile_id] = PromptableSegmentation3D( - self.predictor, sub_volume, tile_embeddings, device=self.device, **self._kwargs + predictor, sub_volume, tile_embeddings, device=device, **self._kwargs ) return self._segmenters[tile_id] @@ -1025,6 +1094,26 @@ def _inner_slices(self, tile_id): glob = tuple(slice(b, e) for b, e in zip(block.inner_block.begin, block.inner_block.end)) return local, glob + def _run_tile_jobs(self, tile_ids: List[int], function: Callable) -> List[Tuple]: + """Run tile jobs concurrently across devices and serially within each device.""" + groups = {} + for tile_id in tile_ids: + worker_id = tile_id % len(self._predictor_devices) + groups.setdefault(worker_id, []).append(tile_id) + + def run_group(group): + return [(tile_id, function(tile_id)) for tile_id in group] + + if len(groups) < 2: + return run_group(next(iter(groups.values()))) if groups else [] + + results = [] + with futures.ThreadPoolExecutor(max_workers=len(groups)) as pool: + tasks = [pool.submit(run_group, group) for group in groups.values()] + for task in tasks: + results.extend(task.result()) + return sorted(results, key=lambda result: result[0]) + def segment_slice(self, frame_idx, points=None, labels=None, boxes=None, masks=None, object_id=1): """Segment a single slice. Points are (x, y), boxes (x0, y0, x1, y1), as passed by the annotator. @@ -1054,8 +1143,7 @@ def segment_slice(self, frame_idx, points=None, labels=None, boxes=None, masks=N None if box_mask is None else _crop_mask_to_tile(self.tiling, self.halo, tid, box_mask) ) - out = np.zeros(self.shape[1:], dtype="uint32") - found = False + tile_jobs = {} for tile_id in sorted(set(tile_points) | set(tile_boxes)): tpoints = np.asarray(tile_points.get(tile_id, [])).reshape(-1, 2) tlabels = np.asarray(tile_labels.get(tile_id, []), dtype=int) @@ -1066,10 +1154,23 @@ def segment_slice(self, frame_idx, points=None, labels=None, boxes=None, masks=N local_points = (tpoints - np.array([x0, y0])) if len(tpoints) else None local_boxes = [b - np.array([x0, y0, x0, y0]) for b in tboxes] if tboxes else None local_masks = tile_masks.get(tile_id) if tboxes else None - tile_seg = self._get_segmenter(tile_id).segment_slice( - frame_idx, points=local_points, labels=(tlabels if len(tlabels) else None), + tile_jobs[tile_id] = ( + local_points, + tlabels if len(tlabels) else None, + local_boxes, + local_masks, + ) + + def segment_tile(tile_id): + local_points, local_labels, local_boxes, local_masks = tile_jobs[tile_id] + return self._get_segmenter(tile_id).segment_slice( + frame_idx, points=local_points, labels=local_labels, boxes=local_boxes, masks=local_masks, object_id=object_id, ) + + out = np.zeros(self.shape[1:], dtype="uint32") + found = False + for tile_id, tile_seg in self._run_tile_jobs(sorted(tile_jobs), segment_tile): if tile_seg is None: continue local, glob = self._inner_slices(tile_id) @@ -1155,10 +1256,20 @@ def predict(self, update_progress=None, early_stop_patience=None, z_range=None): prompted in several tiles keeps one id and is merged across the tile boundaries. """ segmentation = np.zeros(self.shape, dtype="uint64") - for tile_id in sorted(self._segmenters): - tile_seg = self._segmenters[tile_id].predict( - update_progress=update_progress, early_stop_patience=early_stop_patience, z_range=z_range, + progress_lock = threading.Lock() + + def locked_update(value): + if update_progress is None: + return + with progress_lock: + update_progress(value) + + def predict_tile(tile_id): + return self._segmenters[tile_id].predict( + update_progress=locked_update, early_stop_patience=early_stop_patience, z_range=z_range, ) + + for tile_id, tile_seg in self._run_tile_jobs(sorted(self._segmenters), predict_tile): local, glob = self._inner_slices(tile_id) inner = tile_seg[(slice(None),) + local] region = segmentation[(slice(None),) + glob] diff --git a/micro_sam/v2/util.py b/micro_sam/v2/util.py index b3e402263..a767849a8 100644 --- a/micro_sam/v2/util.py +++ b/micro_sam/v2/util.py @@ -3,14 +3,14 @@ import pooch import warnings from pathlib import Path -from typing import Union, Literal, Optional, Tuple +from typing import Union, Literal, Optional, Sequence, Tuple import numpy as np import torch from micro_sam.util import ( - get_device, get_cache_directory, microsam_cachedir, _open_embeddings, _create_dataset_without_data, + get_device, get_cache_directory, microsam_cachedir, _open_embeddings, _configure_mps_memory, ) from micro_sam.v2.models._video_predictor import _build_sam2_video_predictor @@ -20,6 +20,10 @@ from sam2.build_sam import build_sam2 +Device = Optional[Union[str, torch.device]] +Devices = Optional[Union[str, torch.device, Sequence[Union[str, torch.device]]]] + + # NOTE: The model config is expected to be fetched from the module's relative path location. sys.path.append(str(Path(sam2.__file__).parents[0])) @@ -395,172 +399,6 @@ def _compute_2d(input_, predictor, f, save_path, pbar_init, pbar_update): return image_embeddings -def _compute_tiled_2d(input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update): - from micro_sam.util import _create_dataset_with_data - from bioimage_cpp.utils import Blocking - - features = f.require_group("features") - high_res_group = f.require_group("high_res_feats") - - # If the tiled embeddings are already cached we just return the open groups. - if save_path is not None and "shape" in features.attrs: - return {"features": features, "high_res_feats": high_res_group, "input_size": None, "original_size": None} - - tiling = Blocking([0, 0], list(input_.shape[:2]), list(tile_shape)) - n_tiles = tiling.number_of_blocks - - features.attrs["shape"] = list(input_.shape[:2]) - features.attrs["tile_shape"] = list(tile_shape) - features.attrs["halo"] = list(halo) - - pbar_init(n_tiles, "Compute Image Embeddings 2D tiled") - predictor.reset_predictor() - for tile_id in range(n_tiles): - block = tiling.get_block_with_halo(tile_id, list(halo)).outer_block - bb = tuple(slice(begin, end) for begin, end in zip(block.begin, block.end)) - predictor.set_image(to_image(input_[bb])) - - tile_features = predictor.get_image_embedding().cpu().numpy() - high_res_features = [feat.cpu().numpy() for feat in predictor._features["high_res_feats"]] - ds = _create_dataset_with_data(features, str(tile_id), data=tile_features) - ds.attrs["input_size"] = predictor.model.image_size - # Store the original size in the predictor's nested '[[h, w]]' layout (one entry per image), - # so 'set_precomputed' restores '_orig_hw[-1] == (h, w)' for non-square tiles too. - ds.attrs["original_size"] = [list(predictor._orig_hw[0])] - - tile_high_res = high_res_group.require_group(str(tile_id)) - for level, feat in enumerate(high_res_features): - _create_dataset_with_data(tile_high_res, str(level), data=feat) - - predictor.reset_predictor() - pbar_update(1) - - if save_path is not None: - _write_embedding_signature( - f, input_, predictor, tile_shape=tile_shape, halo=halo, input_size=None, original_size=None, - ) - - return {"features": features, "high_res_feats": high_res_group, "input_size": None, "original_size": None} - - -def _compute_tiled_3d(input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update): - from bioimage_cpp.utils import Blocking - - features = f.require_group("features") - pos_enc_group = f.require_group("pos_enc") - fpn_group = f.require_group("fpn") - - # If the tiled embeddings are already cached we just return the open groups. - if save_path is not None and "shape" in features.attrs: - return { - "features": features, "pos_enc": pos_enc_group, "fpn": fpn_group, - "input_size": None, "original_size": None, - } - - # The volume is tiled in-plane (xy); each tile is its own (Z, tile_y, tile_x) sub-volume. - tiling = Blocking([0, 0], list(input_.shape[1:]), list(tile_shape)) - n_tiles = tiling.number_of_blocks - n_slices = input_.shape[0] - - features.attrs["shape"] = list(input_.shape) - features.attrs["tile_shape"] = list(tile_shape) - features.attrs["halo"] = list(halo) - - # Progress is reported per actual patch (tile-column x z slice), not per tile, since each tile - # encodes all z slices and that inner loop is the bulk of the work. - pbar_init(n_tiles * n_slices, "Compute Image Embeddings 3D tiled") - for tile_id in range(n_tiles): - block = tiling.get_block_with_halo(tile_id, list(halo)).outer_block - bb = tuple(slice(begin, end) for begin, end in zip(block.begin, block.end)) - sub_volume = np.asarray(input_[:, bb[0], bb[1]]) - - # Compute the per-slice video-style features for this tile-column (as in '_compute_3d'). - inference_state = predictor.init_state( - volume=sub_volume, volume_embeddings=None, device=predictor.device, ignore_caching_features=True, - ) - - # Stream one slice at a time into per-tile datasets, moving each slice's features off-device - # right away, so peak memory is a single slice rather than the whole tile-column. Buffering all - # slices (and then stacking them) held ~200 MB/slice of high-res features on-device and OOMed on - # deep volumes even when saving to disk. Datasets are created lazily from the first slice's - # shapes with per-slice '(1, ...)' chunking, matching '_compute_3d' and the on-disk layout the - # lazy tiled consumer expects (features/, pos_enc//, fpn//). - feature_ds, pos_enc_dsets, fpn_dsets = None, None, None - tile_pos_group = pos_enc_group.require_group(str(tile_id)) - tile_fpn_group = fpn_group.require_group(str(tile_id)) - input_size, original_size = None, None - for z in range(n_slices): - vision_feats, pos_encs, fpns, original_sizes, input_sizes = _compute_embeddings_batched_3d( - inference_state, predictor, [z], [to_image(sub_volume[z])], pbar_update=pbar_update, - ) - curr_feat = vision_feats[0] - if curr_feat.ndim == 3: - curr_feat = curr_feat.unsqueeze(0) - curr_pos, curr_fpn = pos_encs[0], fpns[0] - - if feature_ds is None: - feature_ds = _create_dataset_without_data( - features, str(tile_id), shape=(n_slices,) + tuple(curr_feat.shape), - dtype="float32", chunks=(1,) + tuple(curr_feat.shape), - ) - pos_enc_dsets = [ - _create_dataset_without_data( - tile_pos_group, str(level), shape=(n_slices,) + tuple(t.shape), - dtype="float32", chunks=(1,) + tuple(t.shape), - ) for level, t in enumerate(curr_pos) - ] - fpn_dsets = [ - _create_dataset_without_data( - tile_fpn_group, str(level), shape=(n_slices,) + tuple(t.shape), - dtype="float32", chunks=(1,) + tuple(t.shape), - ) for level, t in enumerate(curr_fpn) - ] - - feature_ds[z] = curr_feat.detach().cpu().numpy() - for level, t in enumerate(curr_pos): - pos_enc_dsets[level][z] = t.detach().cpu().numpy() - for level, t in enumerate(curr_fpn): - fpn_dsets[level][z] = t.detach().cpu().numpy() - input_size, original_size = input_sizes[-1], original_sizes[-1] - - feature_ds.attrs["input_size"] = input_size - feature_ds.attrs["original_size"] = list(original_size) - - if save_path is not None: - _write_embedding_signature( - f, input_, predictor, tile_shape=tile_shape, halo=halo, input_size=None, original_size=None, - ) - - return { - "features": features, "pos_enc": pos_enc_group, "fpn": fpn_group, - "input_size": None, "original_size": None, - } - - -def _create_list_dataset_without_data(group, prefix_name, tensors, dtype, z_slices): - subgroup = group.require_group(prefix_name) - - ds_list = [] - for i, curr_tensor in enumerate(tensors): - curr_shape = tuple(curr_tensor.shape) - shape = (z_slices,) + curr_shape - chunks = (1,) + curr_shape - name = str(i) - - if name in subgroup: - ds = subgroup[name] - if ds.shape != shape: - raise RuntimeError(f"Invalid shape for {prefix_name}/{name}: expected {shape}, got {ds.shape}") - if getattr(ds, "chunks", None) is not None and ds.chunks != chunks: - raise RuntimeError(f"Invalid chunks for {prefix_name}/{name}: expected {chunks}, got {ds.chunks}") - else: - ds = _create_dataset_without_data(subgroup, name, shape=shape, dtype=dtype, chunks=chunks) - - ds_list.append(ds) - - return ds_list - - def _load_list_datasets(group, prefix_name, lazy_loading): if prefix_name not in group: return [] @@ -575,159 +413,6 @@ def _load_list_datasets(group, prefix_name, lazy_loading): return out -@torch.no_grad -def _compute_embeddings_batched_3d(inference_state, predictor, batched_z, batched_images, pbar_update=None): - batched_vision_features, batched_pos_enc, batched_backbone_fpn, original_sizes, input_sizes = [], [], [], [], [] - - for image, z_id in zip(batched_images, batched_z): - # Run the image encoder to extract relevant features - predictor._get_image_feature(inference_state, frame_idx=z_id, batch_size=1) - - # Let's extract the current 'cached_features' outputs - _, curr_backbone_out = inference_state["cached_features"][z_id] - - # Store the vision transformer outputs and other stuff. - batched_vision_features.append(curr_backbone_out["vision_features"]) - batched_pos_enc.append(curr_backbone_out["vision_pos_enc"]) - batched_backbone_fpn.append(curr_backbone_out["backbone_fpn"]) - original_sizes.append(image.shape[:2]) - input_sizes.append(predictor.image_size) - - if pbar_update is not None: # Advance per slice so a tile-column reports per-patch progress. - pbar_update(1) - - return batched_vision_features, batched_pos_enc, batched_backbone_fpn, original_sizes, input_sizes - - -def _compute_3d(input_, predictor, f, save_path, lazy_loading, pbar_init, pbar_update, batch_size): - # Check if the embeddings are already fully cached. - if save_path is not None and "original_size" in f.attrs: - # In this case we load the embeddings. - features = f["features"] if lazy_loading else f["features"][:] - pos_enc = _load_list_datasets(f, "pos_enc", lazy_loading) - fpn = _load_list_datasets(f, "fpn", lazy_loading) - original_size = f.attrs["original_size"] - input_size = f.attrs["input_size"] - image_embeddings = { - "features": features, - "pos_enc": pos_enc, - "fpn": fpn, - "input_size": input_size, - "original_size": original_size, - } - return image_embeddings - - # Otherwise we have to compute the embeddings. - - # First check if we have a save path or not and set things up accordingly. - if save_path is None: - features, pos_encs, fpns = [], [], [] - save_features = False - partial_features = False - else: - save_features = True - embed_shape = (1, 256, 64, 64) - shape = (input_.shape[0],) + embed_shape - chunks = (1,) + embed_shape - if "features" in f: - partial_features = True - features = f["features"] - if features.shape != shape or features.chunks != chunks: - raise RuntimeError("Invalid partial features") - else: - partial_features = False - from micro_sam.util import _create_dataset_without_data - features = _create_dataset_without_data(f, "features", shape=shape, chunks=chunks, dtype="float32") - - # We create the 'inference_state' object which keeps all important components in memory. - # Pass the predictor's device so encoder inputs match the model when it is not the default device. - inference_state = predictor.init_state( - volume=input_, - volume_embeddings=None, # NOTE: It's a mandatory argument, but with the argument below, passing 'None' doesn't matter. # noqa - device=predictor.device, - ignore_caching_features=True - ) - - # Initialize the pbar and batches. - n_slices = input_.shape[0] - pbar_init(n_slices, "Compute Image Embeddings 3D") - n_batches = int(np.ceil(n_slices / batch_size)) - pos_enc_dsets, fpn_dsets = None, None - - for batch_id in range(n_batches): - z_start = batch_id * batch_size - z_stop = min(z_start + batch_size, n_slices) - - batched_images, batched_z = [], [] - for z in range(z_start, z_stop): - # Skip feature computation in case of partial features in non-zero slice. - if partial_features and np.count_nonzero(features[z]) != 0: - continue - - tile_input = to_image(input_[z]) - batched_images.append(tile_input) - batched_z.append(z) - - ( - batched_vision_features, batched_pos_enc, batched_backbone_fpn, original_sizes, input_sizes - ) = _compute_embeddings_batched_3d(inference_state, predictor, batched_z, batched_images) - - for z, curr_vision_feats, curr_pos_enc, curr_back_fpn in zip( - batched_z, batched_vision_features, batched_pos_enc, batched_backbone_fpn - ): - if curr_vision_feats.ndim == 3: - curr_vision_feats = curr_vision_feats.unsqueeze(0) - - if save_features: - features[z] = curr_vision_feats.detach().cpu().numpy() - if pos_enc_dsets is None: - pos_enc_dsets = _create_list_dataset_without_data( - f, "pos_enc", curr_pos_enc, dtype="float32", z_slices=n_slices - ) - for i, t in enumerate(curr_pos_enc): - pos_enc_dsets[i][z] = t.detach().cpu().numpy() - - if fpn_dsets is None: - fpn_dsets = _create_list_dataset_without_data( - f, "fpn", curr_back_fpn, dtype="float32", z_slices=n_slices - ) - for i, t in enumerate(curr_back_fpn): - fpn_dsets[i][z] = t.detach().cpu().numpy() - - else: - features.append(curr_vision_feats) - pos_encs.append(curr_pos_enc) - fpns.append(curr_back_fpn) - - pbar_update(1) - - if save_features: - _write_embedding_signature( - f, input_, predictor, tile_shape=None, halo=None, - input_size=input_sizes[-1], original_size=original_sizes[-1], - ) - else: - # Concatenate across the z axis for 'vision_features'. - features = torch.cat(features).cpu().numpy() - - # Concatenate across the z axis for other features too. - depth = 3 # Corresponds to the depth of both FPN and Positional Embeddings. - pos_encs = [torch.stack([p[i] for p in pos_encs]) for i in range(depth)] - fpns = [torch.stack([p[i] for p in fpns]) for i in range(depth)] - - pos_enc = _load_list_datasets(f, "pos_enc", lazy_loading) if save_features else pos_encs - fpn = _load_list_datasets(f, "fpn", lazy_loading) if save_features else fpns - - image_embeddings = { - "features": features, - "pos_enc": pos_enc, - "fpn": fpn, - "input_size": input_sizes[-1], - "original_size": original_sizes[-1], - } - return image_embeddings - - def precompute_image_embeddings( predictor, input_: np.ndarray, @@ -737,7 +422,9 @@ def precompute_image_embeddings( tile_shape: Optional[Tuple[int, int]] = None, halo: Optional[Tuple[int, int]] = None, verbose: bool = True, - batch_size: int = 1, + batch_size: Optional[int] = None, + devices: Devices = None, + num_prefetch_workers: int = 4, pbar_init: Optional[callable] = None, pbar_update: Optional[callable] = None, ): @@ -746,7 +433,21 @@ def precompute_image_embeddings( If 'save_path' is given the embeddings will be loaded/saved in a zarr container. Args: - ... + predictor: The SAM2 image or video predictor. + input_: The input image or volume. + save_path: Optional zarr path for persistent embedding storage. + lazy_loading: Return disk-backed arrays for saved 3D embeddings. + ndim: The number of spatial dimensions. By default this is inferred from the input. + tile_shape: Optional in-plane tile shape. + halo: Optional in-plane tile halo. + verbose: Whether to show progress. + batch_size: Number of independent tiles or slices per encoder call. If None, probe the + largest safe value independently on each CUDA device. + devices: Device or devices used for embedding inference. If None and the predictor is on + CUDA, all visible CUDA devices are used. + num_prefetch_workers: Number of threads used to read and preprocess input jobs. + pbar_init: Optional callback to initialize external progress. + pbar_update: Optional callback to update external progress. Returns: The image embeddings. @@ -779,6 +480,8 @@ def precompute_image_embeddings( f.attrs["normalization"] = RAW_NORMALIZATION from micro_sam.util import handle_pbar + from micro_sam.v2.batched_inference import compute_3d, compute_tiled_2d, compute_tiled_3d + _, pbar_init, pbar_update, pbar_close = handle_pbar(verbose, pbar_init, pbar_update) if ndim == 2 and tile_shape is None: @@ -786,13 +489,22 @@ def precompute_image_embeddings( elif ndim == 2 and tile_shape is not None: if halo is None: raise ValueError("To compute tiled embeddings the parameter halo has to be passed.") - embeddings = _compute_tiled_2d(input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update) + embeddings = compute_tiled_2d( + input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update, + batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + ) elif ndim == 3 and tile_shape is None: - embeddings = _compute_3d(input_, predictor, f, save_path, lazy_loading, pbar_init, pbar_update, batch_size) + embeddings = compute_3d( + input_, predictor, f, save_path, lazy_loading, pbar_init, pbar_update, + batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + ) elif ndim == 3 and tile_shape is not None: if halo is None: raise ValueError("To compute tiled embeddings the parameter halo has to be passed.") - embeddings = _compute_tiled_3d(input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update) + embeddings = compute_tiled_3d( + input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update, + batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + ) else: raise ValueError(f"Invalid dimensionality {input_.ndim}, expect 2 or 3 dim data.") diff --git a/setup.cfg b/setup.cfg index 60086002a..b853468c0 100644 --- a/setup.cfg +++ b/setup.cfg @@ -61,7 +61,7 @@ install_requires = superqt timm torch>=2.5 - torch-em>=0.9 + torch-em>=0.10.1 torchvision tqdm trackastra>=0.5.3 diff --git a/test/test_v2_batched_inference.py b/test/test_v2_batched_inference.py new file mode 100644 index 000000000..18430054b --- /dev/null +++ b/test/test_v2_batched_inference.py @@ -0,0 +1,170 @@ +import unittest + +import numpy as np +import torch +import torch.nn.functional as F +from bioimage_cpp.utils import Blocking + +from micro_sam.v2.batched_inference import ( + decode_tiled_3d_embeddings, decode_volume_embeddings, run_batched_pipeline, +) +from micro_sam.v2.prompt_based_segmentation import TiledPromptableSegmentation3D + + +class ArrayWithAttrs: + def __init__(self, data, **attrs): + self.data = np.asarray(data) + self.attrs = attrs + self.shape = self.data.shape + + def __array__(self, dtype=None): + return np.asarray(self.data, dtype=dtype) + + def __getitem__(self, index): + return self.data[index] + + +class Group(dict): + def __init__(self, *args, **attrs): + super().__init__(*args) + self.attrs = attrs + + +class FakeEncoder(torch.nn.Module): + def __init__(self, img_size=8): + super().__init__() + self.img_size = img_size + + +class FakeUNETR(torch.nn.Module): + def __init__(self): + super().__init__() + self.encoder = FakeEncoder() + self.scale = torch.nn.Parameter(torch.ones(1)) + self.batch_sizes = [] + + def forward(self, x): + self.batch_sizes.append(int(x.shape[0])) + features = torch.stack( + [self.encoder(x[:, :, index])[0] for index in range(x.shape[2])], + dim=2, + ) + prediction = features.mean(dim=1, keepdim=True) + prediction = F.interpolate( + prediction, + size=(x.shape[2], *x.shape[-2:]), + mode="trilinear", + align_corners=False, + ) + return prediction.repeat(1, 4, 1, 1, 1) * self.scale + + +class TestBatchedPipeline(unittest.TestCase): + def test_pipeline_batches_and_writes_all_jobs(self): + outputs = {} + progress = [] + + def predict(model, items, device): + self.assertEqual(device, torch.device("cpu")) + return [value + model for value in items] + + run_batched_pipeline( + jobs=range(7), + model_devices=[(3, torch.device("cpu"))], + batch_sizes=[3], + load_fn=lambda value: 2 * value, + predict_fn=predict, + write_fn=outputs.__setitem__, + num_prefetch_workers=2, + update_progress=progress.append, + ) + + self.assertEqual(outputs, {index: 2 * index + 3 for index in range(7)}) + self.assertEqual(sum(progress), 7) + + +class TestBatchedDecoder(unittest.TestCase): + def test_volume_decoder_batches_z_blocks(self): + features = np.arange(8 * 2 * 2 * 2, dtype="float32").reshape(8, 2, 2, 2) + model = FakeUNETR() + + output = decode_volume_embeddings( + model, + {"features": features, "original_size": (8, 8)}, + device="cpu", + z_block=2, + z_halo=0, + batch_size=2, + num_prefetch_workers=2, + ) + + self.assertEqual(output.shape, (4, 8, 8, 8)) + self.assertEqual(model.batch_sizes, [2, 2]) + + def test_tiled_volume_decoder_batches_tiles_and_z_blocks(self): + features = Group( + { + str(tile_id): ArrayWithAttrs( + np.full((4, 1, 2, 2, 2), tile_id + 1, dtype="float32"), + original_size=(4, 4), + ) + for tile_id in range(4) + }, + shape=(4, 8, 8), + tile_shape=(4, 4), + halo=(0, 0), + ) + model = FakeUNETR() + + output = decode_tiled_3d_embeddings( + model, + {"features": features}, + device="cpu", + z_block=2, + z_halo=0, + batch_size=2, + num_prefetch_workers=2, + ) + + self.assertEqual(output.shape, (4, 4, 8, 8)) + self.assertTrue(np.allclose(output[:, :, :4, :4], 1)) + self.assertTrue(np.allclose(output[:, :, :4, 4:], 2)) + self.assertTrue(np.allclose(output[:, :, 4:, :4], 3)) + self.assertTrue(np.allclose(output[:, :, 4:, 4:], 4)) + self.assertEqual(model.batch_sizes, [2, 2, 2, 2]) + + +class FakeInteractiveTile: + def __init__(self, value): + self.value = value + + def predict(self, update_progress=None, **kwargs): + if update_progress is not None: + update_progress(2) + return np.full((2, 4, 4), self.value, dtype="uint64") + + +class TestInteractiveTileScheduling(unittest.TestCase): + def test_active_tile_columns_are_stitched_after_device_jobs(self): + segmenter = TiledPromptableSegmentation3D.__new__(TiledPromptableSegmentation3D) + segmenter.shape = (2, 8, 8) + segmenter.halo = (0, 0) + segmenter.tiling = Blocking([0, 0], [8, 8], [4, 4]) + segmenter._predictor_devices = [(None, torch.device("cpu")), (None, torch.device("cpu"))] + segmenter._segmenters = { + tile_id: FakeInteractiveTile(tile_id + 1) + for tile_id in range(4) + } + progress = [] + + output = segmenter.predict(update_progress=progress.append) + + self.assertTrue(np.all(output[:, :4, :4] == 1)) + self.assertTrue(np.all(output[:, :4, 4:] == 2)) + self.assertTrue(np.all(output[:, 4:, :4] == 3)) + self.assertTrue(np.all(output[:, 4:, 4:] == 4)) + self.assertEqual(sum(progress), 8) + + +if __name__ == "__main__": + unittest.main() From f744db0e5348ffb67ae7eeaf8d99e14f9a8ee999 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 14:29:44 +0200 Subject: [PATCH 02/11] Improve embedding batch sizing and expose annotator controls --- micro_sam/sam_annotator/_state.py | 2 + micro_sam/sam_annotator/_tooltips.py | 6 +- micro_sam/sam_annotator/_widgets.py | 12 ++ micro_sam/v2/automatic_segmentation.py | 5 +- micro_sam/v2/batched_inference.py | 145 +++++++++++++++++++--- micro_sam/v2/instance_segmentation.py | 18 +-- micro_sam/v2/util.py | 7 +- test/test_sam_annotator/test_annotator.py | 10 +- test/test_sam_annotator/test_state.py | 4 +- test/test_v2_batched_inference.py | 65 +++++++++- 10 files changed, 241 insertions(+), 33 deletions(-) diff --git a/micro_sam/sam_annotator/_state.py b/micro_sam/sam_annotator/_state.py index 4014e4cdb..7a02ec2e4 100644 --- a/micro_sam/sam_annotator/_state.py +++ b/micro_sam/sam_annotator/_state.py @@ -160,6 +160,7 @@ def initialize_predictor( decoder_path=None, tile_shape=None, halo=None, + batch_size=1, precompute_autoseg_state=False, prefer_decoder=True, pbar_init=None, @@ -264,6 +265,7 @@ def initialize_predictor( ndim=ndim, tile_shape=tile_shape, halo=halo, + batch_size=batch_size, verbose=True, lazy_loading=lazy_loading, pbar_init=pbar_init, diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py index 21998ed0a..5aae38267 100644 --- a/micro_sam/sam_annotator/_tooltips.py +++ b/micro_sam/sam_annotator/_tooltips.py @@ -2,6 +2,10 @@ tooltips = { "embedding": { + "batch_size": ( + "Number of image slices or tiles encoded together per GPU. Larger values may improve throughput, " + "but can be slower or run out of GPU memory. The safe default is 1." + ), "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.", @@ -15,7 +19,7 @@ "automatic_segmentation_mode": "Select the automatic segmentation mode.", "run_button": "Compute embeddings or load embeddings if embedding_save_path is specified.", "tiling": "Enter tile size for computing tiled embeddings. Enter only x-value for quadratic size or both for non-quadratic.", # noqa - "settings": "Settings for computing the image embeddings: model family and size, tiling, the embedding save path and the compute device.", # noqa + "settings": "Settings for computing the image embeddings: model family and size, tiling, batch size, the embedding save path and the compute device.", # noqa }, "segmentnd": { "box_extension": "Enter factor by which box size is increased when projecting to adjacent slices. Larger factors help if object sizes change between slices.", # noqa diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index d21f3dfd5..00c6b8f5f 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -1788,6 +1788,16 @@ def _create_settings_widget(self): ) setting_values.layout().addLayout(layout) + # Use a conservative default and let users opt into larger per-GPU batches when their + # workload and available memory benefit from them. + self.batch_size = 1 + self.batch_size_param, batch_size_layout = self._add_int_param( + "batch_size", self.batch_size, min_val=1, max_val=64, + title="batch size per GPU:", + tooltip=get_tooltip("embedding", "batch_size"), + ) + setting_values.layout().addLayout(batch_size_layout) + # Create UI for the save path. self.embeddings_save_path = None self.embeddings_save_path_param, save_layout = self._add_path_param( @@ -1922,6 +1932,7 @@ def _reset_inputs_to_defaults(self): self.halo_x_param.setValue(DEFAULT_HALO[0]) self.halo_y_param.setValue(DEFAULT_HALO[1]) self.tiling_dropdown.setCurrentText("no") + self.batch_size_param.setValue(1) self.embeddings_save_path_param.setText("") # Reset the image-dimensionality override back to 'auto' so a new image is re-detected. @@ -2211,6 +2222,7 @@ def pbar_update(update): checkpoint_path=self.custom_weights, tile_shape=tile_shape, halo=halo, + batch_size=self.batch_size, prefer_decoder=True, precompute_autoseg_state=precompute_autoseg_state, pbar_init=pbar_init, diff --git a/micro_sam/v2/automatic_segmentation.py b/micro_sam/v2/automatic_segmentation.py index d3947942f..1a73fc638 100644 --- a/micro_sam/v2/automatic_segmentation.py +++ b/micro_sam/v2/automatic_segmentation.py @@ -97,7 +97,7 @@ def automatic_instance_segmentation( mode: str = "sparse", device: Optional[Union[str, torch.device]] = None, verbose: bool = True, - batch_size: Optional[int] = None, + batch_size: Optional[int] = 1, devices: Devices = None, num_prefetch_workers: int = 4, num_write_workers: int = 1, @@ -122,7 +122,8 @@ def automatic_instance_segmentation( mode: The AIS post-processing mode, 'sparse' (flow) or 'dense' (multicut). Ignored for AMG. device: The device to run inference on. verbose: Whether to print progress. - batch_size: Explicit tile or slice batch size, or None for automatic capacity probing. + batch_size: Explicit tile or slice batch size. Defaults to one; pass None for + throughput-based automatic selection. devices: Inference device or devices. None uses all visible GPUs when the model is on CUDA. num_prefetch_workers: Number of input reading and preprocessing threads. num_write_workers: Number of output writing threads for full tiled inference. diff --git a/micro_sam/v2/batched_inference.py b/micro_sam/v2/batched_inference.py index d698dd9f3..ee593bdea 100644 --- a/micro_sam/v2/batched_inference.py +++ b/micro_sam/v2/batched_inference.py @@ -4,6 +4,7 @@ import os import queue import threading +import time import warnings from collections import defaultdict from copy import deepcopy @@ -113,6 +114,72 @@ def release_model_replicas(model_devices: List[Tuple[torch.nn.Module, torch.devi torch.cuda.empty_cache() +def _is_oom_error(exc: Exception) -> bool: + oom_type = getattr(torch.cuda, "OutOfMemoryError", ()) + return isinstance(exc, oom_type) or "out of memory" in str(exc).lower() + + +def _release_probe_memory(device: torch.device) -> None: + gc.collect() + if device.type == "cuda": + with torch.cuda.device(device): + torch.cuda.empty_cache() + + +def _measure_batch_throughput( + model: torch.nn.Module, + device: torch.device, + batch_size: int, + patch_shape: Tuple[int, ...], + in_channels: int, + prediction_function: Callable[[torch.nn.Module, torch.Tensor], Any], + n_repeats: int = 2, +) -> Optional[float]: + """Measure samples per second for one synthetic batch, returning None on CUDA OOM.""" + inputs = output = None + was_training = model.training + model.eval() + try: + inputs = torch.empty( + (batch_size, in_channels, *patch_shape), dtype=torch.float32, device=device, + ).normal_() + timings = [] + for repeat in range(n_repeats + 1): + if device.type == "cuda": + torch.cuda.synchronize(device) + start = time.perf_counter() + with torch.no_grad(): + output = prediction_function(model, inputs) + if device.type == "cuda": + torch.cuda.synchronize(device) + if repeat > 0: + timings.append(time.perf_counter() - start) + del output + output = None + return float(batch_size) / float(np.median(timings)) + except Exception as exc: + if _is_oom_error(exc): + return None + raise + finally: + model.train(was_training) + del inputs, output + _release_probe_memory(device) + + +def _select_throughput_batch_size( + measurements: Sequence[Tuple[int, float]], min_relative_improvement: float = 0.1, +) -> int: + """Prefer the smallest batch unless a larger one improves throughput materially.""" + if not measurements: + raise ValueError("At least one batch-throughput measurement is required.") + selected_batch, selected_throughput = measurements[0] + for batch_size, throughput in measurements[1:]: + if throughput >= selected_throughput * (1.0 + min_relative_improvement): + selected_batch, selected_throughput = batch_size, throughput + return int(selected_batch) + + def compute_auto_batch_sizes( model_devices: Sequence[Tuple[torch.nn.Module, torch.device]], n_jobs: int, @@ -120,40 +187,56 @@ def compute_auto_batch_sizes( in_channels: int, prediction_function: Callable[[torch.nn.Module, torch.Tensor], Any], ) -> List[int]: - """Probe the largest safe batch size independently on every CUDA device. + """Select a throughput-efficient batch size independently on every CUDA device. - The probe delegates its exponential and binary OOM search to - :func:`torch_em.util.compute_max_batch_size`. CPU and MPS use a conservative batch size of one, - because the torch-em probe intentionally supports CUDA OOM detection only. + Powers of two are benchmarked from one upwards. A larger batch is selected only when it improves + measured throughput by at least 10 percent, and probing stops after two consecutive candidates + fail to do so. CPU and MPS retain the conservative batch size of one. """ if n_jobs < 1: return [1] * len(model_devices) - # Cap the probe at the per-device job count and at 256 (no run needs a larger batch). jobs_per_device = (int(n_jobs) + len(model_devices) - 1) // len(model_devices) upper_bound = min(256, jobs_per_device) batch_sizes = [] - from torch_em.util import compute_max_batch_size - for model, device in model_devices: if device.type != "cuda": batch_sizes.append(1) continue - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", message="The batch size search reached the upper bound.*", category=UserWarning, + + measurements = [] + non_improving_candidates = 0 + candidate = 1 + while candidate <= upper_bound: + previous_selection = ( + _select_throughput_batch_size(measurements) if measurements else None ) - batch_size = compute_max_batch_size( + throughput = _measure_batch_throughput( model=model, + device=device, + batch_size=candidate, patch_shape=patch_shape, in_channels=in_channels, - device=device, - dtype=torch.float32, - safety_factor=0.8, # leave GPU-memory headroom below the probed maximum - max_batch_size=upper_bound, prediction_function=prediction_function, ) - batch_sizes.append(batch_size) + if throughput is None: + break + measurements.append((candidate, throughput)) + selection = _select_throughput_batch_size(measurements) + if previous_selection is not None: + if selection == candidate: + non_improving_candidates = 0 + else: + non_improving_candidates += 1 + if non_improving_candidates >= 2: + break + candidate *= 2 + + if not measurements: + raise RuntimeError( + f"The model does not fit batch size 1 on {device}. Reduce the patch shape or use a smaller model." + ) + batch_sizes.append(_select_throughput_batch_size(measurements)) return batch_sizes @@ -176,6 +259,31 @@ def _safe_put(output_queue: queue.Queue, item: Any, stop_event: threading.Event, raise PipelineAborted() +def _predict_with_oom_backoff( + model: torch.nn.Module, + items: List[Any], + device: torch.device, + predict_fn: Callable[[torch.nn.Module, List[Any], torch.device], List[Any]], +) -> Tuple[List[Any], int]: + """Retry an OOMing batch in halves and return the largest size proven safe.""" + try: + return predict_fn(model, items, device), len(items) + except Exception as exc: + if not _is_oom_error(exc) or len(items) == 1: + raise + + _release_probe_memory(device) + split = len(items) // 2 + warnings.warn( + f"Batch size {len(items)} ran out of memory on {device}; retrying with smaller batches.", + RuntimeWarning, + stacklevel=2, + ) + left, left_size = _predict_with_oom_backoff(model, items[:split], device, predict_fn) + right, right_size = _predict_with_oom_backoff(model, items[split:], device, predict_fn) + return left + right, min(left_size, right_size) + + def run_batched_pipeline( jobs: Iterable[Any], model_devices: Sequence[Tuple[torch.nn.Module, torch.device]], @@ -254,7 +362,10 @@ def consumer(worker_id): if batch: with torch.no_grad(): - predictions = predict_fn(model, [item.data for item in batch], device) + predictions, safe_batch_size = _predict_with_oom_backoff( + model, [item.data for item in batch], device, predict_fn, + ) + batch_size = min(batch_size, safe_batch_size) if len(predictions) != len(batch): raise RuntimeError( f"The batch predictor returned {len(predictions)} outputs for {len(batch)} inputs." diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index cb44fccd5..d97539f7e 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -946,8 +946,8 @@ def _run_full_inference( """Run queued, batched UniSAM2 encoder and decoder inference. Tiled reads and preprocessing, GPU inference, and output writes overlap through the torch-em - prediction pipeline. If `batch_size` is None, the largest safe batch is probed independently - on every selected CUDA device. + prediction pipeline. If `batch_size` is None, candidate sizes are benchmarked and a + throughput-efficient batch is selected independently on every CUDA device. """ from torch_em.util.prediction import predict_with_halo_pipelined from micro_sam.v2.normalization import normalize_raw @@ -985,7 +985,9 @@ def _preprocess(crop): n_jobs=n_blocks, patch_shape=patch_shape, in_channels=3, - prediction_function=lambda this_model, inputs: this_model(inputs), + # The probe's synthetic input bypasses `_preprocess`; clamp it into the [0, 1] + # range the model asserts (values are irrelevant to the memory measurement). + prediction_function=lambda this_model, inputs: this_model(inputs.clamp(0.0, 1.0)), ) batch_size = min(batch_sizes) finally: @@ -1076,7 +1078,7 @@ def initialize( pbar_update: Optional[callable] = None, z_block: Optional[int] = None, z_halo: Optional[int] = None, - batch_size: Optional[int] = None, + batch_size: Optional[int] = 1, devices: Devices = None, num_prefetch_workers: int = 4, num_write_workers: int = 1, @@ -1094,7 +1096,8 @@ def initialize( pbar_update: Callback to update an external progress bar. z_block: Number of slices per decoder z block. z_halo: Overlapping decoder slices used as context. - batch_size: Explicit tile or z-block batch size, or None for automatic capacity probing. + batch_size: Explicit tile or z-block batch size. Defaults to one; pass None for + throughput-based automatic selection. devices: Inference device or devices. None uses all visible GPUs when the model is on CUDA. num_prefetch_workers: Number of input reading and preprocessing threads. num_write_workers: Number of output writing threads for full tiled inference. @@ -1248,14 +1251,15 @@ def initialize( pbar_update: Optional[callable] = None, z_block: Optional[int] = None, z_halo: Optional[int] = None, - batch_size: Optional[int] = None, + batch_size: Optional[int] = 1, devices: Devices = None, num_prefetch_workers: int = 4, num_write_workers: int = 1, ) -> None: """Run tiled UniSAM2 inference and store foreground and distance predictions. - `batch_size=None` probes the largest safe batch independently on each selected CUDA device. + `batch_size=None` benchmarks candidate sizes and selects a throughput-efficient batch + independently on each selected CUDA device. Reads and preprocessing are queued through `num_prefetch_workers`; full inference also overlaps output writes through `num_write_workers`. """ diff --git a/micro_sam/v2/util.py b/micro_sam/v2/util.py index a767849a8..1da6ceac0 100644 --- a/micro_sam/v2/util.py +++ b/micro_sam/v2/util.py @@ -422,7 +422,7 @@ def precompute_image_embeddings( tile_shape: Optional[Tuple[int, int]] = None, halo: Optional[Tuple[int, int]] = None, verbose: bool = True, - batch_size: Optional[int] = None, + batch_size: Optional[int] = 1, devices: Devices = None, num_prefetch_workers: int = 4, pbar_init: Optional[callable] = None, @@ -441,8 +441,9 @@ def precompute_image_embeddings( tile_shape: Optional in-plane tile shape. halo: Optional in-plane tile halo. verbose: Whether to show progress. - batch_size: Number of independent tiles or slices per encoder call. If None, probe the - largest safe value independently on each CUDA device. + batch_size: Number of independent tiles or slices per encoder call. Defaults to one, which + is recommended for 10 GB MIG devices. Pass None to benchmark candidate sizes and select + a throughput-efficient value independently on each CUDA device. devices: Device or devices used for embedding inference. If None and the predictor is on CUDA, all visible CUDA devices are used. num_prefetch_workers: Number of threads used to read and preprocess input jobs. diff --git a/test/test_sam_annotator/test_annotator.py b/test/test_sam_annotator/test_annotator.py index edd55803d..e91b1ff32 100644 --- a/test/test_sam_annotator/test_annotator.py +++ b/test/test_sam_annotator/test_annotator.py @@ -220,7 +220,7 @@ def test_tiling_defaults_and_not_force_enabled(self, make_napari_viewer_proxy): viewer.close() def test_reset_inputs_keeps_optional_paths_unset(self, qapp): - """Clearing inputs on an image switch must not create a blank custom checkpoint path.""" + """Clearing inputs restores safe defaults without creating a blank custom checkpoint path.""" from micro_sam.sam_annotator._widgets import EmbeddingWidget ew = EmbeddingWidget(ndim_choice=True) @@ -231,12 +231,20 @@ def test_reset_inputs_keeps_optional_paths_unset(self, qapp): ew.custom_weights_param.setText(" ") assert ew.custom_weights is None + # Batching is explicitly user-controlled and starts at the safe single-item default. + assert ew.batch_size == 1 + assert ew.batch_size_param.value() == 1 + # Reproduce the input reset used when a different image layer is selected. ew.custom_weights_param.setText("/tmp/custom-weights.pt") + ew.batch_size_param.setValue(4) assert ew.custom_weights == "/tmp/custom-weights.pt" + assert ew.batch_size == 4 ew._reset_inputs_to_defaults() assert ew.custom_weights is None assert ew.custom_weights_param.text() == "" + assert ew.batch_size == 1 + assert ew.batch_size_param.value() == 1 @pytest.mark.parametrize("ndim", [2, 3]) def test_batched_checkbox_hidden_when_tiled(self, make_napari_viewer_proxy, ndim): diff --git a/test/test_sam_annotator/test_state.py b/test/test_sam_annotator/test_state.py index ce5193894..65b518712 100644 --- a/test/test_sam_annotator/test_state.py +++ b/test/test_sam_annotator/test_state.py @@ -54,6 +54,7 @@ def fake_get_sam_model(model_type, ndim, device, checkpoint_path, decoder_path, return object(), {} def fake_precompute_image_embeddings(**kwargs): + captured["batch_size"] = kwargs["batch_size"] return { "features": np.zeros((1, 1, 1, 1), dtype="float32"), "input_size": (8, 8), @@ -71,9 +72,10 @@ def fake_precompute_image_embeddings(**kwargs): checkpoint_path=" ", decoder_path="\t", prefer_decoder=False, + batch_size=4, ) - assert captured == {"checkpoint_path": None, "decoder_path": None} + assert captured == {"checkpoint_path": None, "decoder_path": None, "batch_size": 4} if __name__ == "__main__": diff --git a/test/test_v2_batched_inference.py b/test/test_v2_batched_inference.py index 18430054b..3898b3d00 100644 --- a/test/test_v2_batched_inference.py +++ b/test/test_v2_batched_inference.py @@ -1,4 +1,5 @@ import unittest +from unittest import mock import numpy as np import torch @@ -6,7 +7,11 @@ from bioimage_cpp.utils import Blocking from micro_sam.v2.batched_inference import ( - decode_tiled_3d_embeddings, decode_volume_embeddings, run_batched_pipeline, + _select_throughput_batch_size, + compute_auto_batch_sizes, + decode_tiled_3d_embeddings, + decode_volume_embeddings, + run_batched_pipeline, ) from micro_sam.v2.prompt_based_segmentation import TiledPromptableSegmentation3D @@ -82,6 +87,64 @@ def predict(model, items, device): self.assertEqual(outputs, {index: 2 * index + 3 for index in range(7)}) self.assertEqual(sum(progress), 7) + def test_pipeline_retries_ooming_batches(self): + outputs = {} + attempted_batch_sizes = [] + + def predict(model, items, device): + attempted_batch_sizes.append(len(items)) + if len(items) > 2: + raise torch.cuda.OutOfMemoryError("synthetic test OOM") + return [value + model for value in items] + + with self.assertWarnsRegex(RuntimeWarning, "retrying with smaller batches"): + run_batched_pipeline( + jobs=range(7), + model_devices=[(3, torch.device("cpu"))], + batch_sizes=[4], + load_fn=lambda value: 2 * value, + predict_fn=predict, + write_fn=outputs.__setitem__, + num_prefetch_workers=2, + ) + + self.assertEqual(outputs, {index: 2 * index + 3 for index in range(7)}) + self.assertEqual(attempted_batch_sizes[0], 4) + self.assertTrue(all(batch_size <= 2 for batch_size in attempted_batch_sizes[1:])) + + +class TestAutomaticBatchSizing(unittest.TestCase): + def test_requires_material_throughput_improvement(self): + self.assertEqual( + _select_throughput_batch_size([(1, 10.0), (2, 10.5)]), + 1, + ) + self.assertEqual( + _select_throughput_batch_size([(1, 10.0), (2, 11.5)]), + 2, + ) + + @mock.patch( + "micro_sam.v2.batched_inference._measure_batch_throughput", + side_effect=[10.0, 10.5, 10.6], + ) + def test_stops_after_two_non_improving_candidates(self, measure): + model = torch.nn.Linear(1, 1) + + batch_sizes = compute_auto_batch_sizes( + model_devices=[(model, torch.device("cuda:0"))], + n_jobs=64, + patch_shape=(8, 8), + in_channels=3, + prediction_function=lambda this_model, inputs: this_model(inputs), + ) + + self.assertEqual(batch_sizes, [1]) + self.assertEqual( + [call.kwargs["batch_size"] for call in measure.call_args_list], + [1, 2, 4], + ) + class TestBatchedDecoder(unittest.TestCase): def test_volume_decoder_batches_z_blocks(self): From a6533b48042cceeae8f5efa469276241d2c33a40 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 15:23:25 +0200 Subject: [PATCH 03/11] Minor patch for embedding progress bar over jobs --- micro_sam/sam_annotator/_widgets.py | 13 ++++++--- micro_sam/v2/batched_inference.py | 32 ++++++++++++++++++----- test/test_sam_annotator/test_annotator.py | 28 ++++++++++++++++++++ test/test_v2_batched_inference.py | 9 ++++++- 4 files changed, 72 insertions(+), 10 deletions(-) diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index 00c6b8f5f..e552dcb17 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -601,8 +601,13 @@ def button_clicked(self, label): # Set up the progress bar. We handle this via custom signals that are passed as callbacks to the # function that does the actual work. We need callbacks for initializing the progress bar, # updating it and for stopping the progress bar. -def _create_pbar_for_threadworker(): - pbar = progress() +def _create_pbar_for_threadworker(initial_description=None): + """Create a napari progress bar, optionally visible in an indeterminate preparation state. + + Supplying the description at construction time ensures napari can paint meaningful status before + a synchronous caller reaches the backend callback that provides the final work-item count. + """ + pbar = progress(desc=initial_description) pbar_signals = PBarSignals() pbar_signals.pbar_total.connect( lambda total: setattr(pbar, "total", total) @@ -2190,7 +2195,9 @@ def __call__(self, skip_validate=False): 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() + # Model and decoder preparation happens before the backend knows the tile / slice count, so + # start with an indeterminate status instead of leaving napari's activity display empty. + pbar, pbar_signals = _create_pbar_for_threadworker("Preparing image embeddings") # @thread_worker() def compute_image_embedding(): diff --git a/micro_sam/v2/batched_inference.py b/micro_sam/v2/batched_inference.py index ee593bdea..f7bc2786e 100644 --- a/micro_sam/v2/batched_inference.py +++ b/micro_sam/v2/batched_inference.py @@ -318,6 +318,7 @@ def run_batched_pipeline( input_queue = queue.Queue(maxsize=max(2 * sum(batch_sizes), 2)) output_queue = queue.Queue(maxsize=max(2 * len(model_devices), 2)) + progress_queue = queue.Queue() stop_event = threading.Event() error_box = [] error_lock = threading.Lock() @@ -392,7 +393,7 @@ def writer(): break write_fn(item.spec, item.data) if update_progress is not None: - update_progress(1) + progress_queue.put(1) except PipelineAborted: pass except Exception as exc: # noqa @@ -409,14 +410,33 @@ def writer(): ] threads = [writer_thread, *consumer_threads, *producer_threads] + def forward_progress(block): + increments = 0 + try: + increments = progress_queue.get(timeout=0.05 if block else 0) + except queue.Empty: + return + while True: + try: + increments += progress_queue.get_nowait() + except queue.Empty: + break + update_progress(increments) + try: for thread in threads: thread.start() - for thread in producer_threads: - thread.join() - for thread in consumer_threads: - thread.join() - writer_thread.join() + if update_progress is None: + for thread in threads: + thread.join() + else: + # The writer performs I/O off-thread, but UI progress callbacks must run on the calling + # thread. Poll completed writes here instead of queuing Qt signals behind this wait. + while any(thread.is_alive() for thread in threads): + forward_progress(block=True) + for thread in threads: + thread.join() + forward_progress(block=False) finally: stop_event.set() for thread in threads: diff --git a/test/test_sam_annotator/test_annotator.py b/test/test_sam_annotator/test_annotator.py index e91b1ff32..a78dd1743 100644 --- a/test/test_sam_annotator/test_annotator.py +++ b/test/test_sam_annotator/test_annotator.py @@ -9,6 +9,34 @@ from micro_sam._test_util import check_layer_initialization +def test_progress_bar_initial_description(monkeypatch): + """A progress description supplied at creation is visible before the backend reports a total.""" + from micro_sam.sam_annotator import _widgets + + captured = {} + + class FakeProgress: + def update(self, value): + pass + + def set_description(self, description): + pass + + def close(self): + pass + + def reset(self): + pass + + def fake_progress(**kwargs): + captured["kwargs"] = kwargs + return FakeProgress() + + monkeypatch.setattr(_widgets, "progress", fake_progress) + _widgets._create_pbar_for_threadworker("Preparing image embeddings") + assert captured["kwargs"] == {"desc": "Preparing image embeddings"} + + class TestDetectNdim: """Test the detect_ndim helper function.""" diff --git a/test/test_v2_batched_inference.py b/test/test_v2_batched_inference.py index 3898b3d00..bd431ae17 100644 --- a/test/test_v2_batched_inference.py +++ b/test/test_v2_batched_inference.py @@ -1,3 +1,4 @@ +import threading import unittest from unittest import mock @@ -68,11 +69,16 @@ class TestBatchedPipeline(unittest.TestCase): def test_pipeline_batches_and_writes_all_jobs(self): outputs = {} progress = [] + progress_threads = [] def predict(model, items, device): self.assertEqual(device, torch.device("cpu")) return [value + model for value in items] + def update_progress(update): + progress.append(update) + progress_threads.append(threading.get_ident()) + run_batched_pipeline( jobs=range(7), model_devices=[(3, torch.device("cpu"))], @@ -81,11 +87,12 @@ def predict(model, items, device): predict_fn=predict, write_fn=outputs.__setitem__, num_prefetch_workers=2, - update_progress=progress.append, + update_progress=update_progress, ) self.assertEqual(outputs, {index: 2 * index + 3 for index in range(7)}) self.assertEqual(sum(progress), 7) + self.assertEqual(set(progress_threads), {threading.get_ident()}) def test_pipeline_retries_ooming_batches(self): outputs = {} From 9f51f590a1a4ef35529346e3bf094c92366001a0 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 16:00:27 +0200 Subject: [PATCH 04/11] Make batch_size spin box modular - disable for CPU --- micro_sam/sam_annotator/_widgets.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index e552dcb17..fd5f2e662 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -1794,14 +1794,21 @@ def _create_settings_widget(self): setting_values.layout().addLayout(layout) # Use a conservative default and let users opt into larger per-GPU batches when their - # workload and available memory benefit from them. + # workload and available memory benefit from them. Batching only helps on a GPU, so the + # control is hidden whenever the effective device is the CPU (the default of 1 is still + # used then). Visibility tracks the device dropdown, so it also updates when the user + # switches devices or when 'auto' resolves to the CPU. self.batch_size = 1 self.batch_size_param, batch_size_layout = self._add_int_param( "batch_size", self.batch_size, min_val=1, max_val=64, - title="batch size per GPU:", + title="batch size", tooltip=get_tooltip("embedding", "batch_size"), ) - setting_values.layout().addLayout(batch_size_layout) + self._batch_size_widget = QtWidgets.QWidget() + self._batch_size_widget.setLayout(batch_size_layout) + setting_values.layout().addWidget(self._batch_size_widget) + self.device_dropdown.currentIndexChanged.connect(self._update_batch_size_visibility) + self._update_batch_size_visibility() # Create UI for the save path. self.embeddings_save_path = None @@ -1888,6 +1895,14 @@ def _update_tiling_visibility(self, index=None): self.tiling = self.tiling_dropdown.currentText() self._tiling_widget.setVisible(self.tiling == "yes") + def _update_batch_size_visibility(self, index=None): + # Show the batch size field only when the effective device is a GPU. Batching does not help + # on the CPU, so 'auto' resolving to the CPU (or an explicit CPU selection) hides the field. + device = self.device + if device == "auto": + device = util._get_default_device() + self._batch_size_widget.setVisible(str(device) != "cpu") + def _apply_default_tiling_for_shape(self, shape): # Enable tiling by default for large in-plane images, using the central v2 tiling defaults. # 'shape' is the spatial image shape (channel axis already removed). Shared by the layer-based From f2c37c3c60449e9326911e0e97b36ba1a165b4f0 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 16:17:29 +0200 Subject: [PATCH 05/11] Use FP16 autocast for memory efficient decoder processing --- micro_sam/v2/instance_segmentation.py | 28 ++++++++++++++------------ test/test_v2_automatic_segmentation.py | 14 +++++++++++-- 2 files changed, 27 insertions(+), 15 deletions(-) diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index d97539f7e..8dd2e1cce 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -864,18 +864,28 @@ def forward(self, x): # noqa def _decode_3d_feature_batch(model, features, original_size, device): - """Decode precomputed feature volumes with shape (B, Z, C, H, W).""" + """Decode precomputed feature volumes with shape (B, Z, C, H, W). + + CUDA decoder inference uses FP16 autocast, matching UniSAM2 training and keeping the largest + default annotator tile within a 10 GB MIG partition. Cached predictions remain float32. + """ if features.ndim != 5: raise ValueError(f"Expected batched features with shape (B, Z, C, H, W), got {features.shape}.") + device = features.device if device is None else torch.device(device) img_size = getattr(model.encoder, "img_size", 1024) real_encoder = model.encoder model.encoder = _StubEncoder(features, img_size) try: dummy = torch.zeros((features.shape[0], 3, features.shape[1], *original_size), device=device) - output = model(dummy) + autocast = ( + torch.autocast(device_type="cuda", dtype=torch.float16) + if device.type == "cuda" else contextlib.nullcontext() + ) + with autocast: + output = model(dummy) finally: model.encoder = real_encoder - return output.detach().cpu().numpy() + return output.detach().float().cpu().numpy() def _decode_3d_feature_block(model, feature, original_size, device): @@ -1033,16 +1043,8 @@ def _run_decoder_2d(self, image_embeddings): feature = torch.as_tensor(features, device=self._device).float() original_size = tuple(int(s) for s in np.array(image_embeddings["original_size"]).reshape(-1)[:2]) - model = self._model - img_size = getattr(model.encoder, "img_size", 1024) - real_encoder = model.encoder - model.encoder = _StubEncoder(feature, img_size) - try: - dummy = torch.zeros((1, 3, 1, *original_size), device=self._device) - output = model(dummy) - finally: - model.encoder = real_encoder - return output[0, :, 0].detach().cpu().numpy() + output = _decode_3d_feature_batch(self._model, feature.unsqueeze(1), original_size, self._device) + return output[0, :, 0] @torch.no_grad() def _run_decoder_3d( diff --git a/test/test_v2_automatic_segmentation.py b/test/test_v2_automatic_segmentation.py index 9bc74dcf1..d6083c326 100644 --- a/test/test_v2_automatic_segmentation.py +++ b/test/test_v2_automatic_segmentation.py @@ -90,9 +90,10 @@ class _FakeUNETR: A plain (non-nn.Module) class so the encoder attribute can be freely swapped and restored. """ - def __init__(self, img_size=8): + def __init__(self, img_size=8, output_dtype=torch.float32): self.encoder = types.SimpleNamespace(img_size=img_size) self.seen = [] # features the stub returned on the most recent call + self.output_dtype = output_dtype self.call_z = [] # z size of each decoder call (one per z block) self.call_hw = [] @@ -102,7 +103,7 @@ def __call__(self, x): # x: (1, 3, Z, H, W) dummy self.call_z.append(z) h, w = x.shape[-2:] self.call_hw.append((h, w)) - return torch.zeros((1, 4, z, h, w)) + return torch.zeros((1, 4, z, h, w), dtype=self.output_dtype) def test_decoder_3d_stub_returns_per_slice_features_in_order(): @@ -156,10 +157,19 @@ def test_decoder_2d_uses_original_non_square_shape(): assert model.call_hw == [(4, 8)] +def test_decoder_predictions_are_cached_as_float32(): + features = np.zeros((1, 2, 4, 4), dtype="float32") + model = _FakeUNETR(img_size=8, output_dtype=torch.float16) + out = _run_decoder_2d(model, {"features": features, "original_size": (8, 8)}) + + assert out.dtype == np.float32 + + def test_decoder_3d_zchunks_deep_volume(): # Regression for the z-tiling fix: a deep volume (small in-plane, not tiled in-plane) must be # decoded in bounded z blocks, not all at once, so peak memory stays bounded. z, c, h, w = 10, 2, 4, 4 + feats = np.zeros((z, c, h, w), dtype="float32") model = _FakeUNETR(img_size=8) out = _run_decoder_3d(model, {"features": feats, "original_size": (8, 8)}) From a22b70ea62f0f56455b26a4092254e5b89323d28 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 16:39:21 +0200 Subject: [PATCH 06/11] Extend FPS16 support to MPS --- micro_sam/v2/instance_segmentation.py | 18 ++++++++++------- test/test_v2_automatic_segmentation.py | 27 +++++++++++++++++++++++++- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index 8dd2e1cce..c5ada3e5f 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -863,11 +863,19 @@ def forward(self, x): # noqa return [feature] +def _get_decoder_autocast(device): + """Use FP16 decoder autocast on supported accelerator backends.""" + device_type = torch.device(device).type + if device_type in ("cuda", "mps"): + return torch.autocast(device_type=device_type, dtype=torch.float16) + return contextlib.nullcontext() + + def _decode_3d_feature_batch(model, features, original_size, device): """Decode precomputed feature volumes with shape (B, Z, C, H, W). - CUDA decoder inference uses FP16 autocast, matching UniSAM2 training and keeping the largest - default annotator tile within a 10 GB MIG partition. Cached predictions remain float32. + CUDA and MPS decoder inference use FP16 autocast. This matches UniSAM2 training on CUDA and keeps + the largest default annotator tile within a 10 GB MIG partition. Cached predictions remain float32. """ if features.ndim != 5: raise ValueError(f"Expected batched features with shape (B, Z, C, H, W), got {features.shape}.") @@ -877,11 +885,7 @@ def _decode_3d_feature_batch(model, features, original_size, device): model.encoder = _StubEncoder(features, img_size) try: dummy = torch.zeros((features.shape[0], 3, features.shape[1], *original_size), device=device) - autocast = ( - torch.autocast(device_type="cuda", dtype=torch.float16) - if device.type == "cuda" else contextlib.nullcontext() - ) - with autocast: + with _get_decoder_autocast(device): output = model(dummy) finally: model.encoder = real_encoder diff --git a/test/test_v2_automatic_segmentation.py b/test/test_v2_automatic_segmentation.py index d6083c326..3910ddf0c 100644 --- a/test/test_v2_automatic_segmentation.py +++ b/test/test_v2_automatic_segmentation.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext import types import inspect @@ -7,7 +8,7 @@ from micro_sam.v2.instance_segmentation import ( _block_shape_and_halo, _set_image_predictor_from_backbone, - UniSAM2InstanceSegmentation, TiledUniSAM2InstanceSegmentation, + _get_decoder_autocast, UniSAM2InstanceSegmentation, TiledUniSAM2InstanceSegmentation, get_instance_segmentation_generator, get_decoder, ) from micro_sam.v2.postprocessing import DEFAULT_POSTPROCESSING, run_multicut @@ -165,6 +166,30 @@ def test_decoder_predictions_are_cached_as_float32(): assert out.dtype == np.float32 +@pytest.mark.parametrize("device_type", ("cuda", "mps")) +def test_decoder_autocast_uses_fp16_on_accelerators(device_type, monkeypatch): + calls = [] + + def fake_autocast(device_type, dtype): + calls.append((device_type, dtype)) + return nullcontext() + + monkeypatch.setattr(torch, "autocast", fake_autocast) + with _get_decoder_autocast(torch.device(device_type)): + pass + + assert calls == [(device_type, torch.float16)] + + +def test_decoder_autocast_leaves_cpu_unchanged(monkeypatch): + def fail_autocast(*args, **kwargs): + pytest.fail("CPU decoder inference must not enable autocast.") + + monkeypatch.setattr(torch, "autocast", fail_autocast) + with _get_decoder_autocast(torch.device("cpu")): + pass + + def test_decoder_3d_zchunks_deep_volume(): # Regression for the z-tiling fix: a deep volume (small in-plane, not tiled in-plane) must be # decoded in bounded z blocks, not all at once, so peak memory stays bounded. From e5af2142b2df002911cb9f6229575c612d7827d9 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 18:28:02 +0200 Subject: [PATCH 07/11] Fix volumetric progress reporting and decoder OOM recovery --- micro_sam/sam_annotator/_widgets.py | 4 +- micro_sam/v2/batched_inference.py | 47 ++++++++++++----- micro_sam/v2/prompt_based_segmentation.py | 61 ++++++++++++++++------- test/test_v2_batched_inference.py | 48 +++++++++++++++++- 4 files changed, 127 insertions(+), 33 deletions(-) diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index fd5f2e662..f16086aa0 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -4453,8 +4453,8 @@ def __call__(self): # plain 2d image. Thread workers are disabled in this tool (see top of module), so the run is # synchronous; we drive the bar via callbacks the backends call between units and pump the Qt # event loop with 'processEvents' on each update so it repaints live. It is always closed in - # the 'finally' block. (3d decoder inference runs through a thread pool and is reported as a - # single step, since it cannot update the napari bar live.) + # the 'finally' block. Batched 3d inference forwards completed tile-slice increments from its + # worker threads to this calling thread so napari can repaint them safely. pbar, pbar_signals = _create_pbar_for_threadworker() def pbar_init(total, description): diff --git a/micro_sam/v2/batched_inference.py b/micro_sam/v2/batched_inference.py index f7bc2786e..5a5122f19 100644 --- a/micro_sam/v2/batched_inference.py +++ b/micro_sam/v2/batched_inference.py @@ -126,6 +126,18 @@ def _release_probe_memory(device: torch.device) -> None: torch.cuda.empty_cache() +def _release_fragmented_cuda_cache(device: torch.device) -> None: + """Release inactive cache when it leaves the logical CUDA device with too little free VRAM.""" + if device.type != "cuda": + return + free_memory, total_memory = torch.cuda.mem_get_info(device) + allocated_memory = torch.cuda.memory_allocated(device) + inactive_memory = torch.cuda.memory_reserved(device) - allocated_memory + if free_memory < 0.25 * total_memory and inactive_memory > 0.25 * total_memory: + with torch.cuda.device(device): + torch.cuda.empty_cache() + + def _measure_batch_throughput( model: torch.nn.Module, device: torch.device, @@ -269,10 +281,12 @@ def _predict_with_oom_backoff( try: return predict_fn(model, items, device), len(items) except Exception as exc: - if not _is_oom_error(exc) or len(items) == 1: + if not _is_oom_error(exc): raise _release_probe_memory(device) + if len(items) == 1: + return predict_fn(model, items, device), 1 split = len(items) // 2 warnings.warn( f"Batch size {len(items)} ran out of memory on {device}; retrying with smaller batches.", @@ -293,6 +307,7 @@ def run_batched_pipeline( write_fn: Callable[[Any, Any], None], num_prefetch_workers: int = 4, update_progress: Optional[Callable[[int], None]] = None, + progress_increment: Optional[Callable[[Any], int]] = None, ) -> None: """Run load/preprocess, batched inference, and writing as a bounded threaded pipeline. @@ -393,7 +408,8 @@ def writer(): break write_fn(item.spec, item.data) if update_progress is not None: - progress_queue.put(1) + increment = 1 if progress_increment is None else int(progress_increment(item.spec)) + progress_queue.put(increment) except PipelineAborted: pass except Exception as exc: # noqa @@ -895,7 +911,9 @@ def _predict_jobs(model: torch.nn.Module, items: List[Dict], device: torch.devic raise RuntimeError("Decoder jobs with different output sizes cannot share a batch.") features = torch.stack([item["feature"] for item in items]).to(device, non_blocking=True) output = _decode_3d_feature_batch(model, features, original_size, device) - return [np.array(value) for value in output] + predictions = [np.array(value) for value in output] + _release_fragmented_cuda_cache(device) + return predictions def _feature_shape(job: Dict) -> Tuple[int, ...]: @@ -973,6 +991,8 @@ def _run_decoder_jobs( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + update_progress: Optional[Callable[[int], None]] = None, + progress_increment: Optional[Callable[[Dict], int]] = None, ) -> None: if len(jobs) == 0: return @@ -982,8 +1002,13 @@ def _run_decoder_jobs( for job in jobs: groups[(_feature_shape(job), job["original_size"])].append(job) + # Run the largest shape first so later, smaller groups can reuse its allocator blocks. Starting + # with a boundary z-block and growing to the full context fragmented 10 GB MIG allocators. + group_keys = sorted(groups, key=lambda key: np.prod(key[0]) * np.prod(key[1]), reverse=True) + try: - for group_jobs in groups.values(): + for group_key in group_keys: + group_jobs = groups[group_key] run_batched_pipeline( jobs=group_jobs, model_devices=model_devices, @@ -992,6 +1017,8 @@ def _run_decoder_jobs( predict_fn=_predict_jobs, write_fn=write_fn, num_prefetch_workers=num_prefetch_workers, + update_progress=update_progress, + progress_increment=progress_increment, ) finally: release_model_replicas(model_devices) @@ -1048,12 +1075,11 @@ def write_prediction(job, prediction): inner = job["z0"] - job["c0"] count = job["z1"] - job["z0"] output[:, job["z0"]:job["z1"]] = prediction[:, inner:inner + count] - if pbar_update is not None: - pbar_update(count) _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, progress_increment=lambda job: job["z1"] - job["z0"], ) return output @@ -1106,12 +1132,11 @@ def write_prediction(job, prediction): block.inner_block.begin, block.inner_block.end, )) output[(slice(None),) + inner] = prediction[(slice(None), slice(0, 1)) + local][:, 0] - if pbar_update is not None: - pbar_update(1) _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, ) return output @@ -1173,12 +1198,11 @@ def write_prediction(job, prediction): local_z = job["z0"] - job["c0"] prediction = prediction[:, local_z:local_z + z_count] output[(slice(None), slice(job["z0"], job["z1"])) + inner] = prediction[(slice(None), slice(None)) + local] - if pbar_update is not None: - pbar_update(z_count) _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, progress_increment=lambda job: job["z1"] - job["z0"], ) return output @@ -1220,11 +1244,10 @@ def write_prediction(job, prediction): block.inner_block.begin, block.inner_block.end, )) output[(slice(None),) + inner] = prediction[(slice(None), 0) + local] - if pbar_update is not None: - pbar_update(1) _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + update_progress=pbar_update, ) return output diff --git a/micro_sam/v2/prompt_based_segmentation.py b/micro_sam/v2/prompt_based_segmentation.py index 6be14f2cc..dd0afb2d8 100644 --- a/micro_sam/v2/prompt_based_segmentation.py +++ b/micro_sam/v2/prompt_based_segmentation.py @@ -1,7 +1,7 @@ import gc import ctypes import platform -import threading +import queue from concurrent import futures from copy import copy @@ -1094,24 +1094,54 @@ def _inner_slices(self, tile_id): glob = tuple(slice(b, e) for b, e in zip(block.inner_block.begin, block.inner_block.end)) return local, glob - def _run_tile_jobs(self, tile_ids: List[int], function: Callable) -> List[Tuple]: + def _run_tile_jobs( + self, tile_ids: List[int], function: Callable, update_progress: Optional[Callable[[int], None]] = None, + ) -> List[Tuple]: """Run tile jobs concurrently across devices and serially within each device.""" groups = {} for tile_id in tile_ids: worker_id = tile_id % len(self._predictor_devices) groups.setdefault(worker_id, []).append(tile_id) - def run_group(group): - return [(tile_id, function(tile_id)) for tile_id in group] + def run_group(group, worker_update=None): + results = [] + for tile_id in group: + if worker_update is None: + result = function(tile_id) + else: + result = function(tile_id, worker_update) + results.append((tile_id, result)) + return results if len(groups) < 2: - return run_group(next(iter(groups.values()))) if groups else [] + if not groups: + return [] + group = next(iter(groups.values())) + return run_group(group, update_progress) + + progress_queue = queue.Queue() + + def forward_progress(): + increment = 0 + while True: + try: + increment += int(progress_queue.get_nowait()) + except queue.Empty: + break + if increment: + update_progress(increment) results = [] with futures.ThreadPoolExecutor(max_workers=len(groups)) as pool: - tasks = [pool.submit(run_group, group) for group in groups.values()] - for task in tasks: - results.extend(task.result()) + worker_update = progress_queue.put if update_progress is not None else None + tasks = [pool.submit(run_group, group, worker_update) for group in groups.values()] + pending = set(tasks) + while pending: + done, pending = futures.wait(pending, timeout=0.05, return_when=futures.FIRST_COMPLETED) + forward_progress() + for task in done: + results.extend(task.result()) + forward_progress() return sorted(results, key=lambda result: result[0]) def segment_slice(self, frame_idx, points=None, labels=None, boxes=None, masks=None, object_id=1): @@ -1256,20 +1286,15 @@ def predict(self, update_progress=None, early_stop_patience=None, z_range=None): prompted in several tiles keeps one id and is merged across the tile boundaries. """ segmentation = np.zeros(self.shape, dtype="uint64") - progress_lock = threading.Lock() - def locked_update(value): - if update_progress is None: - return - with progress_lock: - update_progress(value) - - def predict_tile(tile_id): + def predict_tile(tile_id, tile_update=None): return self._segmenters[tile_id].predict( - update_progress=locked_update, early_stop_patience=early_stop_patience, z_range=z_range, + update_progress=tile_update, early_stop_patience=early_stop_patience, z_range=z_range, ) - for tile_id, tile_seg in self._run_tile_jobs(sorted(self._segmenters), predict_tile): + for tile_id, tile_seg in self._run_tile_jobs( + sorted(self._segmenters), predict_tile, update_progress=update_progress, + ): local, glob = self._inner_slices(tile_id) inner = tile_seg[(slice(None),) + local] region = segmentation[(slice(None),) + glob] diff --git a/test/test_v2_batched_inference.py b/test/test_v2_batched_inference.py index bd431ae17..3e2bb1dd1 100644 --- a/test/test_v2_batched_inference.py +++ b/test/test_v2_batched_inference.py @@ -8,6 +8,7 @@ from bioimage_cpp.utils import Blocking from micro_sam.v2.batched_inference import ( + _run_decoder_jobs, _select_throughput_batch_size, compute_auto_batch_sizes, decode_tiled_3d_embeddings, @@ -119,6 +120,25 @@ def predict(model, items, device): self.assertEqual(attempted_batch_sizes[0], 4) self.assertTrue(all(batch_size <= 2 for batch_size in attempted_batch_sizes[1:])) + def test_pipeline_retries_singleton_after_cache_release(self): + outputs = {} + attempts = 0 + + def predict(model, items, device): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise torch.cuda.OutOfMemoryError("synthetic fragmented-cache OOM") + return [value + model for value in items] + + run_batched_pipeline( + jobs=[0], model_devices=[(3, torch.device("cpu"))], batch_sizes=[1], + load_fn=lambda value: 2 * value, predict_fn=predict, write_fn=outputs.__setitem__, + ) + + self.assertEqual(outputs, {0: 3}) + self.assertEqual(attempts, 2) + class TestAutomaticBatchSizing(unittest.TestCase): def test_requires_material_throughput_improvement(self): @@ -157,6 +177,7 @@ class TestBatchedDecoder(unittest.TestCase): def test_volume_decoder_batches_z_blocks(self): features = np.arange(8 * 2 * 2 * 2, dtype="float32").reshape(8, 2, 2, 2) model = FakeUNETR() + progress = [] output = decode_volume_embeddings( model, @@ -166,10 +187,12 @@ def test_volume_decoder_batches_z_blocks(self): z_halo=0, batch_size=2, num_prefetch_workers=2, + pbar_update=progress.append, ) self.assertEqual(output.shape, (4, 8, 8, 8)) self.assertEqual(model.batch_sizes, [2, 2]) + self.assertEqual(sum(progress), 8) def test_tiled_volume_decoder_batches_tiles_and_z_blocks(self): features = Group( @@ -185,6 +208,7 @@ def test_tiled_volume_decoder_batches_tiles_and_z_blocks(self): halo=(0, 0), ) model = FakeUNETR() + progress = [] output = decode_tiled_3d_embeddings( model, @@ -194,6 +218,7 @@ def test_tiled_volume_decoder_batches_tiles_and_z_blocks(self): z_halo=0, batch_size=2, num_prefetch_workers=2, + pbar_update=progress.append, ) self.assertEqual(output.shape, (4, 4, 8, 8)) @@ -202,6 +227,21 @@ def test_tiled_volume_decoder_batches_tiles_and_z_blocks(self): self.assertTrue(np.allclose(output[:, :, 4:, :4], 3)) self.assertTrue(np.allclose(output[:, :, 4:, 4:], 4)) self.assertEqual(model.batch_sizes, [2, 2, 2, 2]) + self.assertEqual(sum(progress), 16) + + def test_largest_decoder_shape_runs_first(self): + jobs = [ + {"source": np.ones((2, 2, 2, 2), dtype="float32"), "original_size": (8, 8), "name": "small"}, + {"source": np.ones((4, 2, 2, 2), dtype="float32"), "original_size": (8, 8), "name": "large"}, + ] + write_order = [] + + _run_decoder_jobs( + FakeUNETR(), jobs, lambda job, prediction: write_order.append(job["name"]), + batch_size=1, devices="cpu", num_prefetch_workers=1, + ) + + self.assertEqual(write_order, ["large", "small"]) class FakeInteractiveTile: @@ -226,14 +266,20 @@ def test_active_tile_columns_are_stitched_after_device_jobs(self): for tile_id in range(4) } progress = [] + progress_threads = [] + + def update_progress(value): + progress.append(value) + progress_threads.append(threading.get_ident()) - output = segmenter.predict(update_progress=progress.append) + output = segmenter.predict(update_progress=update_progress) self.assertTrue(np.all(output[:, :4, :4] == 1)) self.assertTrue(np.all(output[:, :4, 4:] == 2)) self.assertTrue(np.all(output[:, 4:, :4] == 3)) self.assertTrue(np.all(output[:, 4:, 4:] == 4)) self.assertEqual(sum(progress), 8) + self.assertEqual(set(progress_threads), {threading.get_ident()}) if __name__ == "__main__": From b2b34e09fd33aaf9d19959082e1c9d84b38e1bf7 Mon Sep 17 00:00:00 2001 From: anwai98 Date: Wed, 22 Jul 2026 18:32:55 +0200 Subject: [PATCH 08/11] Fix 3d tiled automatic segmentation with in-plane tile shape --- micro_sam/v2/instance_segmentation.py | 17 +++++++++++++--- test/test_v2_automatic_segmentation.py | 28 ++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index c5ada3e5f..1fae5efc1 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -816,7 +816,8 @@ def _block_shape_and_halo(spatial_shape, ndim, tile_shape, halo): Args: spatial_shape: The spatial image shape, (Y, X) for 2d or (Z, Y, X) for 3d. ndim: The number of spatial dimensions (2 or 3). - tile_shape: The in-plane/3d tile shape, or None for no tiling. + tile_shape: The tile shape, or None for no tiling. For 3d data either an in-plane (y, x) + tile, which keeps the default z chunking, or an explicit (z, y, x) tile. halo: The tile halo, or None for no overlap. Returns: @@ -832,8 +833,18 @@ def _block_shape_and_halo(spatial_shape, ndim, tile_shape, halo): block_shape = (1, *spatial_shape) block_halo = (0, 0, 0) elif is_3d: - block_shape = tuple(tile_shape) # (z, y, x) - block_halo = (0, 0, 0) if halo is None else tuple(halo) + # Tiling is in-plane, so the CLI and the annotator pass a 2-entry (y, x) tile. Prepend the + # default z block, keeping z chunked exactly as it is without tiling. A 3-entry (z, y, x) + # tile is used as given. + if len(tile_shape) == 2: + n_slices = spatial_shape[0] + z_block = min(DEFAULT_TILE_Z, n_slices) + block_shape = (z_block, *tile_shape) + z_halo = DEFAULT_HALO_Z if z_block < n_slices else 0 + block_halo = (z_halo, *((0, 0) if halo is None else tuple(halo)[-2:])) + else: + block_shape = tuple(tile_shape) # (z, y, x) + block_halo = (0, 0, 0) if halo is None else tuple(halo) else: block_shape = (1, *tile_shape) # (1, y, x) block_halo = (0, *((0, 0) if halo is None else halo)) diff --git a/test/test_v2_automatic_segmentation.py b/test/test_v2_automatic_segmentation.py index 3910ddf0c..1c9072b00 100644 --- a/test/test_v2_automatic_segmentation.py +++ b/test/test_v2_automatic_segmentation.py @@ -233,6 +233,34 @@ def test_block_shape_3d_tiling_uses_tile(): assert halo == (2, 64, 64) +def test_block_shape_3d_in_plane_tiling_keeps_default_z_chunking(): + # The CLI and the annotator only ever pass an in-plane (y, x) tile; it must be combined with the + # default z block instead of being used as a (z, y, x) block shape. + block, halo = _block_shape_and_halo((50, 1024, 1024), ndim=3, tile_shape=(512, 512), halo=(64, 64)) + assert block == (DEFAULT_TILE_Z, 512, 512) + assert halo == (DEFAULT_HALO_Z, 64, 64) + + +def test_block_shape_3d_in_plane_tiling_shallow_volume(): + # Fewer slices than the default z block -> single z block, no z halo. + block, halo = _block_shape_and_halo((3, 1024, 1024), ndim=3, tile_shape=(512, 512), halo=(64, 64)) + assert block == (3, 512, 512) + assert halo == (0, 64, 64) + + +def test_block_shape_3d_in_plane_tiling_without_halo(): + block, halo = _block_shape_and_halo((50, 1024, 1024), ndim=3, tile_shape=(512, 512), halo=None) + assert block == (DEFAULT_TILE_Z, 512, 512) + assert halo == (DEFAULT_HALO_Z, 0, 0) + + +@pytest.mark.parametrize("tile_shape,halo", [((512, 512), (64, 64)), ((4, 512, 512), (2, 64, 64))]) +def test_block_shape_3d_matches_predict_with_halo_arity(tile_shape, halo): + # predict_with_halo asserts len(block_shape) == len(halo) == ndim. + block, block_halo = _block_shape_and_halo((50, 1024, 1024), ndim=3, tile_shape=tile_shape, halo=halo) + assert len(block) == len(block_halo) == 3 + + def test_block_shape_2d_tiling_uses_tile(): block, halo = _block_shape_and_halo((1024, 1024), ndim=2, tile_shape=(512, 512), halo=(64, 64)) assert block == (1, 512, 512) From 063fcca4744663ddd841f912816faf76237ecef4 Mon Sep 17 00:00:00 2001 From: anwai98 Date: Wed, 22 Jul 2026 19:13:47 +0200 Subject: [PATCH 09/11] Fix tiled 3d interactive segmentation on MPS by comparing device type --- micro_sam/sam_annotator/_widgets.py | 2 +- micro_sam/util.py | 15 +++++++++++++ micro_sam/v1/instance_segmentation.py | 2 +- micro_sam/v2/prompt_based_segmentation.py | 3 ++- test/test_instance_segmentation.py | 15 +++++++++++++ test/test_util.py | 13 +++++++++++ test/test_v2_prompt_based_segmentation.py | 27 +++++++++++++++++++++++ 7 files changed, 74 insertions(+), 3 deletions(-) diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index f16086aa0..6ffb7c35c 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -3887,7 +3887,7 @@ def _allow_segment_3d(self): return True state = AnnotatorState() predictor = state.predictor - if str(predictor.device) == "cpu" or str(predictor.device) == "mps": + if util.device_type(predictor.device) in ("cpu", "mps"): n_slices = self._viewer.layers["auto_segmentation"].data.shape[0] if state.is_sam2: from micro_sam.precompute_state import _has_autoseg_state diff --git a/micro_sam/util.py b/micro_sam/util.py index 6a8f0d6cb..1106f284f 100644 --- a/micro_sam/util.py +++ b/micro_sam/util.py @@ -127,6 +127,21 @@ def _get_default_device(): return device +def device_type(device: Union[str, torch.device]) -> str: + """Get the device type ('cpu', 'cuda' or 'mps'), ignoring any device index. + + Torch reports accelerators with an index (e.g. 'mps:0'), so comparing 'str(device)' against + 'mps' or 'cuda' silently fails. Compare against this instead. + + Args: + device: The device, as a string or torch.device. + + Returns: + The device type. + """ + return torch.device(device).type + + def _configure_mps_memory(device: Union[str, torch.device]) -> None: """Disable the MPS memory watermark so 3d automatic segmentation does not hit a premature OOM. diff --git a/micro_sam/v1/instance_segmentation.py b/micro_sam/v1/instance_segmentation.py index 54d945249..d71d51004 100644 --- a/micro_sam/v1/instance_segmentation.py +++ b/micro_sam/v1/instance_segmentation.py @@ -374,7 +374,7 @@ def __init__( # we set the points per batch to 16 for mps for performance reasons # and otherwise keep them at the default of 64 if points_per_batch is None: - points_per_batch = 16 if str(predictor.device) == "mps" else 64 + points_per_batch = 16 if util.device_type(predictor.device) == "mps" else 64 self._points_per_batch = points_per_batch self._crop_n_layers = crop_n_layers diff --git a/micro_sam/v2/prompt_based_segmentation.py b/micro_sam/v2/prompt_based_segmentation.py index dd0afb2d8..f5e2dee9d 100644 --- a/micro_sam/v2/prompt_based_segmentation.py +++ b/micro_sam/v2/prompt_based_segmentation.py @@ -10,6 +10,7 @@ import numpy as np import torch +from micro_sam.util import device_type from micro_sam.v2.util import Devices from micro_sam.v1.prompt_based_segmentation import _process_box, _compute_logits_from_mask from micro_sam.v2.transforms.resize import resize_longest_side_and_pad_spatial_numpy, ResizeLongestSideTransforms @@ -534,7 +535,7 @@ def __init__( # Offloading frames/state to CPU bounds GPU memory for large volumes on CUDA. On MPS it is off # by default: unified memory saves nothing, and SAM2's CPU->MPS 'non_blocking' transfer of the # consolidated masks races, giving intermittent garbage/NaN masks (patchy interactive results). - is_mps = str(_get_device(device)) == "mps" + is_mps = device_type(_get_device(device)) == "mps" self.offload_video_to_cpu = (not is_mps) if offload_video_to_cpu is None else offload_video_to_cpu self.offload_state_to_cpu = (not is_mps) if offload_state_to_cpu is None else offload_state_to_cpu diff --git a/test/test_instance_segmentation.py b/test/test_instance_segmentation.py index 6d6190807..bb1c48198 100644 --- a/test/test_instance_segmentation.py +++ b/test/test_instance_segmentation.py @@ -1,3 +1,4 @@ +import types import unittest from copy import deepcopy @@ -189,5 +190,19 @@ def test_autoseg_base_hierarchy_and_contract(): assert isinstance(seg, AutoSegBase) and seg.is_initialized is False +@pytest.mark.parametrize( + "device, expected", [("mps", 16), ("mps:0", 16), ("cpu", 64), ("cuda", 64), ("cuda:1", 64)], +) +def test_automatic_mask_generator_points_per_batch_by_device(device, expected): + """MPS needs the smaller batch for performance, including when the device carries an index.""" + import torch + from micro_sam.v1.instance_segmentation import AutomaticMaskGenerator + + predictor = types.SimpleNamespace(device=torch.device(device)) + amg = AutomaticMaskGenerator(predictor) + + assert amg._points_per_batch == expected + + if __name__ == "__main__": unittest.main() diff --git a/test/test_util.py b/test/test_util.py index f11ffc51f..ff2e8b106 100644 --- a/test/test_util.py +++ b/test/test_util.py @@ -430,6 +430,19 @@ def test_get_device(self): self.assertTrue(isinstance(device, torch.device)) self.assertEqual(device.type, "cpu") + def test_device_type(self): + from micro_sam.util import device_type + + self.assertEqual(device_type("cpu"), "cpu") + self.assertEqual(device_type(torch.device("cpu")), "cpu") + + # Indexed accelerators must report the plain type: torch reports a model's parameters as + # living on 'mps:0' / 'cuda:0', so 'str(device) == "mps"' silently fails. + self.assertEqual(device_type("mps"), "mps") + self.assertEqual(device_type(torch.device("mps")), "mps") + self.assertEqual(device_type(torch.device("mps", 0)), "mps") + self.assertEqual(device_type(torch.device("cuda", 3)), "cuda") + def test_configure_mps_memory(self): from micro_sam.util import _configure_mps_memory diff --git a/test/test_v2_prompt_based_segmentation.py b/test/test_v2_prompt_based_segmentation.py index 70c43bd5a..4eccb3165 100644 --- a/test/test_v2_prompt_based_segmentation.py +++ b/test/test_v2_prompt_based_segmentation.py @@ -1,5 +1,6 @@ import numpy as np import pytest +import torch from bioimage_cpp.utils import Blocking from micro_sam.v2.prompt_based_segmentation import ( @@ -57,6 +58,32 @@ def get_segmenter(tile_id): return segmenter +@pytest.mark.parametrize("device", ("mps", torch.device("mps"), torch.device("mps", 0))) +def test_promptable_segmentation_3d_disables_offloading_on_mps(device, monkeypatch): + # Offloading to CPU on MPS races and gives garbage masks, so it must stay off for every way the + # device can be spelled - in particular the indexed 'mps:0' the tiled variant resolves. + monkeypatch.setattr("micro_sam.v2.util._get_device", lambda device=None: device) + monkeypatch.setattr(PromptableSegmentation3D, "init_predictor", lambda self: None) + segmenter = PromptableSegmentation3D( + predictor=None, volume=np.zeros((4, 8, 8), dtype="uint8"), volume_embeddings=None, device=device, + ) + + assert not segmenter.offload_video_to_cpu + assert not segmenter.offload_state_to_cpu + + +def test_promptable_segmentation_3d_keeps_offloading_on_cuda(monkeypatch): + monkeypatch.setattr("micro_sam.v2.util._get_device", lambda device=None: device) + monkeypatch.setattr(PromptableSegmentation3D, "init_predictor", lambda self: None) + segmenter = PromptableSegmentation3D( + predictor=None, volume=np.zeros((4, 8, 8), dtype="uint8"), volume_embeddings=None, + device=torch.device("cuda", 0), + ) + + assert segmenter.offload_video_to_cpu + assert segmenter.offload_state_to_cpu + + def test_promptable_segmentation_3d_progress_total(): segmenter = PromptableSegmentation3D.__new__(PromptableSegmentation3D) segmenter.volume = np.zeros((8, 16, 16), dtype="uint8") From 80165fa233fc1a76075076e7c864dcc258431d04 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Wed, 22 Jul 2026 23:56:59 +0200 Subject: [PATCH 10/11] Address review of batched multi-GPU inference --- micro_sam/_cli.py | 43 +- micro_sam/precompute_state.py | 20 +- .../sam_annotator/_batch_classification.py | 6 +- micro_sam/sam_annotator/_state.py | 6 +- micro_sam/sam_annotator/_widgets.py | 20 +- micro_sam/sam_annotator/annotator_tracking.py | 10 +- micro_sam/sam_annotator/batch_annotator.py | 12 +- micro_sam/sam_annotator/object_classifier.py | 4 + micro_sam/sam_annotator/pixel_classifier.py | 4 + micro_sam/v2/automatic_segmentation.py | 13 +- micro_sam/v2/batched_inference.py | 590 +++++++++++++----- micro_sam/v2/instance_segmentation.py | 73 ++- micro_sam/v2/models/_video_predictor.py | 24 +- micro_sam/v2/prompt_based_segmentation.py | 37 +- micro_sam/v2/util.py | 32 +- test/test_sam_annotator/test_widgets.py | 27 + test/test_v2_automatic_segmentation.py | 65 +- test/test_v2_batched_inference.py | 218 ++++++- 18 files changed, 950 insertions(+), 254 deletions(-) diff --git a/micro_sam/_cli.py b/micro_sam/_cli.py index c79cc5a5b..1eb2cbf8b 100644 --- a/micro_sam/_cli.py +++ b/micro_sam/_cli.py @@ -371,6 +371,24 @@ def _parse_shape(value): return tuple(int(x) for x in value.replace(" ", "").split(",")) +def _parse_batch_size(value): + """Parse the batch size: an integer, or 'auto' to benchmark a throughput-efficient value.""" + if value is None or str(value).lower() == "auto": + return None + try: + return int(value) + except ValueError: + raise click.UsageError(f"Expected an integer or 'auto' for '--batch_size', got '{value}'.") + + +def _parse_devices(value): + """Parse a comma-separated device list like 'cuda:0,cuda:1', or None for all visible devices.""" + if value is None: + return None + devices = [device.strip() for device in value.split(",") if device.strip()] + return devices or None + + def _convert_argval(value): """Best-effort conversion of a pass-through option value to int / float / str.""" try: @@ -472,12 +490,20 @@ def _view_result(image_path, key, segmentation): "-d", "--device", default=None, help="The device for the predictor: 'cuda', 'cpu' or 'mps'. By default the best available is used." ) +@click.option( + "--batch_size", default="1", + help="The number of tiles / slices per model call, or 'auto' to select it by measured throughput." +) +@click.option( + "--devices", default=None, + help="Comma-separated devices for inference, e.g. 'cuda:0,cuda:1'. By default all visible GPUs are used." +) @click.option("--view", is_flag=True, default=False, help="Whether to open the results in napari after segmentation.") @click.option("-v", "--verbose", is_flag=True, default=False, help="Whether to allow verbosity of outputs.") @click.pass_context def inference_segmentation( ctx, input_path, output_path, embedding_path, pattern, key, model_type, checkpoint_path, - tile_shape, halo, ndim, mode, device, view, verbose, + tile_shape, halo, ndim, mode, device, batch_size, devices, view, verbose, ): """Run automatic instance segmentation. @@ -497,6 +523,8 @@ def inference_segmentation( model_type = model_type or DEFAULT_MODEL tile_shape = _parse_shape(tile_shape) halo = _parse_shape(halo) + batch_size = _parse_batch_size(batch_size) + devices = _parse_devices(devices) generate_kwargs = _parse_extra(ctx.args) predictor, segmenter = get_predictor_and_segmenter( @@ -535,6 +563,8 @@ def inference_segmentation( halo=halo, mode=mode, device=device, + batch_size=batch_size, + devices=devices, verbose=verbose, **generate_kwargs, ) @@ -777,9 +807,17 @@ def inference_object_classification( "--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." ) +@click.option( + "--batch_size", default="1", + help="The number of tiles / slices per model call, or 'auto' to select it by measured throughput." +) +@click.option( + "--devices", default=None, + help="Comma-separated devices for inference, e.g. 'cuda:0,cuda:1'. By default all visible GPUs are used." +) def precompute_embeddings( input_path, embedding_path, pattern, key, model_type, checkpoint_path, ndim, - precompute_autoseg_state, prefer_decoder, + precompute_autoseg_state, prefer_decoder, batch_size, devices, ): """Precompute image embeddings (and optionally the automatic-segmentation state).""" from .precompute_state import precompute_state @@ -791,6 +829,7 @@ def precompute_embeddings( pattern=pattern, key=key, ndim=ndim, precompute_autoseg_state=precompute_autoseg_state, prefer_decoder=prefer_decoder, + batch_size=_parse_batch_size(batch_size), devices=_parse_devices(devices), ) diff --git a/micro_sam/precompute_state.py b/micro_sam/precompute_state.py index 96908e777..a13ded898 100644 --- a/micro_sam/precompute_state.py +++ b/micro_sam/precompute_state.py @@ -1,11 +1,12 @@ """Precompute and cache image embeddings for image data (SAM1, SAM2 or VFM encoders). """ +import inspect import os import pickle from glob import glob from pathlib import Path -from typing import Optional, Tuple, Union +from typing import Optional, Sequence, Tuple, Union import h5py import numpy as np @@ -664,6 +665,8 @@ def precompute_state( ndim: Optional[int] = None, precompute_autoseg_state: bool = False, prefer_decoder: bool = True, + batch_size: Optional[int] = 1, + devices: Optional[Union[str, Sequence[str]]] = None, ) -> None: """Precompute and cache the image embeddings (and, optionally, the automatic-segmentation state). @@ -691,6 +694,10 @@ def precompute_state( 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). + batch_size: The number of tiles / slices per model call. Pass None to select a throughput-efficient + value per device. Ignored by the model families that do not support batching (VFM encoders). + devices: The device or devices to compute the embeddings on. By default all visible CUDA devices + are used. Only supported for SAM2 ('hvit_*') models. """ # Imported lazily to avoid a circular import ('_state' imports from this module). from micro_sam.sam_annotator._state import _get_sam_model @@ -707,6 +714,14 @@ def precompute_state( # interface, so the per-family predictor from '_get_sam_model' feeds directly into it. compute_embeddings = util.get_embedding_function(model_type) + # Only SAM2 supports multi-device inference and only SAM1 / SAM2 support batching, so forward + # these settings just to the families whose embedding function accepts them. + supported = inspect.signature(compute_embeddings).parameters + compute_kwargs = { + name: value for name, value in (("batch_size", batch_size), ("devices", devices)) + if name in supported + } + # 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: @@ -741,7 +756,8 @@ def precompute_state( save_path = str(Path(out_path).with_suffix(".zarr")) embeddings = compute_embeddings( - predictor=predictor, input_=image_data, save_path=save_path, ndim=file_ndim, verbose=single + predictor=predictor, input_=image_data, save_path=save_path, ndim=file_ndim, verbose=single, + **compute_kwargs, ) if precompute_autoseg_state: diff --git a/micro_sam/sam_annotator/_batch_classification.py b/micro_sam/sam_annotator/_batch_classification.py index acde27f2b..131cc9c06 100644 --- a/micro_sam/sam_annotator/_batch_classification.py +++ b/micro_sam/sam_annotator/_batch_classification.py @@ -32,7 +32,7 @@ class ClassificationBatchTask(BatchAnnotatorTask): def __init__( self, *, ndim, model_type, embedding_paths=None, tile_shape=None, halo=None, - checkpoint_path=None, device=None, + checkpoint_path=None, device=None, batch_size=1, ): self.ndim = ndim self.model_type = model_type @@ -41,6 +41,7 @@ def __init__( self.halo = halo self.checkpoint_path = checkpoint_path self.device = device + self.batch_size = batch_size def _set_layers(self, viewer, index): """Hook: add or update task-specific layers (e.g. the segmentation layer) for this item.""" @@ -69,7 +70,8 @@ def _init_predictor(self, viewer, image, embedding_path, reuse): state.initialize_predictor( image, model_type=self.model_type, save_path=embedding_path, halo=self.halo, 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, + checkpoint_path=self.checkpoint_path, device=self.device, batch_size=self.batch_size, + skip_load=False, use_cli=True, **kwargs, ) state.image_shape = image.shape if image.ndim == self.ndim else image.shape[:-1] state.ndim = self.ndim diff --git a/micro_sam/sam_annotator/_state.py b/micro_sam/sam_annotator/_state.py index 7a02ec2e4..517f2cc9a 100644 --- a/micro_sam/sam_annotator/_state.py +++ b/micro_sam/sam_annotator/_state.py @@ -252,11 +252,11 @@ def initialize_predictor( save_path = util.make_temp_embedding_path() self.embedding_tmpdir = save_path - # For SAM2 volumes, load the embeddings lazily from the zarr so the high-resolution - # per-slice features stay on disk and are streamed one slice at a time during tracking. + # For SAM2 volumes and tiled images, load the embeddings lazily from the zarr so the + # high-resolution features stay on disk and are streamed one slice / tile at a time. # This keeps memory bounded for large volumes (materialising all slices costs # ~200 MB/slice and OOMs); it only applies when the embeddings are cached on disk. - lazy_loading = self.is_sam2 and ndim == 3 and isinstance(save_path, str) + lazy_loading = needs_disk_cache and isinstance(save_path, str) self.image_embeddings = _comp_embed_fn( predictor=self.predictor, diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index 6ffb7c35c..275448d46 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -408,6 +408,10 @@ def _update_model_type(self): # NOTE: And finally, we should re-enable signals again. self.model_size_dropdown.blockSignals(False) + # Whether the batch size applies depends on the backend, so it follows the model selection. + if getattr(self, "_batch_size_widget", None) is not None: + self._update_batch_size_visibility() + def _create_model_section( self, default_model: Optional[str] = None, @@ -1896,12 +1900,16 @@ def _update_tiling_visibility(self, index=None): self._tiling_widget.setVisible(self.tiling == "yes") def _update_batch_size_visibility(self, index=None): - # Show the batch size field only when the effective device is a GPU. Batching does not help - # on the CPU, so 'auto' resolving to the CPU (or an explicit CPU selection) hides the field. + # Show the batch size field only where it has an effect: batching does not help on the CPU + # (so 'auto' resolving to the CPU hides it), and the VFM encoders offered by the classifiers + # compute their embeddings unbatched. The default of 1 is used whenever it is hidden. + from micro_sam.models.vfm import is_vfm_model + device = self.device if device == "auto": device = util._get_default_device() - self._batch_size_widget.setVisible(str(device) != "cpu") + has_effect = str(device) != "cpu" and not is_vfm_model(getattr(self, "model_type", None)) + self._batch_size_widget.setVisible(has_effect) def _apply_default_tiling_for_shape(self, shape): # Enable tiling by default for large in-plane images, using the central v2 tiling defaults. @@ -2212,7 +2220,7 @@ def __call__(self, skip_validate=False): # Set up progress bar and signals for using it within a threadworker. # Model and decoder preparation happens before the backend knows the tile / slice count, so # start with an indeterminate status instead of leaving napari's activity display empty. - pbar, pbar_signals = _create_pbar_for_threadworker("Preparing image embeddings") + pbar, pbar_signals = _create_pbar_for_threadworker("Computing image embeddings") # @thread_worker() def compute_image_embedding(): @@ -2454,6 +2462,10 @@ def _update_model_type(self): self.model_size_dropdown.setCurrentText(self.model_size) self.model_size_dropdown.update() self.model_size_dropdown.blockSignals(False) + # This branch does not go through the inherited update, so refresh the batch size here too: + # the advanced families include the VFM encoders, which do not support batching. + if getattr(self, "_batch_size_widget", None) is not None: + self._update_batch_size_visibility() def _validate_model_type_and_custom_weights(self): # DINO families always resolve from the registry. DINOv3 weights load from HuggingFace using the diff --git a/micro_sam/sam_annotator/annotator_tracking.py b/micro_sam/sam_annotator/annotator_tracking.py index 5eabe1c0d..89d47b743 100644 --- a/micro_sam/sam_annotator/annotator_tracking.py +++ b/micro_sam/sam_annotator/annotator_tracking.py @@ -486,7 +486,7 @@ 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_autoseg_state=False, + precompute_autoseg_state=False, batch_size=1, ): _validate_tracking_model_type(model_type) self.model_type = model_type @@ -497,6 +497,7 @@ def __init__( self.decoder_path = decoder_path self.device = device self.precompute_autoseg_state = precompute_autoseg_state + self.batch_size = batch_size def result_filename(self, entry, index): if self.have_inputs_as_arrays: @@ -526,7 +527,7 @@ def _init_predictor(self, image, embedding_path, reuse): state.initialize_predictor( 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, + decoder_path=self.decoder_path, device=self.device, batch_size=self.batch_size, precompute_autoseg_state=self.precompute_autoseg_state, use_cli=True, **kwargs, ) @@ -579,6 +580,7 @@ def batch_tracking_annotator( viewer: Optional["napari.viewer.Viewer"] = None, return_viewer: bool = False, skip_done: bool = True, + batch_size: int = 1, ) -> Optional["napari.viewer.Viewer"]: """Run the tracking annotation tool for a batch of timeseries (each item is one TYX video). @@ -599,6 +601,8 @@ def batch_tracking_annotator( 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`. + batch_size: The number of tiles / slices per model call when computing the embeddings. + Only has an effect on a GPU. By default a single tile / slice is used. Returns: The napari viewer, only returned if `return_viewer=True`. @@ -607,7 +611,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_autoseg_state=precompute_autoseg_state, + precompute_autoseg_state=precompute_autoseg_state, batch_size=batch_size, ) 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 8f90b5a8b..0d49c045b 100644 --- a/micro_sam/sam_annotator/batch_annotator.py +++ b/micro_sam/sam_annotator/batch_annotator.py @@ -46,7 +46,7 @@ class SegmentationBatchTask(BatchAnnotatorTask): def __init__( self, *, ndim, model_type, embedding_path, tile_shape, halo, precompute_autoseg_state, checkpoint_path, device, prefer_decoder, - initial_segmentations=None, + initial_segmentations=None, batch_size=1, ): self.ndim = ndim self.model_type = model_type @@ -58,6 +58,7 @@ def __init__( self.device = device self.prefer_decoder = prefer_decoder self.initial_segmentations = initial_segmentations + self.batch_size = batch_size self.predictor = None self.decoder = None @@ -101,7 +102,7 @@ def _init_predictor(self, viewer, image, embedding_path): image, model_type=self.model_type, save_path=embedding_path, halo=self.halo, tile_shape=self.tile_shape, ndim=self.ndim, precompute_autoseg_state=self.precompute_autoseg_state, - checkpoint_path=self.checkpoint_path, + checkpoint_path=self.checkpoint_path, batch_size=self.batch_size, device=self.device, skip_load=False, use_cli=True, **kwargs, ) # Capture the loaded model so subsequent items reuse it instead of reloading. @@ -164,6 +165,7 @@ def batch_annotator( device: Optional[Union[str, torch.device]] = None, prefer_decoder: bool = True, skip_segmented: bool = True, + batch_size: int = 1, ) -> Optional["napari.viewer.Viewer"]: """Run the segmentation annotation tool for a batch of images (2d or 3d). @@ -192,6 +194,8 @@ def batch_annotator( prefer_decoder: Whether to use decoder based instance segmentation if the model used has an additional decoder for instance segmentation. By default, set to 'True'. + batch_size: The number of tiles / slices per model call when computing the embeddings. + Only has an effect on a GPU. By default a single tile / slice is used. skip_segmented: Whether existing output files mark images as completed. If True, resume at the first image without an output and skip any later completed images. If False, start at the first image and load existing segmentations into the 'committed_objects' layer. @@ -217,7 +221,7 @@ def batch_annotator( 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, + initial_segmentations=initial_segmentations, batch_size=batch_size, ) return run_batch( images, output_folder, task, have_inputs_as_arrays=have_inputs_as_arrays, @@ -524,7 +528,7 @@ def __call__(self, skip_validate=False): common = dict( model_type=ew.model_type, tile_shape=tile_shape, halo=halo, - checkpoint_path=ew.custom_weights, device=ew.device, + checkpoint_path=ew.custom_weights, device=ew.device, batch_size=ew.batch_size, viewer=self._viewer, return_viewer=True, ) diff --git a/micro_sam/sam_annotator/object_classifier.py b/micro_sam/sam_annotator/object_classifier.py index a32d0f376..de9a34dbc 100644 --- a/micro_sam/sam_annotator/object_classifier.py +++ b/micro_sam/sam_annotator/object_classifier.py @@ -321,6 +321,7 @@ def batch_object_classifier( viewer: Optional["napari.viewer.Viewer"] = None, return_viewer: bool = False, skip_done: bool = False, + batch_size: int = 1, ) -> Optional["napari.viewer.Viewer"]: """Start the object classifier for a list of images and segmentations. @@ -344,6 +345,8 @@ def batch_object_classifier( ndim: The dimensionality of the data. If not given will be derived from the data. viewer: The viewer to which the functionality should be added. return_viewer: Whether to return the napari viewer instead of starting the event loop. + batch_size: The number of tiles / slices per model call when computing the embeddings. + Only has an effect on a GPU. By default a single tile / slice is used. skip_done: Whether to skip images whose prediction already exists in `output_folder`. Returns: @@ -362,6 +365,7 @@ def batch_object_classifier( task = ObjectClassificationBatchTask( segmentations=segmentations, ndim=ndim, model_type=model_type, embedding_paths=embedding_paths, tile_shape=tile_shape, halo=halo, checkpoint_path=checkpoint_path, device=device, + batch_size=batch_size, ) return run_batch( images, output_folder, task, have_inputs_as_arrays=have_inputs_as_arrays, diff --git a/micro_sam/sam_annotator/pixel_classifier.py b/micro_sam/sam_annotator/pixel_classifier.py index af1f257cf..170960ea8 100644 --- a/micro_sam/sam_annotator/pixel_classifier.py +++ b/micro_sam/sam_annotator/pixel_classifier.py @@ -160,6 +160,7 @@ def batch_pixel_classifier( viewer: Optional["napari.viewer.Viewer"] = None, return_viewer: bool = False, skip_done: bool = False, + batch_size: int = 1, ) -> Optional["napari.viewer.Viewer"]: """Start the pixel classifier for a list of images. @@ -182,6 +183,8 @@ def batch_pixel_classifier( ndim: The dimensionality of the data. If not given will be derived from the data. viewer: The viewer to which the functionality should be added. return_viewer: Whether to return the napari viewer instead of starting the event loop. + batch_size: The number of tiles / slices per model call when computing the embeddings. + Only has an effect on a GPU. By default a single tile / slice is used. skip_done: Whether to skip images whose prediction already exists in `output_folder`. Returns: @@ -195,6 +198,7 @@ def batch_pixel_classifier( task = PixelClassificationBatchTask( ndim=ndim, model_type=model_type, embedding_paths=embedding_paths, tile_shape=tile_shape, halo=halo, checkpoint_path=checkpoint_path, device=device, + batch_size=batch_size, ) return run_batch( images, output_folder, task, have_inputs_as_arrays=have_inputs_as_arrays, diff --git a/micro_sam/v2/automatic_segmentation.py b/micro_sam/v2/automatic_segmentation.py index 1a73fc638..e9a1b143b 100644 --- a/micro_sam/v2/automatic_segmentation.py +++ b/micro_sam/v2/automatic_segmentation.py @@ -122,11 +122,11 @@ def automatic_instance_segmentation( mode: The AIS post-processing mode, 'sparse' (flow) or 'dense' (multicut). Ignored for AMG. device: The device to run inference on. verbose: Whether to print progress. - batch_size: Explicit tile or slice batch size. Defaults to one; pass None for - throughput-based automatic selection. + batch_size: The batch size used when running inference for multiple slices and / or tiles. + Defaults to one; pass None for throughput-based automatic selection. devices: Inference device or devices. None uses all visible GPUs when the model is on CUDA. num_prefetch_workers: Number of input reading and preprocessing threads. - num_write_workers: Number of output writing threads for full tiled inference. + num_write_workers: Number of output writing threads. generate_kwargs: Additional post-processing parameters forwarded to the segmenter's `generate`. Returns: @@ -163,10 +163,13 @@ def automatic_instance_segmentation( tile_shape=tile_shape, halo=halo, verbose=verbose, - lazy_loading=(ndim == 3), + # Volumes and tiled images are streamed from the zarr; only small 2d stays in memory. + lazy_loading=(ndim == 3 or tile_shape is not None), batch_size=batch_size, - devices=devices, + # Without an explicit selection, honor 'device' rather than fanning out over all GPUs. + devices=devices if devices is not None else device, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) segmenter.initialize( raw, diff --git a/micro_sam/v2/batched_inference.py b/micro_sam/v2/batched_inference.py index 5a5122f19..7c8c0380f 100644 --- a/micro_sam/v2/batched_inference.py +++ b/micro_sam/v2/batched_inference.py @@ -1,5 +1,6 @@ """Batched, pipelined, multi-GPU SAM2 inference: scheduling engine, encoder embeddings, and decoder passes.""" +import contextlib import gc import os import queue @@ -14,19 +15,19 @@ import numpy as np import torch -from micro_sam.util import _create_dataset_with_data, _create_dataset_without_data +from micro_sam.util import _create_dataset_without_data from .normalization import to_image -from .util import Device, Devices +from .util import Devices STOP = object() -class PipelineAborted(Exception): +class _PipelineAborted(Exception): """Raised inside workers when another pipeline worker has failed.""" -class AtomicCounter: +class _AtomicCounter: """Lock-guarded counter used to send completion sentinels exactly once.""" def __init__(self, value: int) -> None: @@ -40,7 +41,7 @@ def decrement(self) -> int: @dataclass -class PipelineJob: +class _PipelineJob: """A work item moving from the input loader to inference and output writing.""" spec: Any @@ -61,12 +62,23 @@ def _model_device(model: torch.nn.Module) -> torch.device: return torch.device("cpu") -def resolve_devices(model: torch.nn.Module, devices: Devices = None) -> List[torch.device]: - """Resolve inference devices, using every visible CUDA device by default. +def _resolve_devices(model: torch.nn.Module, devices: Devices = None) -> List[torch.device]: + """Resolve the inference devices, using every visible CUDA device by default. Automatic multi-GPU execution is enabled only when the supplied model already lives on CUDA. - This preserves an explicitly CPU- or MPS-loaded model. Pass a scalar device to force one device, - or a sequence to select an explicit set of devices. + This preserves an explicitly CPU- or MPS-loaded model. + + Args: + model: The model whose device decides the default. Only used when `devices` is None. + devices: A single device to force one device, or a sequence to select an explicit set. + By default all visible CUDA devices are used if the model is on CUDA. + + Returns: + The resolved devices, in the order they will be assigned to inference workers. + + Raises: + ValueError: If no device is given, or if the same device is given more than once. + RuntimeError: If a CUDA device is requested but CUDA is unavailable. """ if devices is None: device = _model_device(model) @@ -88,10 +100,19 @@ def resolve_devices(model: torch.nn.Module, devices: Devices = None) -> List[tor return resolved -def prepare_models( +def _prepare_models( model: torch.nn.Module, devices: Sequence[torch.device], ) -> List[Tuple[torch.nn.Module, torch.device]]: - """Create one eval-mode model replica per device, reusing the original where possible.""" + """Create one eval-mode model replica per device, reusing the original where possible. + + Args: + model: The model to replicate. + devices: The devices to place the replicas on (see `_resolve_devices`). + + Returns: + One (model, device) pair per device. The pair for the model's own device holds the original + model, all others hold a deep copy. Release them with `_release_model_replicas`. + """ source_device = _model_device(model) models = [] for device in devices: @@ -102,11 +123,14 @@ def prepare_models( return models -def release_model_replicas(model_devices: List[Tuple[torch.nn.Module, torch.device]]) -> None: - """Drop the replicas built by :func:`prepare_models` and free their GPU memory. +def _release_model_replicas(model_devices: List[Tuple[torch.nn.Module, torch.device]]) -> None: + """Drop the replicas built by `_prepare_models` and free their GPU memory. Clearing the list releases only the per-device copies; the caller's own model stays alive through its other references. + + Args: + model_devices: The (model, device) pairs returned by `_prepare_models`. Cleared in place. """ model_devices.clear() gc.collect() @@ -192,7 +216,7 @@ def _select_throughput_batch_size( return int(selected_batch) -def compute_auto_batch_sizes( +def _compute_auto_batch_sizes( model_devices: Sequence[Tuple[torch.nn.Module, torch.device]], n_jobs: int, patch_shape: Tuple[int, ...], @@ -204,6 +228,19 @@ def compute_auto_batch_sizes( Powers of two are benchmarked from one upwards. A larger batch is selected only when it improves measured throughput by at least 10 percent, and probing stops after two consecutive candidates fail to do so. CPU and MPS retain the conservative batch size of one. + + Args: + model_devices: The (model, device) pairs to benchmark (see `_prepare_models`). + n_jobs: The total number of jobs; caps the candidates so a batch never exceeds the work. + patch_shape: The spatial shape of a single input to the model. + in_channels: The number of channels of a single input to the model. + prediction_function: Called as `prediction_function(model, inputs)` to run one probe batch. + + Returns: + One batch size per entry in `model_devices`. + + Raises: + RuntimeError: If a device runs out of memory already at batch size one. """ if n_jobs < 1: return [1] * len(model_devices) @@ -258,7 +295,7 @@ def _safe_get(input_queue: queue.Queue, stop_event: threading.Event, timeout: fl return input_queue.get(timeout=timeout) except queue.Empty: continue - raise PipelineAborted() + raise _PipelineAborted() def _safe_put(output_queue: queue.Queue, item: Any, stop_event: threading.Event, timeout: float = 0.2) -> None: @@ -268,7 +305,7 @@ def _safe_put(output_queue: queue.Queue, item: Any, stop_event: threading.Event, return except queue.Full: continue - raise PipelineAborted() + raise _PipelineAborted() def _predict_with_oom_backoff( @@ -298,7 +335,7 @@ def _predict_with_oom_backoff( return left + right, min(left_size, right_size) -def run_batched_pipeline( +def _run_batched_pipeline( jobs: Iterable[Any], model_devices: Sequence[Tuple[torch.nn.Module, torch.device]], batch_sizes: Sequence[int], @@ -306,13 +343,39 @@ def run_batched_pipeline( predict_fn: Callable[[torch.nn.Module, List[Any], torch.device], List[Any]], write_fn: Callable[[Any, Any], None], num_prefetch_workers: int = 4, + num_write_workers: int = 1, update_progress: Optional[Callable[[int], None]] = None, progress_increment: Optional[Callable[[Any], int]] = None, ) -> None: """Run load/preprocess, batched inference, and writing as a bounded threaded pipeline. - There is one inference consumer per device and one writer. Keeping output writes on one thread - is safe for zarr, N5, HDF5, and in-memory outputs, while still overlapping I/O with GPU work. + Loading runs on `num_prefetch_workers` threads, inference on one thread per device, and writing + on `num_write_workers` threads, so input I/O and output I/O overlap with the model forward pass. + The queues between the stages are bounded, so a slow stage throttles the others instead of + growing memory. Jobs are written in completion order, not in input order. If any worker raises, + the whole pipeline is stopped and the first error is re-raised on the calling thread. + + Args: + jobs: The job specifications, e.g. tile ids or slice indices. Passed to `load_fn` and, with + the prediction, to `write_fn`. + model_devices: The (model, device) pairs to run inference on (see `_prepare_models`). + batch_sizes: One batch size per entry in `model_devices` (see `_compute_auto_batch_sizes`). + load_fn: Called as `load_fn(job)` to read and preprocess one job. + predict_fn: Called as `predict_fn(model, items, device)` with the loaded data of one batch. + It must return one output per input. Batches that run out of memory are automatically + retried in halves, and the device's batch size is reduced accordingly. + write_fn: Called as `write_fn(job, prediction)` to store one result. It must be safe to call + from multiple threads if `num_write_workers` is larger than one. + num_prefetch_workers: The number of threads used to read and preprocess jobs. + num_write_workers: The number of threads used to write results. Multiple writers speed up + compressed / chunked outputs; keep it at one for outputs that are not thread-safe. + update_progress: Optional callback advancing external progress by the given number of steps. + It is called on the calling thread, so it is safe to use for Qt / napari progress bars. + progress_increment: Optional callback returning the number of steps a job contributes. + By default every job counts as one step. + + Raises: + ValueError: If the batch sizes do not match the devices, or are not positive. """ jobs = list(jobs) if len(jobs) == 0: @@ -323,6 +386,7 @@ def run_batched_pipeline( raise ValueError(f"Batch sizes must be positive, got {batch_sizes}.") num_prefetch_workers = max(1, min(int(num_prefetch_workers), len(jobs))) + num_write_workers = max(1, min(int(num_write_workers), len(jobs))) batch_sizes = [int(batch_size) for batch_size in batch_sizes] job_queue = queue.Queue() @@ -332,13 +396,13 @@ def run_batched_pipeline( job_queue.put(STOP) input_queue = queue.Queue(maxsize=max(2 * sum(batch_sizes), 2)) - output_queue = queue.Queue(maxsize=max(2 * len(model_devices), 2)) + output_queue = queue.Queue(maxsize=max(2 * len(model_devices), 2 * num_write_workers, 2)) progress_queue = queue.Queue() stop_event = threading.Event() error_box = [] error_lock = threading.Lock() - remaining_producers = AtomicCounter(num_prefetch_workers) - remaining_consumers = AtomicCounter(len(model_devices)) + remaining_producers = _AtomicCounter(num_prefetch_workers) + remaining_consumers = _AtomicCounter(len(model_devices)) def record_error(exc): with error_lock: @@ -352,15 +416,18 @@ def producer(): spec = job_queue.get() if spec is STOP or stop_event.is_set(): break - _safe_put(input_queue, PipelineJob(spec, load_fn(spec)), stop_event) - except PipelineAborted: + _safe_put(input_queue, _PipelineJob(spec, load_fn(spec)), stop_event) + except _PipelineAborted: pass except Exception as exc: # noqa record_error(exc) finally: - if remaining_producers.decrement() == 0 and not stop_event.is_set(): - for _ in model_devices: - input_queue.put(STOP) + # A timed put: a consumer that fails while the queue is full leaves nobody to drain it, + # and a blocking put would then hang this thread (and the join in the outer cleanup). + if remaining_producers.decrement() == 0: + with contextlib.suppress(_PipelineAborted): + for _ in model_devices: + _safe_put(input_queue, STOP, stop_event) def consumer(worker_id): model, device = model_devices[worker_id] @@ -392,13 +459,16 @@ def consumer(worker_id): if got_stop: break - except PipelineAborted: + except _PipelineAborted: pass except Exception as exc: # noqa record_error(exc) finally: - if remaining_consumers.decrement() == 0 and not stop_event.is_set(): - output_queue.put(STOP) + # Timed as well, for the same reason: the writers may already have failed and stopped. + if remaining_consumers.decrement() == 0: + with contextlib.suppress(_PipelineAborted): + for _ in range(num_write_workers): + _safe_put(output_queue, STOP, stop_event) def writer(): try: @@ -410,12 +480,15 @@ def writer(): if update_progress is not None: increment = 1 if progress_increment is None else int(progress_increment(item.spec)) progress_queue.put(increment) - except PipelineAborted: + except _PipelineAborted: pass except Exception as exc: # noqa record_error(exc) - writer_thread = threading.Thread(target=writer, name="sam2-writer") + writer_threads = [ + threading.Thread(target=writer, name=f"sam2-writer-{worker_id}") + for worker_id in range(num_write_workers) + ] consumer_threads = [ threading.Thread(target=consumer, args=(worker_id,), name=f"sam2-consumer-{worker_id}") for worker_id in range(len(model_devices)) @@ -424,7 +497,7 @@ def writer(): threading.Thread(target=producer, name=f"sam2-producer-{worker_id}") for worker_id in range(num_prefetch_workers) ] - threads = [writer_thread, *consumer_threads, *producer_threads] + threads = [*writer_threads, *consumer_threads, *producer_threads] def forward_progress(block): increments = 0 @@ -462,10 +535,6 @@ def forward_progress(block): raise error_box[0] -IMAGE_MEAN = (0.485, 0.456, 0.406) -IMAGE_STD = (0.229, 0.224, 0.225) - - def _clear_group(group) -> None: for key in list(group.keys()): del group[key] @@ -479,11 +548,11 @@ def _prepare_encoder_pipeline( predictor, n_jobs: int, batch_size: Optional[int], devices: Devices, ) -> Tuple[torch.nn.Module, List[Tuple[torch.nn.Module, torch.device]], List[int]]: model = getattr(predictor, "model", predictor) - resolved_devices = resolve_devices(model, devices) - model_devices = prepare_models(model, resolved_devices) + resolved_devices = _resolve_devices(model, devices) + model_devices = _prepare_models(model, resolved_devices) if batch_size is None: image_size = int(getattr(model, "image_size", getattr(predictor, "image_size", 1024))) - batch_sizes = compute_auto_batch_sizes( + batch_sizes = _compute_auto_batch_sizes( model_devices=model_devices, n_jobs=n_jobs, patch_shape=(image_size, image_size), @@ -522,7 +591,7 @@ def _forward_image_batch( ] -def compute_tiled_2d( +def _compute_tiled_2d( input_: np.ndarray, predictor, tile_shape: Tuple[int, int], @@ -534,8 +603,30 @@ def compute_tiled_2d( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> Dict: - """Compute 2D tile embeddings with batched encoders and queued input/output I/O.""" + """Compute 2d tile embeddings with batched encoders and queued input / output I/O. + + Args: + input_: The input image, shape (Y, X) or (Y, X, C). + predictor: The SAM2 image predictor. + tile_shape: The in-plane tile shape. + halo: The in-plane tile halo (overlap). + root: The zarr container the embeddings are written to (see `micro_sam.util._open_embeddings`). + save_path: The path backing `root`, or None for an in-memory container. A complete cache at + this path is returned as is instead of being recomputed. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + batch_size: The number of tiles per encoder call. By default candidate sizes are benchmarked + and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read and preprocess tiles. + num_write_workers: The number of threads used to write embedding tiles. + + Returns: + The tiled image embeddings. 'features' and 'high_res_feats' are the zarr groups holding the + per-tile datasets; 'input_size' and 'original_size' are None because they are stored per tile. + """ from bioimage_cpp.utils import Blocking from micro_sam.v2.util import _write_embedding_signature @@ -567,18 +658,33 @@ def load_tile(tile_id): def predict_tiles(this_model, items, device): return _forward_image_batch(this_model, items, device, feature_sizes) + # Creating a dataset mutates the shared group, so it is serialized; the data (and with it the + # compression) is written outside the lock, which is what multiple write workers speed up. + creation_lock = threading.Lock() + def write_tile(tile_id, result): name = str(tile_id) - dataset = _create_dataset_with_data(features, name, data=result["features"]) - dataset.attrs["input_size"] = model.image_size - dataset.attrs["original_size"] = [list(result["original_size"])] - - tile_high_res = high_res_group.require_group(name) - for level, feature in enumerate(result["high_res_feats"]): - _create_dataset_with_data(tile_high_res, str(level), data=feature) + tile_features = result["features"] + high_res_feats = result["high_res_feats"] + with creation_lock: + dataset = _create_dataset_without_data( + features, name, shape=tile_features.shape, dtype=tile_features.dtype, chunks=tile_features.shape, + ) + dataset.attrs["input_size"] = model.image_size + dataset.attrs["original_size"] = [list(result["original_size"])] + tile_high_res = high_res_group.require_group(name) + high_res_datasets = [ + _create_dataset_without_data( + tile_high_res, str(level), shape=feature.shape, dtype=feature.dtype, chunks=feature.shape, + ) for level, feature in enumerate(high_res_feats) + ] + + dataset[:] = tile_features + for high_res_dataset, feature in zip(high_res_datasets, high_res_feats): + high_res_dataset[:] = feature try: - run_batched_pipeline( + _run_batched_pipeline( jobs=range(n_tiles), model_devices=model_devices, batch_sizes=batch_sizes, @@ -586,10 +692,11 @@ def write_tile(tile_id, result): predict_fn=predict_tiles, write_fn=write_tile, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, update_progress=pbar_update, ) finally: - release_model_replicas(model_devices) + _release_model_replicas(model_devices) if save_path is not None: _write_embedding_signature( @@ -600,12 +707,9 @@ def write_tile(tile_id, result): def _prepare_video_frame(raw: np.ndarray, image_size: int) -> torch.Tensor: - from micro_sam.v2.models._video_predictor import _load_img_as_tensor + from micro_sam.v2.models._video_predictor import _prepare_frame - image, _, _ = _load_img_as_tensor(np.asarray(raw), image_size) - mean = torch.tensor(IMAGE_MEAN, dtype=torch.float32)[:, None, None] - std = torch.tensor(IMAGE_STD, dtype=torch.float32)[:, None, None] - return (image - mean) / std + return _prepare_frame(np.asarray(raw), image_size) def _forward_video_batch(model: torch.nn.Module, items: List[Dict], device: torch.device) -> List[Dict]: @@ -649,7 +753,7 @@ def _load_feature_levels(group, lazy_loading: bool) -> List: return values -def compute_3d( +def _compute_3d( input_: np.ndarray, predictor, root, @@ -660,8 +764,30 @@ def compute_3d( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> Dict: - """Compute volume embeddings by batching slices and overlapping preprocessing and zarr writes.""" + """Compute volume embeddings by batching slices and overlapping preprocessing and zarr writes. + + Args: + input_: The input volume, shape (Z, Y, X). + predictor: The SAM2 video predictor. + root: The zarr container the embeddings are written to (see `micro_sam.util._open_embeddings`). + save_path: The path backing `root`, or None to keep the embeddings in memory. A complete cache + at this path is returned as is instead of being recomputed. + lazy_loading: Whether to return the zarr datasets instead of materializing the embeddings. + Only has an effect if `save_path` is given. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + batch_size: The number of slices per encoder call. By default candidate sizes are benchmarked + and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read and preprocess slices. + num_write_workers: The number of threads used to write embedding slices. + + Returns: + The volume embeddings, with the per-slice 'features', 'pos_enc' and 'fpn' outputs of the + image encoder as well as 'input_size' and 'original_size'. + """ from micro_sam.v2.util import _write_embedding_signature if save_path is not None and "original_size" in root.attrs: @@ -703,6 +829,10 @@ def load_slice(z): "original_size": tuple(int(value) for value in raw.shape[:2]), } + # The datasets are created from the first result, so only their creation is serialized; the + # per-slice writes run in parallel (they go to separate chunks). + creation_lock = threading.Lock() + def write_slice(z, result): nonlocal feature_dataset, pos_datasets, fpn_datasets if save_path is None: @@ -711,10 +841,11 @@ def write_slice(z, result): fpn_values[z] = [torch.from_numpy(value) for value in result["fpn"]] return - if feature_dataset is None: - feature_dataset = _create_feature_dataset(root, "features", n_slices, result["features"]) - pos_datasets = _create_feature_levels(root.require_group("pos_enc"), n_slices, result["pos_enc"]) - fpn_datasets = _create_feature_levels(root.require_group("fpn"), n_slices, result["fpn"]) + with creation_lock: + if feature_dataset is None: + feature_dataset = _create_feature_dataset(root, "features", n_slices, result["features"]) + pos_datasets = _create_feature_levels(root.require_group("pos_enc"), n_slices, result["pos_enc"]) + fpn_datasets = _create_feature_levels(root.require_group("fpn"), n_slices, result["fpn"]) feature_dataset[z] = result["features"] for dataset, value in zip(pos_datasets, result["pos_enc"]): dataset[z] = value @@ -722,7 +853,7 @@ def write_slice(z, result): dataset[z] = value try: - run_batched_pipeline( + _run_batched_pipeline( jobs=range(n_slices), model_devices=model_devices, batch_sizes=batch_sizes, @@ -730,10 +861,11 @@ def write_slice(z, result): predict_fn=_forward_video_batch, write_fn=write_slice, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, update_progress=pbar_update, ) finally: - release_model_replicas(model_devices) + _release_model_replicas(model_devices) original_size = tuple(int(value) for value in input_.shape[-2:]) if save_path is None: @@ -759,7 +891,7 @@ def write_slice(z, result): } -def compute_tiled_3d( +def _compute_tiled_3d( input_: np.ndarray, predictor, tile_shape: Tuple[int, int], @@ -771,8 +903,33 @@ def compute_tiled_3d( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> Dict: - """Compute tile/slice embeddings as one pipelined job stream across all available GPUs.""" + """Compute tile / slice embeddings as one pipelined job stream across all available GPUs. + + Every (tile, slice) pair is a separate job, so tiles and slices are batched together instead of + tile column by tile column. This keeps all devices busy even for few tiles or few slices. + + Args: + input_: The input volume, shape (Z, Y, X). + predictor: The SAM2 video predictor. + tile_shape: The in-plane tile shape. The volume is not tiled along z. + halo: The in-plane tile halo (overlap). + root: The zarr container the embeddings are written to (see `micro_sam.util._open_embeddings`). + save_path: The path backing `root`, or None for an in-memory container. A complete cache at + this path is returned as is instead of being recomputed. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + batch_size: The number of tile slices per encoder call. By default candidate sizes are + benchmarked and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read and preprocess tile slices. + num_write_workers: The number of threads used to write embedding tile slices. + + Returns: + The tiled volume embeddings. 'features', 'pos_enc' and 'fpn' are the zarr groups holding the + per-tile datasets; 'input_size' and 'original_size' are None because they are stored per tile. + """ from bioimage_cpp.utils import Blocking from micro_sam.v2.util import _write_embedding_signature @@ -821,20 +978,25 @@ def load_tile_slice(job): "original_size": tuple(int(value) for value in raw.shape[:2]), } + # A tile's datasets are created from its first slice, so only their creation is serialized; the + # per-slice writes run in parallel (they go to separate chunks). + creation_lock = threading.Lock() + def write_tile_slice(job, result): tile_id, z = job - if tile_id not in tile_datasets: - name = str(tile_id) - feature_dataset = _create_feature_dataset(features, name, n_slices, result["features"]) - feature_dataset.attrs["input_size"] = image_size - feature_dataset.attrs["original_size"] = list(result["original_size"]) - pos_datasets = _create_feature_levels( - pos_enc_group.require_group(name), n_slices, result["pos_enc"], - ) - fpn_datasets = _create_feature_levels( - fpn_group.require_group(name), n_slices, result["fpn"], - ) - tile_datasets[tile_id] = feature_dataset, pos_datasets, fpn_datasets + with creation_lock: + if tile_id not in tile_datasets: + name = str(tile_id) + feature_dataset = _create_feature_dataset(features, name, n_slices, result["features"]) + feature_dataset.attrs["input_size"] = image_size + feature_dataset.attrs["original_size"] = list(result["original_size"]) + pos_datasets = _create_feature_levels( + pos_enc_group.require_group(name), n_slices, result["pos_enc"], + ) + fpn_datasets = _create_feature_levels( + fpn_group.require_group(name), n_slices, result["fpn"], + ) + tile_datasets[tile_id] = feature_dataset, pos_datasets, fpn_datasets feature_dataset, pos_datasets, fpn_datasets = tile_datasets[tile_id] feature_dataset[z] = result["features"] @@ -844,7 +1006,7 @@ def write_tile_slice(job, result): dataset[z] = value try: - run_batched_pipeline( + _run_batched_pipeline( jobs=jobs, model_devices=model_devices, batch_sizes=batch_sizes, @@ -852,10 +1014,11 @@ def write_tile_slice(job, result): predict_fn=_forward_video_batch, write_fn=write_tile_slice, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, update_progress=pbar_update, ) finally: - release_model_replicas(model_devices) + _release_model_replicas(model_devices) if save_path is not None: _write_embedding_signature( @@ -871,7 +1034,7 @@ def write_tile_slice(job, result): } -class EmptyEncoder(torch.nn.Module): +class _EmptyEncoder(torch.nn.Module): """Lightweight placeholder used while replicating decoder-only UniSAM2 models.""" def __init__(self, img_size: int) -> None: @@ -935,42 +1098,52 @@ def _feature_shape(job: Dict) -> Tuple[int, ...]: return shape -def _prepare_decoder_pipeline( - model: torch.nn.Module, jobs: List[Dict], batch_size: Optional[int], devices: Devices, -) -> Tuple[List[Tuple[torch.nn.Module, torch.device]], List[int], torch.nn.Module]: - resolved_devices = resolve_devices(model, devices) +@contextlib.contextmanager +def _decoder_only(model: torch.nn.Module): + """Replace the model's image encoder for the duration of the block, so only the decoder is run. + + The encoder is not needed when decoding precomputed embeddings, and replicating it per device + would waste memory. Restoring it is a context manager (rather than scattered cleanup) so the + caller's model gets its encoder back on every exit path. + """ encoder = model.encoder - model.encoder = EmptyEncoder(getattr(encoder, "img_size", 1024)) + model.encoder = _EmptyEncoder(getattr(encoder, "img_size", 1024)) try: - model_devices = prepare_models(model, resolved_devices) - except Exception: + yield + finally: model.encoder = encoder - raise - if batch_size is not None: - if int(batch_size) < 1: - release_model_replicas(model_devices) - model.encoder = encoder - raise ValueError(f"batch_size must be positive or None, got {batch_size}.") - return model_devices, [int(batch_size)] * len(model_devices), encoder - representative = max( - jobs, - key=lambda job: np.prod(_feature_shape(job)) * np.prod(job["original_size"]), - ) - z, channels, height, width = _feature_shape(representative) - original_size = representative["original_size"] +def _prepare_decoder_pipeline( + model: torch.nn.Module, + jobs: List[Dict], + batch_size: Optional[int], + resolved_devices: Sequence[torch.device], +) -> Tuple[List[Tuple[torch.nn.Module, torch.device]], List[int]]: + """Replicate the decoder and select its batch size. Must run inside `_decoder_only`.""" + model_devices = _prepare_models(model, resolved_devices) + try: + if batch_size is not None: + if int(batch_size) < 1: + raise ValueError(f"batch_size must be positive or None, got {batch_size}.") + return model_devices, [int(batch_size)] * len(model_devices) + + representative = max( + jobs, + key=lambda job: np.prod(_feature_shape(job)) * np.prod(job["original_size"]), + ) + z, channels, height, width = _feature_shape(representative) + original_size = representative["original_size"] - def prediction_function(this_model, inputs): - from micro_sam.v2.instance_segmentation import _decode_3d_feature_batch + def prediction_function(this_model, inputs): + from micro_sam.v2.instance_segmentation import _decode_3d_feature_batch - features = inputs.permute(0, 2, 1, 3, 4) - return _decode_3d_feature_batch( - this_model, features, original_size, inputs.device, - ) + features = inputs.permute(0, 2, 1, 3, 4) + return _decode_3d_feature_batch( + this_model, features, original_size, inputs.device, + ) - try: - batch_sizes = compute_auto_batch_sizes( + batch_sizes = _compute_auto_batch_sizes( model_devices=model_devices, n_jobs=len(jobs), patch_shape=(z, height, width), @@ -978,10 +1151,9 @@ def prediction_function(this_model, inputs): prediction_function=prediction_function, ) except Exception: - release_model_replicas(model_devices) - model.encoder = encoder + _release_model_replicas(model_devices) raise - return model_devices, batch_sizes, encoder + return model_devices, batch_sizes def _run_decoder_jobs( @@ -991,44 +1163,58 @@ def _run_decoder_jobs( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, update_progress: Optional[Callable[[int], None]] = None, progress_increment: Optional[Callable[[Dict], int]] = None, ) -> None: if len(jobs) == 0: return - model_devices, batch_sizes, encoder = _prepare_decoder_pipeline(model, jobs, batch_size, devices) - groups = defaultdict(list) - for job in jobs: - groups[(_feature_shape(job), job["original_size"])].append(job) + # Resolved before the encoder is swapped out, so the device is read off the full model. + resolved_devices = _resolve_devices(model, devices) + with _decoder_only(model): + groups = defaultdict(list) + for job in jobs: + groups[(_feature_shape(job), job["original_size"])].append(job) - # Run the largest shape first so later, smaller groups can reuse its allocator blocks. Starting - # with a boundary z-block and growing to the full context fragmented 10 GB MIG allocators. - group_keys = sorted(groups, key=lambda key: np.prod(key[0]) * np.prod(key[1]), reverse=True) + # Run the largest shape first so later, smaller groups can reuse its allocator blocks. Starting + # with a boundary z-block and growing to the full context fragmented 10 GB MIG allocators. + group_keys = sorted(groups, key=lambda key: np.prod(key[0]) * np.prod(key[1]), reverse=True) - try: - for group_key in group_keys: - group_jobs = groups[group_key] - run_batched_pipeline( - jobs=group_jobs, - model_devices=model_devices, - batch_sizes=batch_sizes, - load_fn=_load_job, - predict_fn=_predict_jobs, - write_fn=write_fn, - num_prefetch_workers=num_prefetch_workers, - update_progress=update_progress, - progress_increment=progress_increment, - ) - finally: - release_model_replicas(model_devices) - model.encoder = encoder + model_devices, batch_sizes = _prepare_decoder_pipeline(model, jobs, batch_size, resolved_devices) + try: + for group_key in group_keys: + group_jobs = groups[group_key] + _run_batched_pipeline( + jobs=group_jobs, + model_devices=model_devices, + batch_sizes=batch_sizes, + load_fn=_load_job, + predict_fn=_predict_jobs, + write_fn=write_fn, + num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, + update_progress=update_progress, + progress_increment=progress_increment, + ) + finally: + _release_model_replicas(model_devices) + + +def _resolve_z_blocking(z_block: Optional[int], z_halo: Optional[int]) -> Tuple[int, int]: + """Resolve the z block / halo used to decode a volume, falling back to the defaults.""" + from micro_sam.v2.util import DEFAULT_HALO_Z, DEFAULT_TILE_Z + + z_block = DEFAULT_TILE_Z if z_block is None else int(z_block) + z_halo = DEFAULT_HALO_Z if z_halo is None else int(z_halo) + if z_block < 1 or z_halo < 0: + raise ValueError(f"z_block must be positive and z_halo non-negative, got {z_block}, {z_halo}.") + return z_block, z_halo -def decode_volume_embeddings( +def _decode_volume_embeddings( model: torch.nn.Module, image_embeddings: Dict, - device: Device = None, z_block: Optional[int] = None, z_halo: Optional[int] = None, pbar_init: Optional[Callable] = None, @@ -1036,10 +1222,33 @@ def decode_volume_embeddings( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> np.ndarray: - """Decode z blocks from non-tiled volume embeddings with queued reads and batched inference.""" - from micro_sam.v2.util import DEFAULT_HALO_Z, DEFAULT_TILE_Z - + """Decode z blocks from non-tiled volume embeddings with queued reads and batched inference. + + The volume is decoded in blocks of `z_block` slices, each extended by `z_halo` context slices + that are cropped away again when the block is written. + + Args: + model: The UniSAM2 model. Its encoder is bypassed, only the decoder is run. + image_embeddings: The volume embeddings (see `_compute_3d`). + device: Unused, kept for interface compatibility. Use `devices` instead. + z_block: The number of slices decoded per block. By default `micro_sam.v2.util.DEFAULT_TILE_Z`. + z_halo: The number of context slices per block. By default `micro_sam.v2.util.DEFAULT_HALO_Z`. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + batch_size: The number of z blocks per decoder call. By default candidate sizes are + benchmarked and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read embedding blocks. + num_write_workers: The number of threads used to write decoded blocks. + + Returns: + The decoder predictions, shape (4, Z, Y, X): foreground and the three distance channels. + + Raises: + ValueError: If the features are not 3d, if `z_block` is not positive or `z_halo` is negative. + """ features = image_embeddings["features"] n_dims = len(features.shape) if n_dims not in (4, 5): @@ -1048,10 +1257,7 @@ def decode_volume_embeddings( ) n_slices = int(features.shape[0]) - z_block = DEFAULT_TILE_Z if z_block is None else int(z_block) - z_halo = DEFAULT_HALO_Z if z_halo is None else int(z_halo) - if z_block < 1 or z_halo < 0: - raise ValueError(f"z_block must be positive and z_halo non-negative, got {z_block}, {z_halo}.") + z_block, z_halo = _resolve_z_blocking(z_block, z_halo) original_size = tuple(int(value) for value in np.asarray(image_embeddings["original_size"]).reshape(-1)[:2]) output = np.zeros((4, n_slices, *original_size), dtype="float32") @@ -1078,7 +1284,7 @@ def write_prediction(job, prediction): _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, - num_prefetch_workers=num_prefetch_workers, + num_prefetch_workers=num_prefetch_workers, num_write_workers=num_write_workers, update_progress=pbar_update, progress_increment=lambda job: job["z1"] - job["z0"], ) return output @@ -1096,17 +1302,32 @@ def _tiled_metadata(image_embeddings: Dict, is_3d: bool) -> Tuple: return features, shape, halo, tiling -def decode_tiled_2d_embeddings( +def _decode_tiled_2d_embeddings( model: torch.nn.Module, image_embeddings: Dict, - device: Device = None, pbar_init: Optional[Callable] = None, pbar_update: Optional[Callable] = None, batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> np.ndarray: - """Decode and stitch 2D embedding tiles in automatically sized batches.""" + """Decode and stitch 2d embedding tiles in automatically sized batches. + + Args: + model: The UniSAM2 model. Its encoder is bypassed, only the decoder is run. + image_embeddings: The tiled image embeddings (see `_compute_tiled_2d`). + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + batch_size: The number of tiles per decoder call. By default candidate sizes are benchmarked + and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read embedding tiles. + num_write_workers: The number of threads used to stitch decoded tiles into the output. + + Returns: + The stitched decoder predictions, shape (4, Y, X): foreground and the three distance channels. + """ features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=False) output = np.zeros((4, *shape), dtype="float32") jobs = [] @@ -1135,7 +1356,7 @@ def write_prediction(job, prediction): _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, - num_prefetch_workers=num_prefetch_workers, + num_prefetch_workers=num_prefetch_workers, num_write_workers=num_write_workers, update_progress=pbar_update, ) return output @@ -1161,10 +1382,9 @@ def _tiled_3d_jobs(features, tiling, n_slices: int, z_block: int, z_halo: int) - return jobs -def decode_tiled_3d_embeddings( +def _decode_tiled_3d_embeddings( model: torch.nn.Module, image_embeddings: Dict, - device: Device = None, pbar_init: Optional[Callable] = None, pbar_update: Optional[Callable] = None, z_block: Optional[int] = None, @@ -1172,14 +1392,32 @@ def decode_tiled_3d_embeddings( batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> np.ndarray: - """Batch over both tile columns and z blocks while decoding tiled 3D embeddings.""" - from micro_sam.v2.util import DEFAULT_HALO_Z, DEFAULT_TILE_Z - + """Batch over both tile columns and z blocks while decoding tiled 3d embeddings. + + Args: + model: The UniSAM2 model. Its encoder is bypassed, only the decoder is run. + image_embeddings: The tiled volume embeddings (see `_compute_tiled_3d`). + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + z_block: The number of slices decoded per block. By default `micro_sam.v2.util.DEFAULT_TILE_Z`. + z_halo: The number of context slices per block. By default `micro_sam.v2.util.DEFAULT_HALO_Z`. + batch_size: The number of (tile, z block) jobs per decoder call. By default candidate sizes + are benchmarked and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read embedding blocks. + num_write_workers: The number of threads used to stitch decoded blocks into the output. + + Returns: + The stitched decoder predictions, shape (4, Z, Y, X): foreground and the three distance channels. + + Raises: + ValueError: If `z_block` is not positive or `z_halo` is negative. + """ features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=True) n_slices = shape[0] - z_block = DEFAULT_TILE_Z if z_block is None else int(z_block) - z_halo = DEFAULT_HALO_Z if z_halo is None else int(z_halo) + z_block, z_halo = _resolve_z_blocking(z_block, z_halo) jobs = _tiled_3d_jobs(features, tiling, n_slices, z_block, z_halo) output = np.zeros((4, *shape), dtype="float32") @@ -1201,25 +1439,53 @@ def write_prediction(job, prediction): _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, - num_prefetch_workers=num_prefetch_workers, + num_prefetch_workers=num_prefetch_workers, num_write_workers=num_write_workers, update_progress=pbar_update, progress_increment=lambda job: job["z1"] - job["z0"], ) return output -def decode_tiled_3d_slice( +def _decode_tiled_3d_slice( model: torch.nn.Module, image_embeddings: Dict, index: int, - device: Device = None, pbar_init: Optional[Callable] = None, pbar_update: Optional[Callable] = None, batch_size: Optional[int] = None, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, ) -> np.ndarray: - """Decode one volume slice across all embedding tiles in batches.""" + """Decode one volume slice across all embedding tiles in batches. + + Used for slice-wise automatic segmentation of a tiled volume, e.g. when only the current slice + is segmented in the annotator. + + Args: + model: The UniSAM2 model. Its encoder is bypassed, only the decoder is run. + image_embeddings: The tiled volume embeddings (see `_compute_tiled_3d`). + index: The index of the slice to decode. + pbar_init: Callback to initialize an external progress bar. + pbar_update: Callback to update an external progress bar. + batch_size: The number of tiles per decoder call. By default candidate sizes are benchmarked + and a throughput-efficient value is selected on each CUDA device. + devices: The device or devices to run inference on (see `_resolve_devices`). + num_prefetch_workers: The number of threads used to read embedding tiles. + num_write_workers: The number of threads used to stitch decoded tiles into the output. + + Returns: + The stitched decoder predictions for the slice, shape (4, Y, X). + + Raises: + ValueError: If `index` is outside the volume. + """ features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=True) + n_slices = shape[0] + index = int(index) + # A negative index would silently decode from the end, an index past the volume nothing at all. + if not 0 <= index < n_slices: + raise ValueError(f"The slice index must be in [0, {n_slices}), got {index}.") + output = np.zeros((4, *shape[1:]), dtype="float32") jobs = [] for tile_id in range(tiling.number_of_blocks): @@ -1247,7 +1513,7 @@ def write_prediction(job, prediction): _run_decoder_jobs( model, jobs, write_prediction, batch_size=batch_size, devices=devices, - num_prefetch_workers=num_prefetch_workers, + num_prefetch_workers=num_prefetch_workers, num_write_workers=num_write_workers, update_progress=pbar_update, ) return output diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index 1fae5efc1..45832ec7d 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -387,7 +387,8 @@ def initialize( if image.ndim == 3 and image.shape[-1] != 3 and i is not None: image = image[i] image_embeddings = precompute_image_embeddings( - predictor, image, save_path=save_path, ndim=2, tile_shape=tile_shape, halo=halo, verbose=verbose, + predictor, image, save_path=save_path, ndim=2, tile_shape=tile_shape, halo=halo, + verbose=verbose, lazy_loading=True, ) feats = image_embeddings["features"] @@ -900,7 +901,9 @@ def _decode_3d_feature_batch(model, features, original_size, device): output = model(dummy) finally: model.encoder = real_encoder - return output.detach().float().cpu().numpy() + # Move before casting: an fp32 copy of the whole output on the accelerator would cancel out the + # memory that fp16 inference saves (and can trigger the OOM backoff). + return output.detach().cpu().float().numpy() def _decode_3d_feature_block(model, feature, original_size, device): @@ -964,6 +967,10 @@ def __init__(self, model: torch.nn.Module, device: Optional[Union[str, torch.dev self._prediction = None self._is_initialized = False + def _inference_devices(self, devices: Devices) -> Devices: + """Fall back to the configured device, so it is not overridden by the multi-GPU default.""" + return self._device if devices is None else devices + def _run_full_inference( self, raw, ndim, tile_shape=None, halo=None, pbar_init=None, pbar_update=None, batch_size=None, devices: Devices = None, num_prefetch_workers=4, num_write_workers=1, @@ -977,7 +984,7 @@ def _run_full_inference( from torch_em.util.prediction import predict_with_halo_pipelined from micro_sam.v2.normalization import normalize_raw from micro_sam.v2.batched_inference import ( - compute_auto_batch_sizes, prepare_models, release_model_replicas, resolve_devices, + _compute_auto_batch_sizes, _prepare_models, _release_model_replicas, _resolve_devices, ) def _preprocess(crop): @@ -999,13 +1006,13 @@ def _preprocess(crop): img_size = getattr(getattr(self._model, "encoder", None), "img_size", 1024) resize_model = ResizeLongestSideWrapper(self._model, img_size) - resolved_devices = resolve_devices(resize_model, devices) + resolved_devices = _resolve_devices(resize_model, self._inference_devices(devices)) if batch_size is None: - model_devices = prepare_models(resize_model, resolved_devices) + model_devices = _prepare_models(resize_model, resolved_devices) patch_shape = tuple(block + 2 * overlap for block, overlap in zip(block_shape, block_halo)) try: - batch_sizes = compute_auto_batch_sizes( + batch_sizes = _compute_auto_batch_sizes( model_devices=model_devices, n_jobs=n_blocks, patch_shape=patch_shape, @@ -1016,7 +1023,7 @@ def _preprocess(crop): ) batch_size = min(batch_sizes) finally: - release_model_replicas(model_devices) + _release_model_replicas(model_devices) elif int(batch_size) < 1: raise ValueError(f"batch_size must be positive or None, got {batch_size}.") @@ -1064,22 +1071,22 @@ def _run_decoder_2d(self, image_embeddings): @torch.no_grad() def _run_decoder_3d( self, image_embeddings, z_block=None, z_halo=None, pbar_init=None, pbar_update=None, - batch_size=None, devices: Devices = None, num_prefetch_workers=4, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, num_write_workers=1, ): """Decode queued z blocks from precomputed 3d embeddings.""" - from micro_sam.v2.batched_inference import decode_volume_embeddings + from micro_sam.v2.batched_inference import _decode_volume_embeddings - return decode_volume_embeddings( + return _decode_volume_embeddings( model=self._model, image_embeddings=image_embeddings, - device=self._device, z_block=z_block, z_halo=z_halo, pbar_init=pbar_init, pbar_update=pbar_update, batch_size=batch_size, - devices=devices, + devices=self._inference_devices(devices), num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) @torch.no_grad() @@ -1113,11 +1120,11 @@ def initialize( pbar_update: Callback to update an external progress bar. z_block: Number of slices per decoder z block. z_halo: Overlapping decoder slices used as context. - batch_size: Explicit tile or z-block batch size. Defaults to one; pass None for - throughput-based automatic selection. + batch_size: The batch size used when running inference for multiple z-blocks and / or tiles. + Defaults to one; pass None for throughput-based automatic selection. devices: Inference device or devices. None uses all visible GPUs when the model is on CUDA. num_prefetch_workers: Number of input reading and preprocessing threads. - num_write_workers: Number of output writing threads for full tiled inference. + num_write_workers: Number of output writing threads. """ if image_embeddings is not None and ndim == 3: self._prediction = self._run_decoder_3d( @@ -1129,6 +1136,7 @@ def initialize( batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) elif image_embeddings is not None: if pbar_init is not None: @@ -1198,61 +1206,61 @@ class TiledUniSAM2InstanceSegmentation(UniSAM2InstanceSegmentation): @torch.no_grad() def _run_decoder_tiled_2d( self, image_embeddings, pbar_init=None, pbar_update=None, - batch_size=None, devices: Devices = None, num_prefetch_workers=4, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, num_write_workers=1, ): """Decode and stitch queued 2d embedding tiles.""" - from micro_sam.v2.batched_inference import decode_tiled_2d_embeddings + from micro_sam.v2.batched_inference import _decode_tiled_2d_embeddings - return decode_tiled_2d_embeddings( + return _decode_tiled_2d_embeddings( model=self._model, image_embeddings=image_embeddings, - device=self._device, pbar_init=pbar_init, pbar_update=pbar_update, batch_size=batch_size, - devices=devices, + devices=self._inference_devices(devices), num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) @torch.no_grad() def _run_decoder_tiled_3d( self, image_embeddings, pbar_init=None, pbar_update=None, z_block=None, z_halo=None, - batch_size=None, devices: Devices = None, num_prefetch_workers=4, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, num_write_workers=1, ): """Decode batches spanning tile columns and z blocks, then stitch them.""" - from micro_sam.v2.batched_inference import decode_tiled_3d_embeddings + from micro_sam.v2.batched_inference import _decode_tiled_3d_embeddings - return decode_tiled_3d_embeddings( + return _decode_tiled_3d_embeddings( model=self._model, image_embeddings=image_embeddings, - device=self._device, pbar_init=pbar_init, pbar_update=pbar_update, z_block=z_block, z_halo=z_halo, batch_size=batch_size, - devices=devices, + devices=self._inference_devices(devices), num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) @torch.no_grad() def _run_decoder_tiled_3d_slice( self, image_embeddings, i, pbar_init=None, pbar_update=None, - batch_size=None, devices: Devices = None, num_prefetch_workers=4, + batch_size=None, devices: Devices = None, num_prefetch_workers=4, num_write_workers=1, ): """Decode one volume slice across all embedding tiles in batches.""" - from micro_sam.v2.batched_inference import decode_tiled_3d_slice + from micro_sam.v2.batched_inference import _decode_tiled_3d_slice - return decode_tiled_3d_slice( + return _decode_tiled_3d_slice( model=self._model, image_embeddings=image_embeddings, index=i, - device=self._device, pbar_init=pbar_init, pbar_update=pbar_update, batch_size=batch_size, - devices=devices, + devices=self._inference_devices(devices), num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) @torch.no_grad() @@ -1277,8 +1285,8 @@ def initialize( `batch_size=None` benchmarks candidate sizes and selects a throughput-efficient batch independently on each selected CUDA device. - Reads and preprocessing are queued through `num_prefetch_workers`; full inference also overlaps - output writes through `num_write_workers`. + Reads and preprocessing are queued through `num_prefetch_workers`, output writes through + `num_write_workers`. """ decoder_kwargs = { "pbar_init": pbar_init, @@ -1286,6 +1294,7 @@ def initialize( "batch_size": batch_size, "devices": devices, "num_prefetch_workers": num_prefetch_workers, + "num_write_workers": num_write_workers, } if image_embeddings is not None and ndim == 2 and "fpn" in image_embeddings and i is not None: self._prediction = self._run_decoder_tiled_3d_slice( diff --git a/micro_sam/v2/models/_video_predictor.py b/micro_sam/v2/models/_video_predictor.py index f7f71058f..899ac74d0 100644 --- a/micro_sam/v2/models/_video_predictor.py +++ b/micro_sam/v2/models/_video_predictor.py @@ -16,6 +16,10 @@ # so repeatedly segmenting the same / nearby slice reuses the upload; small so memory stays bounded. MAX_CACHED_FRAMES = 8 +# The ImageNet statistics SAM2 normalizes its frames with (sam2.utils.misc only has them as defaults). +IMG_MEAN = (0.485, 0.456, 0.406) +IMG_STD = (0.229, 0.224, 0.225) + def _load_img_as_tensor(img_path, image_size): """Load a single frame as a float32 [0, 1] tensor of shape (3, image_size, image_size). @@ -51,6 +55,22 @@ def _load_img_as_tensor(img_path, image_size): return img, video_height, video_width +def _prepare_frame(raw, image_size): + """Resize and ImageNet-normalize one frame, exactly as the video predictor loads its frames. + + Args: + raw: The frame, either a numpy array or a path to an image file. + image_size: The size the longest side is resized to. + + Returns: + The frame as a (3, image_size, image_size) float32 tensor on the CPU. + """ + image, _, _ = _load_img_as_tensor(raw, image_size) + mean = torch.tensor(IMG_MEAN, dtype=torch.float32)[:, None, None] + std = torch.tensor(IMG_STD, dtype=torch.float32)[:, None, None] + return (image - mean) / std + + class _LazyVideoFrames: """Produce per-slice frame tensors from a numpy volume on demand, without stacking the whole volume. @@ -84,8 +104,8 @@ def _load_video_frames_from_images( volume, image_size, offload_video_to_cpu, - img_mean=(0.485, 0.456, 0.406), - img_std=(0.229, 0.224, 0.225), + img_mean=IMG_MEAN, + img_std=IMG_STD, async_loading_frames=False, compute_device=torch.device("cuda"), verbosity=True, diff --git a/micro_sam/v2/prompt_based_segmentation.py b/micro_sam/v2/prompt_based_segmentation.py index f5e2dee9d..a2461ffdf 100644 --- a/micro_sam/v2/prompt_based_segmentation.py +++ b/micro_sam/v2/prompt_based_segmentation.py @@ -120,9 +120,9 @@ def _clone_image_predictor(predictor, model: torch.nn.Module): def _get_image_predictor_devices(predictor, devices: Devices) -> List[Tuple]: """Create and cache one image-predictor replica per selected device.""" - from micro_sam.v2.batched_inference import prepare_models, resolve_devices + from micro_sam.v2.batched_inference import _prepare_models, _resolve_devices - resolved_devices = resolve_devices(predictor.model, devices) + resolved_devices = _resolve_devices(predictor.model, devices) cache_key = tuple(str(device) for device in resolved_devices) cached = getattr(predictor, "_micro_sam_predictor_devices", None) if cached is not None and cached[0] == cache_key: @@ -130,7 +130,7 @@ def _get_image_predictor_devices(predictor, devices: Devices) -> List[Tuple]: predictor_devices = [ (predictor if model is predictor.model else _clone_image_predictor(predictor, model), device) - for model, device in prepare_models(predictor.model, resolved_devices) + for model, device in _prepare_models(predictor.model, resolved_devices) ] predictor._micro_sam_predictor_devices = (cache_key, predictor_devices) return predictor_devices @@ -392,9 +392,11 @@ def tiled_promptable_segmentation_2d( tile_jobs[tile_id] = local_points, tlabels, local_boxes, local_masks predictor_devices = _get_image_predictor_devices(predictor, devices) + # Spread the active tiles round-robin. Their ids are sparse, so mapping by 'tile_id % n_devices' + # would run e.g. the tiles 0, 2, 4 on one device and leave the others idle. groups = {} - for tile_id in tile_jobs: - worker_id = tile_id % len(predictor_devices) + for position, tile_id in enumerate(sorted(tile_jobs)): + worker_id = position % len(predictor_devices) groups.setdefault(worker_id, []).append(tile_id) def run_group(worker_id, tile_ids): @@ -1013,10 +1015,10 @@ class TiledPromptableSegmentation3D: def __init__(self, predictor, volume, volume_embeddings, device=None, devices: Devices = None, **kwargs): from bioimage_cpp.utils import Blocking - from micro_sam.v2.batched_inference import prepare_models, resolve_devices + from micro_sam.v2.batched_inference import _prepare_models, _resolve_devices - resolved_devices = resolve_devices(predictor, devices) - self._predictor_devices = prepare_models(predictor, resolved_devices) + resolved_devices = _resolve_devices(predictor, devices) + self._predictor_devices = _prepare_models(predictor, resolved_devices) self.predictor = predictor self.volume = volume @@ -1032,6 +1034,8 @@ def __init__(self, predictor, volume, volume_embeddings, device=None, devices: D # Per-tile state, built lazily for the tiles that actually receive prompts. self._segmenters = {} + # Device assigned to each active tile, see '_worker_id'. + self._tile_workers = {} def init_predictor(self): # Per-tile inference states are created lazily in '_get_segmenter'. @@ -1043,6 +1047,7 @@ def reset_predictor(self): for segmenter in self._segmenters.values(): segmenter.reset_predictor() self._segmenters = {} + self._tile_workers = {} _free_device_memory() def get_progress_total(self, z_range=None): @@ -1058,6 +1063,17 @@ def _tile_index(self, y, x): return tile_id return 0 + def _worker_id(self, tile_id): + """Return the device a tile is bound to, assigned round-robin when the tile is first used. + + The tile state (inference state + embedding cache) lives on one device, so the affinity is + kept for the tile's lifetime. Assigning by 'tile_id % n_devices' instead would leave devices + idle, because the active tile ids are sparse (e.g. the tiles 0, 2, 4 share one residue). + """ + if tile_id not in self._tile_workers: + self._tile_workers[tile_id] = len(self._tile_workers) % len(self._predictor_devices) + return self._tile_workers[tile_id] + def _outer_offset(self, tile_id): """Return the (y0, x0) origin of the tile's outer (halo-included) block.""" outer = self.tiling.get_block_with_halo(tile_id, list(self.halo)).outer_block @@ -1082,7 +1098,7 @@ def _get_segmenter(self, tile_id): "input_size": tile_dataset.attrs["input_size"], "original_size": tile_dataset.attrs["original_size"], } - predictor, device = self._predictor_devices[tile_id % len(self._predictor_devices)] + predictor, device = self._predictor_devices[self._worker_id(tile_id)] self._segmenters[tile_id] = PromptableSegmentation3D( predictor, sub_volume, tile_embeddings, device=device, **self._kwargs ) @@ -1101,8 +1117,7 @@ def _run_tile_jobs( """Run tile jobs concurrently across devices and serially within each device.""" groups = {} for tile_id in tile_ids: - worker_id = tile_id % len(self._predictor_devices) - groups.setdefault(worker_id, []).append(tile_id) + groups.setdefault(self._worker_id(tile_id), []).append(tile_id) def run_group(group, worker_update=None): results = [] diff --git a/micro_sam/v2/util.py b/micro_sam/v2/util.py index 1da6ceac0..d44ed8e56 100644 --- a/micro_sam/v2/util.py +++ b/micro_sam/v2/util.py @@ -11,7 +11,7 @@ from micro_sam.util import ( get_device, get_cache_directory, microsam_cachedir, _open_embeddings, - _configure_mps_memory, + _configure_mps_memory, make_temp_embedding_path, ) from micro_sam.v2.models._video_predictor import _build_sam2_video_predictor from micro_sam.v2.normalization import RAW_NORMALIZATION, to_image @@ -425,6 +425,7 @@ def precompute_image_embeddings( batch_size: Optional[int] = 1, devices: Devices = None, num_prefetch_workers: int = 4, + num_write_workers: int = 1, pbar_init: Optional[callable] = None, pbar_update: Optional[callable] = None, ): @@ -436,17 +437,21 @@ def precompute_image_embeddings( predictor: The SAM2 image or video predictor. input_: The input image or volume. save_path: Optional zarr path for persistent embedding storage. - lazy_loading: Return disk-backed arrays for saved 3D embeddings. + lazy_loading: Stream the embeddings from the zarr instead of materializing them in memory. + Applies to volumes and to tiled images; if no `save_path` is given an ephemeral on-disk + cache is used, which is removed at process exit. ndim: The number of spatial dimensions. By default this is inferred from the input. tile_shape: Optional in-plane tile shape. halo: Optional in-plane tile halo. verbose: Whether to show progress. - batch_size: Number of independent tiles or slices per encoder call. Defaults to one, which - is recommended for 10 GB MIG devices. Pass None to benchmark candidate sizes and select - a throughput-efficient value independently on each CUDA device. + batch_size: The batch size used when running inference for multiple slices and / or tiles. + Defaults to one, which is recommended for 10 GB MIG devices. Pass None to benchmark + candidate sizes and select a throughput-efficient value independently on each CUDA device. devices: Device or devices used for embedding inference. If None and the predictor is on CUDA, all visible CUDA devices are used. num_prefetch_workers: Number of threads used to read and preprocess input jobs. + num_write_workers: Number of threads used to write the embeddings. Only has an effect for + volumes and tiled images, which are written incrementally. pbar_init: Optional callback to initialize external progress. pbar_update: Optional callback to update external progress. @@ -457,6 +462,12 @@ def precompute_image_embeddings( if ndim == 2: configure_image_predictor(predictor) + is_streamed = ndim == 3 or tile_shape is not None + if lazy_loading and is_streamed and save_path is None: + # Without a save path the zarr is held in memory, which defeats lazy loading. Back it by an + # ephemeral on-disk store so the tiles / slices are streamed from disk instead. + save_path = make_temp_embedding_path() + # Handle the embedding save_path. # We don't have a save path, open in memory zarr file to hold tiled embeddings. if save_path is None: @@ -481,7 +492,7 @@ def precompute_image_embeddings( f.attrs["normalization"] = RAW_NORMALIZATION from micro_sam.util import handle_pbar - from micro_sam.v2.batched_inference import compute_3d, compute_tiled_2d, compute_tiled_3d + from micro_sam.v2.batched_inference import _compute_3d, _compute_tiled_2d, _compute_tiled_3d _, pbar_init, pbar_update, pbar_close = handle_pbar(verbose, pbar_init, pbar_update) @@ -490,21 +501,24 @@ def precompute_image_embeddings( elif ndim == 2 and tile_shape is not None: if halo is None: raise ValueError("To compute tiled embeddings the parameter halo has to be passed.") - embeddings = compute_tiled_2d( + embeddings = _compute_tiled_2d( input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) elif ndim == 3 and tile_shape is None: - embeddings = compute_3d( + embeddings = _compute_3d( input_, predictor, f, save_path, lazy_loading, pbar_init, pbar_update, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) elif ndim == 3 and tile_shape is not None: if halo is None: raise ValueError("To compute tiled embeddings the parameter halo has to be passed.") - embeddings = compute_tiled_3d( + embeddings = _compute_tiled_3d( input_, predictor, tile_shape, halo, f, save_path, pbar_init, pbar_update, batch_size=batch_size, devices=devices, num_prefetch_workers=num_prefetch_workers, + num_write_workers=num_write_workers, ) else: raise ValueError(f"Invalid dimensionality {input_.ndim}, expect 2 or 3 dim data.") diff --git a/test/test_sam_annotator/test_widgets.py b/test/test_sam_annotator/test_widgets.py index 34174c86d..3ae60c133 100644 --- a/test/test_sam_annotator/test_widgets.py +++ b/test/test_sam_annotator/test_widgets.py @@ -68,3 +68,30 @@ def test_embedding_widget(make_napari_viewer, tmp_path): # Close the viewer at the end of the test. viewer.close() + + +@pytest.mark.gui +def test_batch_size_visibility_follows_device_and_model(make_napari_viewer_proxy): + """The batch size control is only shown where it has an effect (GPU, and not a VFM encoder).""" + from micro_sam.sam_annotator._widgets import ClassificationEmbeddingWidget + + make_napari_viewer_proxy() + widget = ClassificationEmbeddingWidget() + + widget.device = "cpu" + widget.model_type = "hvit_t_cells" + widget._update_batch_size_visibility() + assert widget._batch_size_widget.isHidden() + + widget.device = "cuda" + widget._update_batch_size_visibility() + assert not widget._batch_size_widget.isHidden() + + # The VFM encoders offered by the classifiers compute their embeddings unbatched. + widget.model_type = "vit_b_dinov2" + widget._update_batch_size_visibility() + assert widget._batch_size_widget.isHidden() + + widget.model_type = "vit_b_lm" + widget._update_batch_size_visibility() + assert not widget._batch_size_widget.isHidden() diff --git a/test/test_v2_automatic_segmentation.py b/test/test_v2_automatic_segmentation.py index 1c9072b00..b3ca4a647 100644 --- a/test/test_v2_automatic_segmentation.py +++ b/test/test_v2_automatic_segmentation.py @@ -7,7 +7,7 @@ import pytest from micro_sam.v2.instance_segmentation import ( - _block_shape_and_halo, _set_image_predictor_from_backbone, + _block_shape_and_halo, _set_image_predictor_from_backbone, _decode_3d_feature_batch, _get_decoder_autocast, UniSAM2InstanceSegmentation, TiledUniSAM2InstanceSegmentation, get_instance_segmentation_generator, get_decoder, ) @@ -351,3 +351,66 @@ def test_tiled_decoder_2d_stitches_all_tiles(): out = segmenter._run_decoder_tiled_2d({"features": feats_group}) assert out.shape == (4, *shape) assert len(model.call_z) == 4 # one decoder pass per tile + + +@pytest.mark.parametrize("devices,expected", [(None, "cuda:1"), ("cpu", "cpu")]) +def test_configured_device_is_used_when_no_devices_are_given(devices, expected, monkeypatch): + # Regression: the batched decoder paths only forwarded 'devices', so the device the segmenter was + # constructed with was silently replaced by the model's device (or by all visible GPUs). + import micro_sam.v2.batched_inference as batched_inference + + forwarded = {} + + def fake_decode(model, image_embeddings, **kwargs): + forwarded["devices"] = kwargs["devices"] + return np.zeros((4, 1, 8, 8), dtype="float32") + + monkeypatch.setattr(batched_inference, "_decode_volume_embeddings", fake_decode) + segmenter = UniSAM2InstanceSegmentation(_FakeUNETR(img_size=8), device="cuda:1") + segmenter._run_decoder_3d({"features": np.zeros((1, 2, 4, 4), dtype="float32")}, devices=devices) + assert forwarded["devices"] == expected + + +class _RecordingOutput: + """Stands in for the decoder output tensor and records the conversions applied to it.""" + + def __init__(self, array, calls): + self._array = array + self._calls = calls + + def detach(self): + self._calls.append("detach") + return self + + def cpu(self): + self._calls.append("cpu") + return self + + def float(self): + self._calls.append("float") + return self + + def numpy(self): + return self._array + + +class _RecordingModel: + def __init__(self, output): + self.encoder = types.SimpleNamespace(img_size=8) + self._output = output + + def __call__(self, x): + return self._output + + +def test_decoder_output_is_moved_to_cpu_before_the_float_cast(): + # Casting on the device would hold an fp32 copy of the whole output next to the fp16 one, + # cancelling out the memory that fp16 inference saves. + calls = [] + array = np.zeros((1, 4, 1, 8, 8), dtype="float32") + model = _RecordingModel(_RecordingOutput(array, calls)) + + out = _decode_3d_feature_batch(model, torch.zeros((1, 1, 2, 4, 4)), (8, 8), "cpu") + + assert out is array + assert calls == ["detach", "cpu", "float"] diff --git a/test/test_v2_batched_inference.py b/test/test_v2_batched_inference.py index 3e2bb1dd1..c908b67e9 100644 --- a/test/test_v2_batched_inference.py +++ b/test/test_v2_batched_inference.py @@ -1,4 +1,5 @@ import threading +import time import unittest from unittest import mock @@ -8,12 +9,13 @@ from bioimage_cpp.utils import Blocking from micro_sam.v2.batched_inference import ( + _compute_auto_batch_sizes, + _decode_tiled_3d_embeddings, + _decode_tiled_3d_slice, + _decode_volume_embeddings, + _run_batched_pipeline, _run_decoder_jobs, _select_throughput_batch_size, - compute_auto_batch_sizes, - decode_tiled_3d_embeddings, - decode_volume_embeddings, - run_batched_pipeline, ) from micro_sam.v2.prompt_based_segmentation import TiledPromptableSegmentation3D @@ -80,7 +82,7 @@ def update_progress(update): progress.append(update) progress_threads.append(threading.get_ident()) - run_batched_pipeline( + _run_batched_pipeline( jobs=range(7), model_devices=[(3, torch.device("cpu"))], batch_sizes=[3], @@ -95,6 +97,97 @@ def update_progress(update): self.assertEqual(sum(progress), 7) self.assertEqual(set(progress_threads), {threading.get_ident()}) + def test_pipeline_writes_with_multiple_workers(self): + outputs = {} + writer_threads = set() + lock = threading.Lock() + + def predict(model, items, device): + return [value + model for value in items] + + def write(job, prediction): + # Slow writes so the queue actually spreads over the writers. + time.sleep(0.01) + with lock: + outputs[job] = prediction + writer_threads.add(threading.get_ident()) + + _run_batched_pipeline( + jobs=range(16), + model_devices=[(3, torch.device("cpu"))], + batch_sizes=[2], + load_fn=lambda value: 2 * value, + predict_fn=predict, + write_fn=write, + num_prefetch_workers=2, + num_write_workers=3, + ) + + self.assertEqual(outputs, {index: 2 * index + 3 for index in range(16)}) + self.assertGreater(len(writer_threads), 1) + self.assertLessEqual(len(writer_threads), 3) + self.assertNotIn(threading.get_ident(), writer_threads) + + def test_pipeline_does_not_deadlock_when_prediction_fails(self): + # Regression: the sentinels were sent with a blocking put. A consumer that failed while the + # input queue was full left nobody to drain it, so the producer (and the join) hung forever. + def predict(model, items, device): + time.sleep(0.3) # let the producers reach their sentinel put with a full queue + raise RuntimeError("synthetic prediction failure") + + errors = [] + + def run(): + try: + _run_batched_pipeline( + jobs=range(3), + model_devices=[(3, torch.device("cpu"))], + batch_sizes=[1], + load_fn=lambda value: value, + predict_fn=predict, + write_fn=lambda job, prediction: None, + num_prefetch_workers=3, + ) + except Exception as exc: # noqa + errors.append(exc) + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=30) + + self.assertFalse(thread.is_alive(), "the pipeline deadlocked after a worker failure") + self.assertEqual([str(error) for error in errors], ["synthetic prediction failure"]) + + def test_pipeline_does_not_deadlock_when_a_write_fails(self): + # The symmetric case: the consumers' sentinels go to the writers, which may already have died. + def write(job, prediction): + time.sleep(0.02) + raise RuntimeError("synthetic write failure") + + errors = [] + + def run(): + try: + _run_batched_pipeline( + jobs=range(12), + model_devices=[(1, torch.device("cpu"))], + batch_sizes=[1], + load_fn=lambda value: value, + predict_fn=lambda model, items, device: list(items), + write_fn=write, + num_prefetch_workers=3, + num_write_workers=4, + ) + except Exception as exc: # noqa + errors.append(exc) + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=30) + + self.assertFalse(thread.is_alive(), "the pipeline deadlocked after a write failure") + self.assertEqual([str(error) for error in errors], ["synthetic write failure"]) + def test_pipeline_retries_ooming_batches(self): outputs = {} attempted_batch_sizes = [] @@ -106,7 +199,7 @@ def predict(model, items, device): return [value + model for value in items] with self.assertWarnsRegex(RuntimeWarning, "retrying with smaller batches"): - run_batched_pipeline( + _run_batched_pipeline( jobs=range(7), model_devices=[(3, torch.device("cpu"))], batch_sizes=[4], @@ -131,7 +224,7 @@ def predict(model, items, device): raise torch.cuda.OutOfMemoryError("synthetic fragmented-cache OOM") return [value + model for value in items] - run_batched_pipeline( + _run_batched_pipeline( jobs=[0], model_devices=[(3, torch.device("cpu"))], batch_sizes=[1], load_fn=lambda value: 2 * value, predict_fn=predict, write_fn=outputs.__setitem__, ) @@ -158,7 +251,7 @@ def test_requires_material_throughput_improvement(self): def test_stops_after_two_non_improving_candidates(self, measure): model = torch.nn.Linear(1, 1) - batch_sizes = compute_auto_batch_sizes( + batch_sizes = _compute_auto_batch_sizes( model_devices=[(model, torch.device("cuda:0"))], n_jobs=64, patch_shape=(8, 8), @@ -179,10 +272,9 @@ def test_volume_decoder_batches_z_blocks(self): model = FakeUNETR() progress = [] - output = decode_volume_embeddings( + output = _decode_volume_embeddings( model, {"features": features, "original_size": (8, 8)}, - device="cpu", z_block=2, z_halo=0, batch_size=2, @@ -210,10 +302,9 @@ def test_tiled_volume_decoder_batches_tiles_and_z_blocks(self): model = FakeUNETR() progress = [] - output = decode_tiled_3d_embeddings( + output = _decode_tiled_3d_embeddings( model, {"features": features}, - device="cpu", z_block=2, z_halo=0, batch_size=2, @@ -244,6 +335,40 @@ def test_largest_decoder_shape_runs_first(self): self.assertEqual(write_order, ["large", "small"]) +class TestDecoderEncoderRestoration(unittest.TestCase): + def test_encoder_is_restored_when_job_inspection_fails(self): + # Regression: the encoder was swapped for the decoder-only replica before the job shapes were + # inspected, so a malformed job left the caller's model with the placeholder encoder. + model = FakeUNETR() + encoder = model.encoder + malformed = [{"source": np.ones((2,) * 6, dtype="float32"), "original_size": (8, 8)}] + + with self.assertRaises(ValueError): + _run_decoder_jobs(model, malformed, lambda job, prediction: None, batch_size=1, devices="cpu") + self.assertIs(model.encoder, encoder) + + def test_encoder_is_restored_when_batch_size_is_invalid(self): + model = FakeUNETR() + encoder = model.encoder + jobs = [{"source": np.ones((2, 2, 2, 2), dtype="float32"), "original_size": (8, 8)}] + + with self.assertRaisesRegex(ValueError, "batch_size must be positive"): + _run_decoder_jobs(model, jobs, lambda job, prediction: None, batch_size=0, devices="cpu") + self.assertIs(model.encoder, encoder) + + def test_encoder_is_restored_when_a_write_fails(self): + model = FakeUNETR() + encoder = model.encoder + jobs = [{"source": np.ones((2, 2, 2, 2), dtype="float32"), "original_size": (8, 8)}] + + def write(job, prediction): + raise RuntimeError("synthetic write failure") + + with self.assertRaisesRegex(RuntimeError, "synthetic write failure"): + _run_decoder_jobs(model, jobs, write, batch_size=1, devices="cpu") + self.assertIs(model.encoder, encoder) + + class FakeInteractiveTile: def __init__(self, value): self.value = value @@ -261,6 +386,7 @@ def test_active_tile_columns_are_stitched_after_device_jobs(self): segmenter.halo = (0, 0) segmenter.tiling = Blocking([0, 0], [8, 8], [4, 4]) segmenter._predictor_devices = [(None, torch.device("cpu")), (None, torch.device("cpu"))] + segmenter._tile_workers = {} segmenter._segmenters = { tile_id: FakeInteractiveTile(tile_id + 1) for tile_id in range(4) @@ -282,5 +408,73 @@ def update_progress(value): self.assertEqual(set(progress_threads), {threading.get_ident()}) +class TestZBlockValidation(unittest.TestCase): + def _tiled_embeddings(self): + features = Group( + { + "0": ArrayWithAttrs(np.zeros((4, 1, 2, 2, 2), dtype="float32"), original_size=(4, 4)), + }, + shape=(4, 4, 4), + tile_shape=(4, 4), + halo=(0, 0), + ) + return {"features": features} + + def test_tiled_decoder_rejects_invalid_z_blocking(self): + # Regression: a negative z_block made the job range empty, so an all-zero prediction was + # returned without ever running the decoder. + model = FakeUNETR() + for z_block, z_halo in [(-1, 0), (0, 0), (2, -1)]: + with self.assertRaisesRegex(ValueError, "z_block must be positive"): + _decode_tiled_3d_embeddings( + model, self._tiled_embeddings(), z_block=z_block, z_halo=z_halo, batch_size=1, + ) + + def test_volume_decoder_rejects_invalid_z_blocking(self): + model = FakeUNETR() + embeddings = {"features": np.zeros((4, 2, 2, 2), dtype="float32"), "original_size": (8, 8)} + for z_block, z_halo in [(-1, 0), (0, 0), (2, -1)]: + with self.assertRaisesRegex(ValueError, "z_block must be positive"): + _decode_volume_embeddings(model, embeddings, z_block=z_block, z_halo=z_halo, batch_size=1) + + def test_tiled_slice_decoder_rejects_out_of_range_index(self): + # A negative index would decode from the end, an index past the volume nothing at all. + model = FakeUNETR() + for index in (-1, 4, 10): + with self.assertRaisesRegex(ValueError, "slice index must be in"): + _decode_tiled_3d_slice(model, self._tiled_embeddings(), index=index, batch_size=1) + + def test_tiled_slice_decoder_accepts_valid_index(self): + model = FakeUNETR() + output = _decode_tiled_3d_slice(model, self._tiled_embeddings(), index=3, batch_size=1) + self.assertEqual(output.shape, (4, 4, 4)) + + +class TestTileDeviceAffinity(unittest.TestCase): + def _segmenter(self, n_devices): + segmenter = TiledPromptableSegmentation3D.__new__(TiledPromptableSegmentation3D) + segmenter._predictor_devices = [(None, torch.device("cpu"))] * n_devices + segmenter._tile_workers = {} + return segmenter + + def test_sparse_tiles_are_spread_over_devices(self): + # Regression: 'tile_id % n_devices' mapped the tiles 0, 2, 4, 6 to the same device. + segmenter = self._segmenter(2) + segmenter._run_tile_jobs([0, 2, 4, 6], lambda tile_id: tile_id) + self.assertEqual([segmenter._worker_id(tile_id) for tile_id in (0, 2, 4, 6)], [0, 1, 0, 1]) + + def test_tile_affinity_is_kept_across_jobs(self): + # The tile state lives on one device, so its assignment must not change between jobs. + segmenter = self._segmenter(3) + segmenter._run_tile_jobs([3, 6], lambda tile_id: tile_id) + assignment = dict(segmenter._tile_workers) + self.assertEqual(sorted(assignment.values()), [0, 1]) + + segmenter._run_tile_jobs([6, 3, 9], lambda tile_id: tile_id) + for tile_id, worker_id in assignment.items(): + self.assertEqual(segmenter._worker_id(tile_id), worker_id) + self.assertEqual(segmenter._worker_id(9), 2) + + if __name__ == "__main__": unittest.main() From 90f3602ec341db81554089acca056c126796e315 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Thu, 23 Jul 2026 01:24:00 +0200 Subject: [PATCH 11/11] Try to fix Windows CI --- test/test_sam_annotator/test_widgets.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_sam_annotator/test_widgets.py b/test/test_sam_annotator/test_widgets.py index 3ae60c133..2f47a81fa 100644 --- a/test/test_sam_annotator/test_widgets.py +++ b/test/test_sam_annotator/test_widgets.py @@ -71,12 +71,12 @@ def test_embedding_widget(make_napari_viewer, tmp_path): @pytest.mark.gui -def test_batch_size_visibility_follows_device_and_model(make_napari_viewer_proxy): +def test_batch_size_visibility_follows_device_and_model(qtbot): """The batch size control is only shown where it has an effect (GPU, and not a VFM encoder).""" from micro_sam.sam_annotator._widgets import ClassificationEmbeddingWidget - make_napari_viewer_proxy() widget = ClassificationEmbeddingWidget() + qtbot.addWidget(widget) widget.device = "cpu" widget.model_type = "hvit_t_cells"