diff --git a/environment.yaml b/environment.yaml index 9e5553dc..b5d52e37 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 9c21a049..d3947942 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 00000000..d698dd9f --- /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 a0537ec8..cb44fccd 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 5564d9cf..6be14f2c 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 b3e40226..a767849a 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 60086002..b853468c 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 00000000..18430054 --- /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()