micro_sam.automatic_segmentation

  1import os
  2import warnings
  3from glob import glob
  4from tqdm import tqdm
  5from pathlib import Path
  6from functools import partial
  7from typing import Dict, List, Optional, Union, Tuple, Literal
  8
  9import numpy as np
 10import imageio.v3 as imageio
 11
 12from torch_em.data.datasets.util import split_kwargs
 13
 14from . import util
 15from .instance_segmentation import (
 16    get_instance_segmentation_generator, get_decoder, AMGBase,
 17    AutomaticMaskGenerator, TiledAutomaticMaskGenerator,
 18    AutomaticPromptGenerator, TiledAutomaticPromptGenerator,
 19    InstanceSegmentationWithDecoder, TiledInstanceSegmentationWithDecoder,
 20    DEFAULT_SEGMENTATION_MODE_WITH_DECODER,
 21)
 22from .multi_dimensional_segmentation import automatic_3d_segmentation, automatic_tracking_implementation
 23
 24
 25def get_predictor_and_segmenter(
 26    model_type: str,
 27    checkpoint: Optional[Union[os.PathLike, str]] = None,
 28    device: str = None,
 29    segmentation_mode: Optional[Literal["amg", "ais", "apg"]] = None,
 30    is_tiled: bool = False,
 31    predictor=None,
 32    state=None,
 33    **kwargs,
 34) -> Tuple[util.SamPredictor, Union[AMGBase, InstanceSegmentationWithDecoder]]:
 35    f"""Get the Segment Anything model and class for automatic instance segmentation.
 36
 37    Args:
 38        model_type: The Segment Anything model choice.
 39        checkpoint: The filepath to the stored model checkpoints.
 40        device: The torch device. By default, automatically chooses the best available device.
 41        segmentation_mode: The segmentation mode. One of 'amg', 'ais', or 'apg'.
 42            By default, '{DEFAULT_SEGMENTATION_MODE_WITH_DECODER}' is used
 43            if a decoder is passed, otherwise 'amg' is used.
 44        is_tiled: Whether to return segmenter for performing segmentation in tiling window style.
 45            By default, set to 'False'.
 46        predictor: The pre-loaded predictor (optional).
 47        state: The pre-loaded state (optional).
 48        kwargs: Keyword arguments for the automatic mask generation class.
 49
 50    Returns:
 51        The Segment Anything model.
 52        The automatic instance segmentation class.
 53    """
 54    # Get the predictor and state for Segment Anything Model.
 55    if predictor is None:
 56        device = util.get_device(device=device)
 57        predictor, state = util.get_sam_model(
 58            model_type=model_type, device=device, checkpoint_path=checkpoint, return_state=True
 59        )
 60    else:
 61        assert state is not None
 62
 63    if segmentation_mode in (None, "auto"):
 64        segmentation_mode = DEFAULT_SEGMENTATION_MODE_WITH_DECODER if "decoder_state" in state else "amg"
 65
 66    if segmentation_mode.lower() == "amg":
 67        decoder = None
 68    else:
 69        if "decoder_state" not in state:
 70            raise RuntimeError(
 71                f"You have passed 'segmentation_mode={segmentation_mode}', but your model does not contain a decoder."
 72            )
 73        decoder_state = state["decoder_state"]
 74        decoder = get_decoder(image_encoder=predictor.model.image_encoder, decoder_state=decoder_state, device=device)
 75
 76    segmenter = get_instance_segmentation_generator(
 77        predictor=predictor, is_tiled=is_tiled, decoder=decoder, segmentation_mode=segmentation_mode, **kwargs
 78    )
 79    return predictor, segmenter
 80
 81
 82def _add_suffix_to_output_path(output_path: Union[str, os.PathLike], suffix: str) -> str:
 83    fpath = Path(output_path).resolve()
 84    fext = fpath.suffix if fpath.suffix else ".tif"
 85    return str(fpath.with_name(f"{fpath.stem}{suffix}{fext}"))
 86
 87
 88def automatic_tracking(
 89    predictor: util.SamPredictor,
 90    segmenter: Union[AMGBase, InstanceSegmentationWithDecoder],
 91    input_path: Union[Union[os.PathLike, str], np.ndarray],
 92    output_path: Optional[Union[os.PathLike, str]] = None,
 93    embedding_path: Optional[Union[os.PathLike, str]] = None,
 94    key: Optional[str] = None,
 95    tile_shape: Optional[Tuple[int, int]] = None,
 96    halo: Optional[Tuple[int, int]] = None,
 97    verbose: bool = True,
 98    return_embeddings: bool = False,
 99    annotate: bool = False,
100    batch_size: int = 1,
101    mode: str = "greedy",
102    tracking_model: str = "general_2d",
103    **generate_kwargs
104) -> Tuple[np.ndarray, List[Dict]]:
105    """Run automatic tracking for the input timeseries.
106
107    Args:
108        predictor: The Segment Anything model.
109        segmenter: The automatic instance segmentation class.
110        input_path: input_path: The input image file(s). Can either be a single image file (e.g. tif or png),
111            or a container file (e.g. hdf5 or zarr).
112        output_path: The folder where the tracking outputs will be saved in CTC format.
113        embedding_path: The path where the embeddings are cached already / will be saved.
114        key: The key to the input file. This is needed for container files (eg. hdf5 or zarr)
115            or to load several images as 3d volume. Provide a glob patterm, eg. "*.tif", for this case.
116        tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
117        halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
118        verbose: Verbosity flag. By default, set to 'True'.
119        return_embeddings: Whether to return the precomputed image embeddings.
120            By default, does not return the embeddings.
121        annotate: Whether to activate the annotator for continue annotation process.
122            By default, does not activate the annotator.
123        batch_size: The batch size to compute image embeddings over tiles / z-planes.
124            By default, does it sequentially, i.e. one after the other.
125        mode: The trackastra linking solver. One of 'greedy_nodiv', 'greedy' or 'ilp'.
126            'ilp' uses the motile solver. By default, set to 'greedy'.
127        tracking_model: The pretrained trackastra model to use. By default, set to 'general_2d'.
128        generate_kwargs: optional keyword arguments for the generate function of the AMG, APG, or AIS class.
129
130    Returns:
131        The tracking result as a timeseries, where each object is labeled by its track id.
132        The lineages representing cell divisions, stored as a dictionary.
133    """
134    # Load the input image file.
135    # We assume that it has to be read from file if it is a str or pathlike.
136    # Otherwise we assume it is a numpy array like object.
137    image_data = util.load_image_data(input_path, key) if isinstance(input_path, (str, os.PathLike)) else input_path
138
139    if (image_data.ndim != 3) and (image_data.ndim != 4 and image_data.shape[-1] != 3):
140        raise ValueError(f"The inputs does not match the shape expectation of 3d inputs: {image_data.shape}")
141
142    gap_closing, min_time_extent = generate_kwargs.get("gap_closing"), generate_kwargs.get("min_time_extent")
143    segmentation, lineage, image_embeddings = automatic_tracking_implementation(
144        image_data,
145        predictor,
146        segmenter,
147        embedding_path=embedding_path,
148        gap_closing=gap_closing,
149        min_time_extent=min_time_extent,
150        tile_shape=tile_shape,
151        halo=halo,
152        verbose=verbose,
153        batch_size=batch_size,
154        mode=mode,
155        tracking_model=tracking_model,
156        return_embeddings=True,
157        output_folder=output_path,
158        **generate_kwargs,
159    )
160
161    if annotate:
162        # TODO We need to support initialization of the tracking annotator with the tracking result for this.
163        raise NotImplementedError("Annotation after running the automated tracking is currently not supported.")
164
165    if return_embeddings:
166        return segmentation, lineage, image_embeddings
167    else:
168        return segmentation, lineage
169
170
171def automatic_instance_segmentation(
172    predictor: util.SamPredictor,
173    segmenter: Union[AMGBase, InstanceSegmentationWithDecoder],
174    input_path: Union[Union[os.PathLike, str], np.ndarray],
175    output_path: Optional[Union[os.PathLike, str]] = None,
176    embedding_path: Optional[Union[os.PathLike, str]] = None,
177    mask_path: Optional[Union[Union[os.PathLike, str], np.ndarray]] = None,
178    key: Optional[str] = None,
179    mask_key: Optional[str] = None,
180    ndim: Optional[int] = None,
181    tile_shape: Optional[Tuple[int, int]] = None,
182    halo: Optional[Tuple[int, int]] = None,
183    verbose: bool = True,
184    return_embeddings: bool = False,
185    annotate: bool = False,
186    batch_size: int = 1,
187    **generate_kwargs
188) -> np.ndarray:
189    """Run automatic segmentation for the input image.
190
191    Args:
192        predictor: The Segment Anything model.
193        segmenter: The automatic instance segmentation class.
194        input_path: input_path: The input image file(s). Can either be a single image file (e.g. tif or png),
195            or a container file (e.g. hdf5 or zarr).
196        output_path: The output path where the instance segmentations will be saved.
197        embedding_path: The path where the embeddings are cached already / will be saved.
198        mask_path: The path to an optional foreground mask. Areas outside of the foreground will not be processed.
199        key: The key to the input file. This is needed for container files (eg. hdf5 or zarr)
200            or to load several images as 3d volume. Provide a glob patterm, eg. "*.tif", for this case.
201        mask_key: The key to the (optional) foreground mask.
202        ndim: The dimensionality of the data. By default the dimensionality of the data will be used.
203            If you have RGB data you have to specify this explicitly, e.g. pass ndim=2 for 2d segmentation of RGB.
204        tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
205        halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
206        verbose: Verbosity flag. By default, set to 'True'.
207        return_embeddings: Whether to return the precomputed image embeddings.
208            By default, does not return the embeddings.
209        annotate: Whether to activate the annotator for continue annotation process.
210            By default, does not activate the annotator.
211        batch_size: The batch size to compute image embeddings over tiles / z-planes.
212            By default, does it sequentially, i.e. one after the other.
213        generate_kwargs: optional keyword arguments for the generate function of the AMG or AIS class.
214
215    Returns:
216        The segmentation result.
217    """
218    # Avoid overwriting already stored segmentations.
219    if output_path is not None:
220        output_path = Path(output_path).with_suffix(".tif")
221        if os.path.exists(output_path):
222            print(f"The segmentation results are already stored at '{os.path.abspath(output_path)}'.")
223            return
224
225    # We assume that it has to be read from file if it is a str or pathlike.
226    # Otherwise we assume it is a numpy array like object.
227    image_data = util.load_image_data(input_path, key) if isinstance(input_path, (str, os.PathLike)) else input_path
228
229    ndim = image_data.ndim if ndim is None else ndim
230
231    # Load the mask defining foreground if it was given.
232    if mask_path is None:
233        mask = None
234    else:
235        mask = util.load_image_data(mask_path, mask_key) if isinstance(mask_path, (str, os.PathLike)) else mask_path
236
237    if ndim == 2:
238        if (image_data.ndim != 2) and (image_data.ndim != 3 and image_data.shape[-1] != 3):
239            raise ValueError(f"The inputs does not match the shape expectation of 2d inputs: {image_data.shape}")
240
241        # Precompute the image embeddings.
242        image_embeddings = util.precompute_image_embeddings(
243            predictor=predictor,
244            input_=image_data,
245            save_path=embedding_path,
246            ndim=ndim,
247            tile_shape=tile_shape,
248            halo=halo,
249            verbose=verbose,
250            batch_size=batch_size,
251            mask=mask,
252        )
253        initialize_kwargs = dict(image=image_data, image_embeddings=image_embeddings, verbose=verbose)
254        if mask is not None:
255            initialize_kwargs["mask"] = mask
256
257        # If we run AIS with tiling then we use the same tile shape for the watershed postprocessing.
258        # In this case, we also add the batch size to the initialize kwargs,
259        # so that the segmentation decoder can be applied in a batched fashion.
260        if isinstance(segmenter, InstanceSegmentationWithDecoder) and tile_shape is not None:
261            if not isinstance(segmenter, TiledAutomaticPromptGenerator):
262                generate_kwargs.update({"tile_shape": tile_shape, "halo": halo})
263            initialize_kwargs["batch_size"] = batch_size
264
265        segmenter.initialize(**initialize_kwargs)
266        instances = segmenter.generate(**generate_kwargs)
267
268    else:
269        if (image_data.ndim != 3) and (image_data.ndim != 4 and image_data.shape[-1] != 3):
270            raise ValueError(f"The inputs does not match the shape expectation of 3d inputs: {image_data.shape}")
271        if mask is not None:
272            raise NotImplementedError
273
274        instances, image_embeddings = automatic_3d_segmentation(
275            volume=image_data,
276            predictor=predictor,
277            segmentor=segmenter,
278            embedding_path=embedding_path,
279            tile_shape=tile_shape,
280            halo=halo,
281            verbose=verbose,
282            return_embeddings=True,
283            batch_size=batch_size,
284            **generate_kwargs
285        )
286
287    # Before starting to annotate, if at all desired, store the automatic segmentations in the first stage.
288    if output_path is not None:
289        _output_path = _add_suffix_to_output_path(output_path, "_automatic") if annotate else output_path
290        imageio.imwrite(_output_path, instances, compression="zlib")
291        if verbose:
292            print(f"The automatic segmentation results are stored at '{os.path.abspath(_output_path)}'.")
293
294    # Allow opening the automatic segmentation in the annotator for further annotation, if desired.
295    if annotate:
296        from micro_sam.sam_annotator import annotator_2d, annotator_3d
297        annotator_function = annotator_2d if ndim == 2 else annotator_3d
298
299        viewer = annotator_function(
300            image=image_data,
301            model_type=predictor.model_name,
302            embedding_path=image_embeddings,  # Providing the precomputed image embeddings.
303            segmentation_result=instances,  # Initializes the automatic segmentation to the annotator.
304            tile_shape=tile_shape,
305            halo=halo,
306            return_viewer=True,  # Returns the viewer, which allows the user to store the updated segmentations.
307        )
308
309        # Start the GUI here
310        import napari
311        napari.run()
312
313        # We extract the segmentation in "committed_objects" layer, where the user either:
314        # a) Performed interactive segmentation / corrections and committed them, OR
315        # b) Did not do anything and closed the annotator, i.e. keeps the segmentations as it is.
316        instances = viewer.layers["committed_objects"].data
317
318        # Save the instance segmentation, if 'output_path' provided.
319        if output_path is not None:
320            imageio.imwrite(output_path, instances, compression="zlib")
321            if verbose:
322                print(f"The final segmentation results are stored at '{os.path.abspath(output_path)}'.")
323
324    if return_embeddings:
325        return instances, image_embeddings
326    else:
327        return instances
328
329
330def _get_inputs_from_paths(paths, pattern):
331    "Function to get all filepaths in a directory."
332
333    if isinstance(paths, str):
334        paths = [paths]
335
336    fpaths = []
337    for path in paths:
338        if os.path.isfile(path):  # It is just one filepath.
339            fpaths.append(path)
340        else:  # Otherwise, if the path is a directory, fetch all inputs provided with a pattern.
341            assert pattern is not None, \
342                f"You must provide a pattern to search for files in the directory: '{os.path.abspath(path)}'."
343            fpaths.extend(glob(os.path.join(path, pattern)))
344
345    return fpaths
346
347
348def main():
349    """@private"""
350    import argparse
351
352    available_models = list(util.get_model_names())
353    available_models = ", ".join(available_models)
354
355    parser = argparse.ArgumentParser(
356        description="Run automatic segmentation or tracking for 2d, 3d or timeseries data.\n"
357        "Either a single input file or multiple input files are supported. You can specify multiple files "
358        "by either providing multiple filepaths to the '--i/--input_paths' argument, or by providing an argument "
359        "to '--pattern' to use a wildcard pattern ('*') for selecting multiple files.\n"
360        "NOTE: for automatic 3d segmentation or tracking the data has to be stored as volume / timeseries, "
361        "stacking individual tif images is not supported.\n"
362        "Segmentation is performed using one of the three modes supported by micro_sam: \n"
363        "automatic instance segmentation (AIS), automatic prompt generation (APG) or automatic mask generation (AMG).\n"
364        "In addition to the options listed below, "
365        "you can also passed additional arguments for these three segmentation modes:\n"
366        "For AIS: '--center_distance_threshold', '--boundary_distance_threshold' and other arguments of `InstanceSegmentationWithDecoder.generate`.\n"  # noqa
367        "FOR APG: '--center_distance_threshold', '--boundary_distance_threshold' and other arguments of `AutomaticPromptGenerator.generate`.\n"  # noqa
368        "For AMG: '--pred_iou_thresh', '--stability_score_thresh' and other arguments of `AutomaticMaskGenerator.generate`."  # noqa
369    )
370    parser.add_argument(
371        "-i", "--input_path", required=True, type=str, nargs="+",
372        help="The filepath(s) to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "  # noqa
373        "or elf.io.open_file (e.g. hdf5, zarr, mrc). For the latter you also need to pass the 'key' parameter."
374    )
375    parser.add_argument(
376        "-o", "--output_path", required=True, type=str,
377        help="The filepath to store the results. If multiple inputs are provied, "
378        "this should be a folder. For a single image, you should provide the path to a tif file for the output segmentation."  # noqa
379        "NOTE: Segmentation results are stored as tif files, tracking results in the CTC fil format ."
380    )
381    parser.add_argument(
382        "-e", "--embedding_path", default=None, type=str,
383        help="An optional path where the embeddings will be saved. If multiple inputs are provided, "
384        "this should be a folder. Otherwise you can store embeddings in single zarr file."
385    )
386    parser.add_argument(
387        "--pattern", type=str, help="Pattern / wildcard for selecting files in a folder. To select all files use '*'."
388    )
389    parser.add_argument(
390        "-k", "--key", default=None, type=str,
391        help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
392        "for an image stack it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
393    )
394    parser.add_argument(
395        "-m", "--model_type", default=util._DEFAULT_MODEL, type=str,
396        help=f"The segment anything model that will be used, one of {available_models}."
397    )
398    parser.add_argument(
399        "-c", "--checkpoint", default=None, type=str, help="Checkpoint from which the SAM model will be loaded."
400    )
401    parser.add_argument(
402        "--tile_shape", nargs="+", type=int, help="The tile shape for using tiled prediction.", default=None
403    )
404    parser.add_argument(
405        "--halo", nargs="+", type=int, help="The halo for using tiled prediction.", default=None
406    )
407    parser.add_argument(
408        "-n", "--ndim", default=None, type=int,
409        help="The number of spatial dimensions in the data. Please specify this if your data has a channel dimension."
410    )
411    parser.add_argument(
412        "--mode", default="auto", type=str,
413        help="The choice of automatic segmentation mode. Either 'auto', 'amg', 'apg', or 'ais'."
414    )
415    parser.add_argument(
416        "--annotate", action="store_true",
417        help="Whether to continue annotation after the automatic segmentation is generated."
418    )
419    parser.add_argument(
420        "-d", "--device", default=None, type=str,
421        help="The device to use for the predictor. Can be one of 'cuda', 'cpu' or 'mps' (only MAC)."
422        "By default the most performant available device will be selected."
423    )
424    parser.add_argument(
425        "--batch_size", type=int, default=1,
426        help="The batch size for computing image embeddings over tiles or z-plane. "
427        "By default, computes the image embeddings for one tile / z-plane at a time."
428    )
429    parser.add_argument(
430        "--tracking", action="store_true", help="Run automatic tracking instead of instance segmentation. "
431        "NOTE: It is only supported for timeseries inputs."
432    )
433    parser.add_argument(
434        "-v", "--verbose", action="store_true", help="Whether to allow verbosity of outputs."
435    )
436
437    args, parameter_args = parser.parse_known_args()
438
439    def _convert_argval(value):
440        # The values for the parsed arguments need to be in the expected input structure as provided.
441        # i.e. integers and floats should be in their original types.
442        try:
443            return int(value)
444        except ValueError:
445            return float(value)
446
447    # NOTE: the script below allows the possibility to catch additional parsed arguments which correspond to
448    # the automatic segmentation post-processing parameters (eg. 'center_distance_threshold' in AIS)
449    extra_kwargs = {
450        parameter_args[i].lstrip("--"): _convert_argval(parameter_args[i + 1]) for i in range(0, len(parameter_args), 2)
451    }
452
453    # Separate extra arguments as per where they should be passed in the automatic segmentation class.
454    # This is done to ensure the extra arguments are allocated to the desired location.
455    # eg. for AMG, 'points_per_side' is expected by '__init__',
456    # and 'stability_score_thresh' is expected in 'generate' method.
457    mode = args.mode
458    if mode in ("auto", None):
459        # We have to load the state to see if we have a decoder in this case.
460        device = util.get_device(device=args.device)
461        predictor, state = util.get_sam_model(
462            model_type=args.model_type, device=device, checkpoint_path=args.checkpoint, return_state=True
463        )
464        mode = DEFAULT_SEGMENTATION_MODE_WITH_DECODER if "decoder_state" in state else "amg"
465    else:
466        predictor, state = None, None
467
468    if mode.lower() == "amg":
469        segmenter_class = AutomaticMaskGenerator if args.tile_shape is None else TiledAutomaticMaskGenerator
470    elif mode.lower() == "ais":
471        segmenter_class = InstanceSegmentationWithDecoder if args.tile_shape is None else\
472            TiledInstanceSegmentationWithDecoder
473    elif mode.lower() == "apg":
474        segmenter_class = AutomaticPromptGenerator if args.tile_shape is None else TiledAutomaticPromptGenerator
475    else:
476        raise ValueError(f"Invalid segmentation_mode: {mode}. Choose one of 'amg', 'ais', or 'apg'.")
477    init_kwargs, generate_kwargs = split_kwargs(segmenter_class, **extra_kwargs)
478
479    predictor, segmenter = get_predictor_and_segmenter(
480        model_type=args.model_type,
481        checkpoint=args.checkpoint,
482        device=args.device,
483        segmentation_mode=mode,
484        is_tiled=args.tile_shape is not None,
485        predictor=predictor,
486        state=state,
487        **init_kwargs,
488    )
489
490    # Get the filepaths to input images (and other paths to store stuff, eg. segmentations and embeddings)
491    # Check whether the inputs are as expected, otherwise assort them.
492    input_paths = _get_inputs_from_paths(args.input_path, args.pattern)
493    assert len(input_paths) > 0, "'micro-sam' could not extract any image data internally."
494
495    output_path = args.output_path
496    embedding_path = args.embedding_path
497    has_one_input = len(input_paths) == 1
498
499    instance_seg_function = automatic_tracking if args.tracking else partial(
500        automatic_instance_segmentation, ndim=args.ndim
501    )
502
503    # Run automatic segmentation per image.
504    for input_path in tqdm(input_paths, desc="Run automatic " + ("tracking" if args.tracking else "segmentation")):
505        if has_one_input:  # When we have only one image / volume.
506            _embedding_fpath = embedding_path  # Either folder or zarr file, would work for both.
507
508            output_fdir = os.path.splitext(output_path)[0]
509            os.makedirs(output_fdir, exist_ok=True)
510
511            # For tracking, we ensure that the output path is a folder,
512            # i.e. does not have an extension. We throw a warning if the user provided an extension.
513            if args.tracking:
514                if os.path.splitext(output_path)[-1]:
515                    warnings.warn(
516                        f"The output folder has an extension '{os.path.splitext(output_path)[-1]}'. "
517                        "We remove it and treat it as a folder to store tracking outputs in CTC format."
518                    )
519                _output_fpath = output_fdir
520            else:  # Otherwise, we can store outputs for user directly in the provided filepath, ensuring extension .tif
521                _output_fpath = f"{output_fdir}.tif"
522
523        else:  # When we have multiple images.
524            # Get the input filename, without the extension.
525            input_name = str(Path(input_path).stem)
526
527            # Let's check the 'embedding_path'.
528            if embedding_path is None:  # For computing embeddings on-the-fly, we don't care about the path logic.
529                _embedding_fpath = embedding_path
530            else:  # Otherwise, store each embeddings inside a folder.
531                embedding_folder = os.path.splitext(embedding_path)[0]  # Treat the provided embedding path as folder.
532                os.makedirs(embedding_folder, exist_ok=True)
533                _embedding_fpath = os.path.join(embedding_folder, f"{input_name}.zarr")  # Create each embedding file.
534
535            # Get the output folder name.
536            output_folder = os.path.splitext(output_path)[0]
537            os.makedirs(output_folder, exist_ok=True)
538
539            # Next, let's check for output file to store segmentation (or tracks).
540            if args.tracking:  # For tracking, we store CTC outputs in subfolders, with input_name as folder.
541                _output_fpath = os.path.join(output_folder, input_name)
542            else:  # Otherwise, store each result inside a folder.
543                _output_fpath = os.path.join(output_folder, f"{input_name}.tif")
544
545        instance_seg_function(
546            predictor=predictor,
547            segmenter=segmenter,
548            input_path=input_path,
549            output_path=_output_fpath,
550            embedding_path=_embedding_fpath,
551            key=args.key,
552            tile_shape=args.tile_shape,
553            halo=args.halo,
554            annotate=args.annotate,
555            verbose=args.verbose,
556            batch_size=args.batch_size,
557            **generate_kwargs,
558        )
def get_predictor_and_segmenter( model_type: str, checkpoint: Union[str, os.PathLike, NoneType] = None, device: str = None, segmentation_mode: Optional[Literal['amg', 'ais', 'apg']] = None, is_tiled: bool = False, predictor=None, state=None, **kwargs) -> Tuple[mobile_sam.predictor.SamPredictor, Union[micro_sam.instance_segmentation.AMGBase, micro_sam.instance_segmentation.InstanceSegmentationWithDecoder]]:
26def get_predictor_and_segmenter(
27    model_type: str,
28    checkpoint: Optional[Union[os.PathLike, str]] = None,
29    device: str = None,
30    segmentation_mode: Optional[Literal["amg", "ais", "apg"]] = None,
31    is_tiled: bool = False,
32    predictor=None,
33    state=None,
34    **kwargs,
35) -> Tuple[util.SamPredictor, Union[AMGBase, InstanceSegmentationWithDecoder]]:
36    f"""Get the Segment Anything model and class for automatic instance segmentation.
37
38    Args:
39        model_type: The Segment Anything model choice.
40        checkpoint: The filepath to the stored model checkpoints.
41        device: The torch device. By default, automatically chooses the best available device.
42        segmentation_mode: The segmentation mode. One of 'amg', 'ais', or 'apg'.
43            By default, '{DEFAULT_SEGMENTATION_MODE_WITH_DECODER}' is used
44            if a decoder is passed, otherwise 'amg' is used.
45        is_tiled: Whether to return segmenter for performing segmentation in tiling window style.
46            By default, set to 'False'.
47        predictor: The pre-loaded predictor (optional).
48        state: The pre-loaded state (optional).
49        kwargs: Keyword arguments for the automatic mask generation class.
50
51    Returns:
52        The Segment Anything model.
53        The automatic instance segmentation class.
54    """
55    # Get the predictor and state for Segment Anything Model.
56    if predictor is None:
57        device = util.get_device(device=device)
58        predictor, state = util.get_sam_model(
59            model_type=model_type, device=device, checkpoint_path=checkpoint, return_state=True
60        )
61    else:
62        assert state is not None
63
64    if segmentation_mode in (None, "auto"):
65        segmentation_mode = DEFAULT_SEGMENTATION_MODE_WITH_DECODER if "decoder_state" in state else "amg"
66
67    if segmentation_mode.lower() == "amg":
68        decoder = None
69    else:
70        if "decoder_state" not in state:
71            raise RuntimeError(
72                f"You have passed 'segmentation_mode={segmentation_mode}', but your model does not contain a decoder."
73            )
74        decoder_state = state["decoder_state"]
75        decoder = get_decoder(image_encoder=predictor.model.image_encoder, decoder_state=decoder_state, device=device)
76
77    segmenter = get_instance_segmentation_generator(
78        predictor=predictor, is_tiled=is_tiled, decoder=decoder, segmentation_mode=segmentation_mode, **kwargs
79    )
80    return predictor, segmenter
def automatic_tracking( predictor: mobile_sam.predictor.SamPredictor, segmenter: Union[micro_sam.instance_segmentation.AMGBase, micro_sam.instance_segmentation.InstanceSegmentationWithDecoder], input_path: Union[os.PathLike, str, numpy.ndarray], output_path: Union[str, os.PathLike, NoneType] = None, embedding_path: Union[str, os.PathLike, NoneType] = None, key: Optional[str] = None, tile_shape: Optional[Tuple[int, int]] = None, halo: Optional[Tuple[int, int]] = None, verbose: bool = True, return_embeddings: bool = False, annotate: bool = False, batch_size: int = 1, mode: str = 'greedy', tracking_model: str = 'general_2d', **generate_kwargs) -> Tuple[numpy.ndarray, List[Dict]]:
 89def automatic_tracking(
 90    predictor: util.SamPredictor,
 91    segmenter: Union[AMGBase, InstanceSegmentationWithDecoder],
 92    input_path: Union[Union[os.PathLike, str], np.ndarray],
 93    output_path: Optional[Union[os.PathLike, str]] = None,
 94    embedding_path: Optional[Union[os.PathLike, str]] = None,
 95    key: Optional[str] = None,
 96    tile_shape: Optional[Tuple[int, int]] = None,
 97    halo: Optional[Tuple[int, int]] = None,
 98    verbose: bool = True,
 99    return_embeddings: bool = False,
100    annotate: bool = False,
101    batch_size: int = 1,
102    mode: str = "greedy",
103    tracking_model: str = "general_2d",
104    **generate_kwargs
105) -> Tuple[np.ndarray, List[Dict]]:
106    """Run automatic tracking for the input timeseries.
107
108    Args:
109        predictor: The Segment Anything model.
110        segmenter: The automatic instance segmentation class.
111        input_path: input_path: The input image file(s). Can either be a single image file (e.g. tif or png),
112            or a container file (e.g. hdf5 or zarr).
113        output_path: The folder where the tracking outputs will be saved in CTC format.
114        embedding_path: The path where the embeddings are cached already / will be saved.
115        key: The key to the input file. This is needed for container files (eg. hdf5 or zarr)
116            or to load several images as 3d volume. Provide a glob patterm, eg. "*.tif", for this case.
117        tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
118        halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
119        verbose: Verbosity flag. By default, set to 'True'.
120        return_embeddings: Whether to return the precomputed image embeddings.
121            By default, does not return the embeddings.
122        annotate: Whether to activate the annotator for continue annotation process.
123            By default, does not activate the annotator.
124        batch_size: The batch size to compute image embeddings over tiles / z-planes.
125            By default, does it sequentially, i.e. one after the other.
126        mode: The trackastra linking solver. One of 'greedy_nodiv', 'greedy' or 'ilp'.
127            'ilp' uses the motile solver. By default, set to 'greedy'.
128        tracking_model: The pretrained trackastra model to use. By default, set to 'general_2d'.
129        generate_kwargs: optional keyword arguments for the generate function of the AMG, APG, or AIS class.
130
131    Returns:
132        The tracking result as a timeseries, where each object is labeled by its track id.
133        The lineages representing cell divisions, stored as a dictionary.
134    """
135    # Load the input image file.
136    # We assume that it has to be read from file if it is a str or pathlike.
137    # Otherwise we assume it is a numpy array like object.
138    image_data = util.load_image_data(input_path, key) if isinstance(input_path, (str, os.PathLike)) else input_path
139
140    if (image_data.ndim != 3) and (image_data.ndim != 4 and image_data.shape[-1] != 3):
141        raise ValueError(f"The inputs does not match the shape expectation of 3d inputs: {image_data.shape}")
142
143    gap_closing, min_time_extent = generate_kwargs.get("gap_closing"), generate_kwargs.get("min_time_extent")
144    segmentation, lineage, image_embeddings = automatic_tracking_implementation(
145        image_data,
146        predictor,
147        segmenter,
148        embedding_path=embedding_path,
149        gap_closing=gap_closing,
150        min_time_extent=min_time_extent,
151        tile_shape=tile_shape,
152        halo=halo,
153        verbose=verbose,
154        batch_size=batch_size,
155        mode=mode,
156        tracking_model=tracking_model,
157        return_embeddings=True,
158        output_folder=output_path,
159        **generate_kwargs,
160    )
161
162    if annotate:
163        # TODO We need to support initialization of the tracking annotator with the tracking result for this.
164        raise NotImplementedError("Annotation after running the automated tracking is currently not supported.")
165
166    if return_embeddings:
167        return segmentation, lineage, image_embeddings
168    else:
169        return segmentation, lineage

Run automatic tracking for the input timeseries.

Arguments:
  • predictor: The Segment Anything model.
  • segmenter: The automatic instance segmentation class.
  • input_path: input_path: The input image file(s). Can either be a single image file (e.g. tif or png), or a container file (e.g. hdf5 or zarr).
  • output_path: The folder where the tracking outputs will be saved in CTC format.
  • embedding_path: The path where the embeddings are cached already / will be saved.
  • 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.
  • tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
  • halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
  • verbose: Verbosity flag. By default, set to 'True'.
  • return_embeddings: Whether to return the precomputed image embeddings. By default, does not return the embeddings.
  • annotate: Whether to activate the annotator for continue annotation process. By default, does not activate the annotator.
  • batch_size: The batch size to compute image embeddings over tiles / z-planes. By default, does it sequentially, i.e. one after the other.
  • mode: The trackastra linking solver. One of 'greedy_nodiv', 'greedy' or 'ilp'. 'ilp' uses the motile solver. By default, set to 'greedy'.
  • tracking_model: The pretrained trackastra model to use. By default, set to 'general_2d'.
  • generate_kwargs: optional keyword arguments for the generate function of the AMG, APG, or AIS class.
Returns:

The tracking result as a timeseries, where each object is labeled by its track id. The lineages representing cell divisions, stored as a dictionary.

def automatic_instance_segmentation( predictor: mobile_sam.predictor.SamPredictor, segmenter: Union[micro_sam.instance_segmentation.AMGBase, micro_sam.instance_segmentation.InstanceSegmentationWithDecoder], input_path: Union[os.PathLike, str, numpy.ndarray], output_path: Union[str, os.PathLike, NoneType] = None, embedding_path: Union[str, os.PathLike, NoneType] = None, mask_path: Union[os.PathLike, str, numpy.ndarray, NoneType] = 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, verbose: bool = True, return_embeddings: bool = False, annotate: bool = False, batch_size: int = 1, **generate_kwargs) -> numpy.ndarray:
172def automatic_instance_segmentation(
173    predictor: util.SamPredictor,
174    segmenter: Union[AMGBase, InstanceSegmentationWithDecoder],
175    input_path: Union[Union[os.PathLike, str], np.ndarray],
176    output_path: Optional[Union[os.PathLike, str]] = None,
177    embedding_path: Optional[Union[os.PathLike, str]] = None,
178    mask_path: Optional[Union[Union[os.PathLike, str], np.ndarray]] = None,
179    key: Optional[str] = None,
180    mask_key: Optional[str] = None,
181    ndim: Optional[int] = None,
182    tile_shape: Optional[Tuple[int, int]] = None,
183    halo: Optional[Tuple[int, int]] = None,
184    verbose: bool = True,
185    return_embeddings: bool = False,
186    annotate: bool = False,
187    batch_size: int = 1,
188    **generate_kwargs
189) -> np.ndarray:
190    """Run automatic segmentation for the input image.
191
192    Args:
193        predictor: The Segment Anything model.
194        segmenter: The automatic instance segmentation class.
195        input_path: input_path: The input image file(s). Can either be a single image file (e.g. tif or png),
196            or a container file (e.g. hdf5 or zarr).
197        output_path: The output path where the instance segmentations will be saved.
198        embedding_path: The path where the embeddings are cached already / will be saved.
199        mask_path: The path to an optional foreground mask. Areas outside of the foreground will not be processed.
200        key: The key to the input file. This is needed for container files (eg. hdf5 or zarr)
201            or to load several images as 3d volume. Provide a glob patterm, eg. "*.tif", for this case.
202        mask_key: The key to the (optional) foreground mask.
203        ndim: The dimensionality of the data. By default the dimensionality of the data will be used.
204            If you have RGB data you have to specify this explicitly, e.g. pass ndim=2 for 2d segmentation of RGB.
205        tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
206        halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
207        verbose: Verbosity flag. By default, set to 'True'.
208        return_embeddings: Whether to return the precomputed image embeddings.
209            By default, does not return the embeddings.
210        annotate: Whether to activate the annotator for continue annotation process.
211            By default, does not activate the annotator.
212        batch_size: The batch size to compute image embeddings over tiles / z-planes.
213            By default, does it sequentially, i.e. one after the other.
214        generate_kwargs: optional keyword arguments for the generate function of the AMG or AIS class.
215
216    Returns:
217        The segmentation result.
218    """
219    # Avoid overwriting already stored segmentations.
220    if output_path is not None:
221        output_path = Path(output_path).with_suffix(".tif")
222        if os.path.exists(output_path):
223            print(f"The segmentation results are already stored at '{os.path.abspath(output_path)}'.")
224            return
225
226    # We assume that it has to be read from file if it is a str or pathlike.
227    # Otherwise we assume it is a numpy array like object.
228    image_data = util.load_image_data(input_path, key) if isinstance(input_path, (str, os.PathLike)) else input_path
229
230    ndim = image_data.ndim if ndim is None else ndim
231
232    # Load the mask defining foreground if it was given.
233    if mask_path is None:
234        mask = None
235    else:
236        mask = util.load_image_data(mask_path, mask_key) if isinstance(mask_path, (str, os.PathLike)) else mask_path
237
238    if ndim == 2:
239        if (image_data.ndim != 2) and (image_data.ndim != 3 and image_data.shape[-1] != 3):
240            raise ValueError(f"The inputs does not match the shape expectation of 2d inputs: {image_data.shape}")
241
242        # Precompute the image embeddings.
243        image_embeddings = util.precompute_image_embeddings(
244            predictor=predictor,
245            input_=image_data,
246            save_path=embedding_path,
247            ndim=ndim,
248            tile_shape=tile_shape,
249            halo=halo,
250            verbose=verbose,
251            batch_size=batch_size,
252            mask=mask,
253        )
254        initialize_kwargs = dict(image=image_data, image_embeddings=image_embeddings, verbose=verbose)
255        if mask is not None:
256            initialize_kwargs["mask"] = mask
257
258        # If we run AIS with tiling then we use the same tile shape for the watershed postprocessing.
259        # In this case, we also add the batch size to the initialize kwargs,
260        # so that the segmentation decoder can be applied in a batched fashion.
261        if isinstance(segmenter, InstanceSegmentationWithDecoder) and tile_shape is not None:
262            if not isinstance(segmenter, TiledAutomaticPromptGenerator):
263                generate_kwargs.update({"tile_shape": tile_shape, "halo": halo})
264            initialize_kwargs["batch_size"] = batch_size
265
266        segmenter.initialize(**initialize_kwargs)
267        instances = segmenter.generate(**generate_kwargs)
268
269    else:
270        if (image_data.ndim != 3) and (image_data.ndim != 4 and image_data.shape[-1] != 3):
271            raise ValueError(f"The inputs does not match the shape expectation of 3d inputs: {image_data.shape}")
272        if mask is not None:
273            raise NotImplementedError
274
275        instances, image_embeddings = automatic_3d_segmentation(
276            volume=image_data,
277            predictor=predictor,
278            segmentor=segmenter,
279            embedding_path=embedding_path,
280            tile_shape=tile_shape,
281            halo=halo,
282            verbose=verbose,
283            return_embeddings=True,
284            batch_size=batch_size,
285            **generate_kwargs
286        )
287
288    # Before starting to annotate, if at all desired, store the automatic segmentations in the first stage.
289    if output_path is not None:
290        _output_path = _add_suffix_to_output_path(output_path, "_automatic") if annotate else output_path
291        imageio.imwrite(_output_path, instances, compression="zlib")
292        if verbose:
293            print(f"The automatic segmentation results are stored at '{os.path.abspath(_output_path)}'.")
294
295    # Allow opening the automatic segmentation in the annotator for further annotation, if desired.
296    if annotate:
297        from micro_sam.sam_annotator import annotator_2d, annotator_3d
298        annotator_function = annotator_2d if ndim == 2 else annotator_3d
299
300        viewer = annotator_function(
301            image=image_data,
302            model_type=predictor.model_name,
303            embedding_path=image_embeddings,  # Providing the precomputed image embeddings.
304            segmentation_result=instances,  # Initializes the automatic segmentation to the annotator.
305            tile_shape=tile_shape,
306            halo=halo,
307            return_viewer=True,  # Returns the viewer, which allows the user to store the updated segmentations.
308        )
309
310        # Start the GUI here
311        import napari
312        napari.run()
313
314        # We extract the segmentation in "committed_objects" layer, where the user either:
315        # a) Performed interactive segmentation / corrections and committed them, OR
316        # b) Did not do anything and closed the annotator, i.e. keeps the segmentations as it is.
317        instances = viewer.layers["committed_objects"].data
318
319        # Save the instance segmentation, if 'output_path' provided.
320        if output_path is not None:
321            imageio.imwrite(output_path, instances, compression="zlib")
322            if verbose:
323                print(f"The final segmentation results are stored at '{os.path.abspath(output_path)}'.")
324
325    if return_embeddings:
326        return instances, image_embeddings
327    else:
328        return instances

Run automatic segmentation for the input image.

Arguments:
  • predictor: The Segment Anything model.
  • segmenter: The automatic instance segmentation class.
  • input_path: input_path: The input image file(s). Can either be a single image file (e.g. tif or png), 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.
  • halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
  • verbose: Verbosity flag. By default, set to 'True'.
  • return_embeddings: Whether to return the precomputed image embeddings. By default, does not return the embeddings.
  • annotate: Whether to activate the annotator for continue annotation process. By default, does not activate the annotator.
  • batch_size: The batch size to compute image embeddings over tiles / z-planes. By default, does it sequentially, i.e. one after the other.
  • generate_kwargs: optional keyword arguments for the generate function of the AMG or AIS class.
Returns:

The segmentation result.