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.