Merge pull request #1171 from computational-cell-analytics/dev

Bump version

Author: constantinpape

doc/faq.md

@@ -187,6 +187,11 @@ You can then use those models with the custom checkpoint option, see answer 15 f
 ### 18. I would like to evaluate the instance segmentation quantitatively. Can you suggest how to do that?
 `micro-sam` supports a `micro_sam.evaluate` CLI, which computes the mean segmentation accuracy (introduced in the Pascal VOC challenge) of the predicted instance segmentation with the corresponding ground-truth annotations. Please see our paper (`Methods` -> `Inference and Evaluation` for more details about it) and `$ micro_sam.evaluate -h` for more details about the evaluation CLI.
 
+### 19. I get `RuntimeError: GET was unable to find an engine to execute this computation` on a V100 GPU (*"or any older GPU"*).
+This is a known issue for a combination of older generation GPUs (eg. V100s) and pytorch compiled with the latest CUDA Toolkit (eg. CUDA 12.9 and PyTorch 2.8 has been tested to throw this error on V100s).
+Here's what you can do to solve this issue:
+- Use a PyTorch/CUDA build that is known to work with V100, for example CUDA 12.1 or 11.8 with a compatible PyTorch version (please check your installed CUDA drivers).
+- Run on CPU (slower, but works).
 
 ## Fine-tuning questions
 

micro_sam/__version__.py

@@ -1 +1 @@
-__version__ = "1.7.1"
+__version__ = "1.7.2"

micro_sam/automatic_segmentation.py

@@ -127,10 +127,9 @@ def automatic_tracking(
         The lineages representing cell divisions, stored as a dictionary.
     """
     # Load the input image file.
-    if isinstance(input_path, np.ndarray):
-        image_data = input_path
-    else:
-        image_data = util.load_image_data(input_path, key)
+    # We assume that it has to be read from file if it is a str or pathlike.
+    # Otherwise we assume it is a numpy array like object.
+    image_data = util.load_image_data(input_path, key) if isinstance(input_path, (str, os.PathLike)) else input_path
 
     if (image_data.ndim != 3) and (image_data.ndim != 4 and image_data.shape[-1] != 3):
         raise ValueError(f"The inputs does not match the shape expectation of 3d inputs: {image_data.shape}")
@@ -168,7 +167,9 @@ def automatic_instance_segmentation(
     input_path: Union[Union[os.PathLike, str], np.ndarray],
     output_path: Optional[Union[os.PathLike, str]] = None,
     embedding_path: Optional[Union[os.PathLike, str]] = None,
+    mask_path: Optional[Union[Union[os.PathLike, str], np.ndarray]] = None,
     key: Optional[str] = None,
+    mask_key: Optional[str] = None,
     ndim: Optional[int] = None,
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
@@ -187,8 +188,10 @@ def automatic_instance_segmentation(
             or a container file (e.g. hdf5 or zarr).
         output_path: The output path where the instance segmentations will be saved.
         embedding_path: The path where the embeddings are cached already / will be saved.
+        mask_path: The path to an optional foreground mask. Areas outside of the foreground will not be processed.
         key: The key to the input file. This is needed for container files (eg. hdf5 or zarr)
             or to load several images as 3d volume. Provide a glob patterm, eg. "*.tif", for this case.
+        mask_key: The key to the (optional) foreground mask.
         ndim: The dimensionality of the data. By default the dimensionality of the data will be used.
             If you have RGB data you have to specify this explicitly, e.g. pass ndim=2 for 2d segmentation of RGB.
         tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
@@ -212,11 +215,9 @@ def automatic_instance_segmentation(
             print(f"The segmentation results are already stored at '{os.path.abspath(output_path)}'.")
             return
 
-    # Load the input image file.
-    if isinstance(input_path, np.ndarray):
-        image_data = input_path
-    else:
-        image_data = util.load_image_data(input_path, key)
+    # We assume that it has to be read from file if it is a str or pathlike.
+    # Otherwise we assume it is a numpy array like object.
+    image_data = util.load_image_data(input_path, key) if isinstance(input_path, (str, os.PathLike)) else input_path
 
     ndim = image_data.ndim if ndim is None else ndim
 
@@ -244,6 +245,11 @@ def automatic_instance_segmentation(
             generate_kwargs.update({"tile_shape": tile_shape, "halo": halo})
             initialize_kwargs["batch_size"] = batch_size
 
+        # Load the mask defining foreground if it was given.
+        if mask_path is not None:
+            mask = util.load_image_data(mask_path, mask_key) if isinstance(mask_path, (str, os.PathLike)) else mask_path
+            initialize_kwargs["mask"] = mask
+
         segmenter.initialize(**initialize_kwargs)
         instances = segmenter.generate(**generate_kwargs)
 

micro_sam/instance_segmentation.py

@@ -4,17 +4,25 @@
 """
 
 import os
+import shutil
+import tempfile
 import warnings
 from abc import ABC
+from contextlib import contextmanager
 from copy import deepcopy
 from collections import OrderedDict
 from typing import Any, Dict, Literal, List, Optional, Tuple, Union
 
-import vigra
 import numpy as np
+import zarr
 from skimage.measure import regionprops
 from skimage.segmentation import find_boundaries
 
+try:
+    import fastfilters as filter_impl
+except ImportError:
+    import vigra.filters as filter_impl
+
 import torch
 from torchvision.ops.boxes import batched_nms, box_area
 
@@ -23,6 +31,8 @@
 
 import elf.parallel as parallel_impl
 from elf.parallel.filters import apply_filter
+from elf.wrapper.base import MultiTransformationWrapper
+from elf.wrapper.generic import ThresholdWrapper
 
 from nifty.tools import blocking
 
@@ -853,6 +863,23 @@ def get_predictor_and_decoder(
     return predictor, decoder
 
 
+@contextmanager
+def _array_or_zarr(shape, dtype, chunks, use_zarr=False):
+    if not use_zarr:
+        yield np.zeros(shape, dtype=dtype)
+        return
+
+    tmpdir = tempfile.mkdtemp(prefix="tmp-zarr-")
+    try:
+        store_path = os.path.join(tmpdir, "tmp.zarr")
+        root = zarr.open_group(store_path, mode="w")
+        arr = root.create_dataset(name="data", shape=shape, dtype=dtype, chunks=chunks)
+        yield arr
+
+    finally:
+        shutil.rmtree(tmpdir, ignore_errors=True)
+
+
 def _watershed_from_center_and_boundary_distances_parallel(
     center_distances,
     boundary_distances,
@@ -866,6 +893,8 @@ def _watershed_from_center_and_boundary_distances_parallel(
     halo,
     n_threads,
     verbose=False,
+    optimize_memory=False,
+    segmentation=None,
 ):
     center_distances = apply_filter(
         center_distances, "gaussianSmoothing", sigma=distance_smoothing,
@@ -876,30 +905,45 @@ def _watershed_from_center_and_boundary_distances_parallel(
         block_shape=tile_shape, n_threads=n_threads
     )
 
-    fg_mask = foreground_map > foreground_threshold
+    fg_mask = ThresholdWrapper(foreground_map, foreground_threshold, operator=np.greater)
 
-    marker_map = np.logical_and(
-        center_distances < center_distance_threshold, boundary_distances < boundary_distance_threshold
+    marker_map = MultiTransformationWrapper(
+        np.logical_and,
+        ThresholdWrapper(center_distances, center_distance_threshold, operator=np.less),
+        ThresholdWrapper(boundary_distances, boundary_distance_threshold, operator=np.less),
     )
-    marker_map[~fg_mask] = 0
+    marker_map = MultiTransformationWrapper(np.logical_and, marker_map, fg_mask)
 
-    markers = np.zeros(marker_map.shape, dtype="uint64")
-    markers = parallel_impl.label(
-        marker_map, out=markers, block_shape=tile_shape, n_threads=n_threads, verbose=verbose,
-    )
+    with _array_or_zarr(marker_map.shape, dtype="uint64", chunks=tile_shape, use_zarr=optimize_memory) as markers:
+        markers = parallel_impl.label(
+            marker_map, out=markers, block_shape=tile_shape, n_threads=n_threads, verbose=verbose,
+        )
 
-    seg = np.zeros_like(markers, dtype="uint64")
-    seg = parallel_impl.seeded_watershed(
-        boundary_distances, seeds=markers, out=seg, block_shape=tile_shape,
-        halo=halo, n_threads=n_threads, verbose=verbose, mask=fg_mask,
-    )
+        if segmentation is None:
+            segmentation = np.zeros(markers.shape, dtype="uint64")
+        segmentation = parallel_impl.seeded_watershed(
+            boundary_distances, seeds=markers, out=segmentation, block_shape=tile_shape,
+            halo=halo, n_threads=n_threads, verbose=verbose, mask=fg_mask,
+        )
 
-    out = np.zeros_like(seg, dtype="uint64")
-    out = parallel_impl.size_filter(
-        seg, out=out, min_size=min_size, block_shape=tile_shape, n_threads=n_threads, verbose=verbose
-    )
+    if min_size > 0:
+        segmentation = parallel_impl.size_filter(
+            segmentation, out=segmentation, min_size=min_size,
+            block_shape=tile_shape, n_threads=n_threads, verbose=verbose
+        )
 
-    return out
+    return segmentation
+
+
+def _apply_smoothing(foreground, foreground_smoothing, tile_shape, n_threads):
+    if tile_shape is None:
+        foreground = filter_impl.gaussianSmoothing(foreground, foreground_smoothing)
+    else:
+        foreground = apply_filter(
+            foreground, "gaussianSmoothing", sigma=foreground_smoothing,
+            block_shape=tile_shape, n_threads=n_threads
+        )
+    return foreground
 
 
 class InstanceSegmentationWithDecoder:
@@ -1042,6 +1086,8 @@ def generate(
         tile_shape: Optional[Tuple[int, int]] = None,
         halo: Optional[Tuple[int, int]] = None,
         n_threads: Optional[int] = None,
+        optimize_memory: bool = False,
+        segmentation: Optional[np.ndarray] = None,
     ) -> Union[List[Dict[str, Any]], np.ndarray]:
         """Generate instance segmentation for the currently initialized image.
 
@@ -1067,6 +1113,8 @@ def generate(
                 If not given then post-processing will not be parallelized.
             halo: Halo for parallel post-processing. See also `tile_shape`.
             n_threads: Number of threads for parallel post-processing. See also `tile_shape`.
+            optimize_memory: Whether to optimize the memory consumption by allocating intermediate files.
+            segmentation: Optional pre-allocated segmentation.
 
         Returns:
             The segmentation masks.
@@ -1075,7 +1123,7 @@ def generate(
             raise RuntimeError("InstanceSegmentationWithDecoder has not been initialized. Call initialize first.")
 
         if foreground_smoothing > 0:
-            foreground = vigra.filters.gaussianSmoothing(self._foreground, foreground_smoothing)
+            foreground = _apply_smoothing(self._foreground, foreground_smoothing, tile_shape, n_threads)
         else:
             foreground = self._foreground
 
@@ -1106,6 +1154,8 @@ def generate(
                 halo=halo,
                 n_threads=n_threads,
                 verbose=False,
+                optimize_memory=optimize_memory,
+                segmentation=segmentation,
             )
 
         if output_mode != "instance_segmentation":

micro_sam/sam_annotator/_annotator.py

@@ -34,14 +34,14 @@ def _require_layers(self, layer_choices: Optional[List[str]] = None):
                 widgets._validation_window_for_missing_layer("current_object")
             self._viewer.add_labels(data=dummy_data, name="current_object")
             if image_scale is not None:
-                self.layers["current_objects"].scale = image_scale
+                self._viewer.layers["current_object"].scale = image_scale
 
         if "auto_segmentation" not in self._viewer.layers:
             if layer_choices and "auto_segmentation" in layer_choices:  # Check at 'commit' call button.
                 widgets._validation_window_for_missing_layer("auto_segmentation")
             self._viewer.add_labels(data=dummy_data, name="auto_segmentation")
             if image_scale is not None:
-                self.layers["auto_segmentation"].scale = image_scale
+                self._viewer.layers["auto_segmentation"].scale = image_scale
 
         if "committed_objects" not in self._viewer.layers:
             if layer_choices and "committed_objects" in layer_choices:  # Check at 'commit' call button.
@@ -50,7 +50,7 @@ def _require_layers(self, layer_choices: Optional[List[str]] = None):
             # Randomize colors so it is easy to see when object committed.
             self._viewer.layers["committed_objects"].new_colormap()
             if image_scale is not None:
-                self.layers["committed_objects"].scale = image_scale
+                self._viewer.layers["committed_objects"].scale = image_scale
 
         # Add the point layer for point prompts.
         self._point_labels = ["positive", "negative"]

micro_sam/sam_annotator/_state.py

@@ -94,6 +94,7 @@ def initialize_predictor(
         predictor=None,
         decoder=None,
         checkpoint_path=None,
+        decoder_path=None,
         tile_shape=None,
         halo=None,
         precompute_amg_state=False,
@@ -113,7 +114,7 @@ def progress_bar_factory(model_type):
 
             self.predictor, state = util.get_sam_model(
                 device=device, model_type=model_type,
-                checkpoint_path=checkpoint_path, return_state=True,
+                checkpoint_path=checkpoint_path, decoder_path=decoder_path, return_state=True,
                 progress_bar_factory=None if use_cli else progress_bar_factory,
             )
             if prefer_decoder and "decoder_state" in state and model_type != "vit_b_medical_imaging":

micro_sam/sam_annotator/annotator_2d.py

@@ -48,6 +48,7 @@ def annotator_2d(
     viewer: Optional["napari.viewer.Viewer"] = None,
     precompute_amg_state: bool = False,
     checkpoint_path: Optional[str] = None,
+    decoder_path: Optional[str] = None,
     device: Optional[Union[str, torch.device]] = None,
     prefer_decoder: bool = True,
 ) -> Optional["napari.viewer.Viewer"]:
@@ -73,6 +74,7 @@ def annotator_2d(
             This will take more time when precomputing embeddings, but will then make
             automatic mask generation much faster. By default, set to 'False'.
         checkpoint_path: Path to a custom checkpoint from which to load the SAM model.
+        decoder_path: Path to a custom decoder checkpoint from which to load the 'micro-sam` decoder.
         device: The computational device to use for the SAM model.
             By default, automatically chooses the best available device.
         prefer_decoder: Whether to use decoder based instance segmentation if
@@ -89,7 +91,8 @@ def annotator_2d(
     state.initialize_predictor(
         image, model_type=model_type, save_path=embedding_path,
         halo=halo, tile_shape=tile_shape, precompute_amg_state=precompute_amg_state,
-        ndim=2, checkpoint_path=checkpoint_path, device=device, prefer_decoder=prefer_decoder,
+        ndim=2, checkpoint_path=checkpoint_path, decoder_path=decoder_path,
+        device=device, prefer_decoder=prefer_decoder,
         skip_load=False, use_cli=True,
     )
 
@@ -137,5 +140,5 @@ def main():
         segmentation_result=segmentation_result,
         model_type=args.model_type, tile_shape=args.tile_shape, halo=args.halo,
         precompute_amg_state=args.precompute_amg_state, checkpoint_path=args.checkpoint,
-        device=args.device, prefer_decoder=args.prefer_decoder,
+        decoder_path=args.decoder_path, device=args.device, prefer_decoder=args.prefer_decoder,
     )

micro_sam/sam_annotator/annotator_3d.py

@@ -58,6 +58,7 @@ def annotator_3d(
     viewer: Optional["napari.viewer.Viewer"] = None,
     precompute_amg_state: bool = False,
     checkpoint_path: Optional[str] = None,
+    decoder_path: Optional[str] = None,
     device: Optional[Union[str, torch.device]] = None,
     prefer_decoder: bool = True,
 ) -> Optional["napari.viewer.Viewer"]:
@@ -83,6 +84,7 @@ def annotator_3d(
             This will take more time when precomputing embeddings, but will then make
             automatic mask generation much faster. By default, set to 'False'.
         checkpoint_path: Path to a custom checkpoint from which to load the SAM model.
+        decoder_path: Path to a custom decoder checkpoint from which to load the 'micro-sam` decoder.
         device: The computational device to use for the SAM model.
             By default, automatically chooses the best available device.
         prefer_decoder: Whether to use decoder based instance segmentation if
@@ -99,7 +101,8 @@ def annotator_3d(
     state.initialize_predictor(
         image, model_type=model_type, save_path=embedding_path,
         halo=halo, tile_shape=tile_shape, ndim=3, precompute_amg_state=precompute_amg_state,
-        checkpoint_path=checkpoint_path, device=device, prefer_decoder=prefer_decoder,
+        checkpoint_path=checkpoint_path, decoder_path=decoder_path,
+        device=device, prefer_decoder=prefer_decoder,
         use_cli=True,
     )
 
@@ -148,4 +151,5 @@ def main():
         model_type=args.model_type, tile_shape=args.tile_shape, halo=args.halo,
         checkpoint_path=args.checkpoint, device=args.device,
         precompute_amg_state=args.precompute_amg_state, prefer_decoder=args.prefer_decoder,
+        decoder_path=args.decoder_path,
     )