Update automatic tracking for storing tracks in CTC format (#1045)

anwai98 · web-flow · commit 83de9c5e8f98 · 2025-05-28T18:45:12.000+02:00
Enable saving automatic tracking results and update automatic segmentation doc
diff --git a/.gitignore b/.gitignore
@@ -195,6 +195,7 @@ iterative_prompting_results/
 *.tif
 *.zip
 *MACOSX
+hela_ctc
 clf-test-data
 
 # Related to i2k workshop folders.
diff --git a/examples/README.md b/examples/README.md
@@ -1,10 +1,14 @@
 # Examples
 
 Examples for using the `micro_sam` annotation tools:
-- `annotator_2d.py`: run the interactive 2d annotation tool.
-- `annotator_3d.py`: run the interactive 3d annotation tool.
-- `annotator_tracking.py`: run the interactive tracking annotation tool.
-- `image_series_annotator.py`: run the annotation tool for a series of images.
+- `annotator_2d.py`: Run the interactive 2d annotation tool.
+- `annotator_3d.py`: Run the interactive 3d annotation tool.
+- `annotator_tracking.py`: Run the interactive tracking annotation tool.
+- `image_series_annotator.py`: Run the annotation tool for a series of images.
+
+And python scripts for automatic segmentation and tracking:
+- `automatic_segmentation.py`: Run automatic segmentation on 2d images.
+- `automatic_tracking.py`: Run automatic tracking on 2d timeseries images.
 
 We provide Jupyter Notebooks for using automatic segmentation and finetuning on some example data in the [notebooks](../notebooks/) folder.
 
diff --git a/examples/automatic_tracking.py b/examples/automatic_tracking.py
@@ -0,0 +1,56 @@
+import os
+
+from elf.io import open_file
+
+from micro_sam.util import get_cache_directory
+from micro_sam.sample_data import fetch_tracking_example_data
+from micro_sam.automatic_segmentation import automatic_tracking, get_predictor_and_segmenter
+
+
+DATA_CACHE = os.path.join(get_cache_directory(), "sample_data")
+EMBEDDING_CACHE = os.path.join(get_cache_directory(), "embeddings")
+os.makedirs(EMBEDDING_CACHE, exist_ok=True)
+
+
+def example_automatic_tracking(use_finetuned_model):
+    """Run automatic tracking for data from the cell tracking challenge.
+    """
+    # Download the example tracking data.
+    example_data = fetch_tracking_example_data(DATA_CACHE)
+
+    # Load the example data (load the sequence of tif files as timeseries)
+    with open_file(example_data, mode="r") as f:
+        timeseries = f["*.tif"]
+
+    if use_finetuned_model:
+        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-ctc-vit_b_lm.zarr")
+        model_type = "vit_b_lm"
+    else:
+        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-ctc.zarr")
+        model_type = "vit_h"
+
+    predictor, segmenter = get_predictor_and_segmenter(model_type=model_type, amg=False)
+
+    masks_tracked, _ = automatic_tracking(
+        predictor=predictor,
+        segmenter=segmenter,
+        input_path=timeseries[:],
+        output_path="./hela_ctc",
+        embedding_path=embedding_path,
+    )
+
+    import napari
+    v = napari.Viewer()
+    v.add_image(timeseries)
+    v.add_labels(masks_tracked)
+    napari.run()
+
+
+def main():
+    # Whether to use the fine-tuned SAM model.
+    use_finetuned_model = True
+    example_automatic_tracking(use_finetuned_model)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/micro_sam/automatic_segmentation.py b/micro_sam/automatic_segmentation.py
@@ -1,8 +1,9 @@
 import os
-from functools import partial
+import warnings
 from glob import glob
 from tqdm import tqdm
 from pathlib import Path
+from functools import partial
 from typing import Dict, List, Optional, Union, Tuple
 
 import numpy as np
@@ -95,7 +96,7 @@ def automatic_tracking(
         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.
+        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.
@@ -111,11 +112,9 @@ def automatic_tracking(
         generate_kwargs: optional keyword arguments for the generate function of the AMG 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.
     """
-    if output_path is not None:
-        # TODO implement saving tracking results in CTC format and use it to save the result here.
-        raise NotImplementedError("Saving the tracking result to file is currently not supported.")
-
     # Load the input image file.
     if isinstance(input_path, np.ndarray):
         image_data = input_path
@@ -142,7 +141,8 @@ def automatic_tracking(
         halo=halo,
         verbose=verbose,
         batch_size=batch_size,
-        return_image_embeddings=True,
+        return_embeddings=True,
+        output_folder=output_path,
         **generate_kwargs,
     )
 
@@ -335,12 +335,6 @@ def _get_inputs_from_paths(paths, pattern):
     return fpaths
 
 
-def _has_extension(fpath: Union[os.PathLike, str]) -> bool:
-    "Returns whether the provided path has an extension or not."
-    breakpoint()
-    return bool(os.path.splitext(fpath)[1])
-
-
 def main():
     """@private"""
     import argparse
@@ -349,23 +343,34 @@ def main():
     available_models = ", ".join(available_models)
 
     parser = argparse.ArgumentParser(
-        description="Run automatic segmentation for an image using either automatic instance segmentation (AIS) \n"
-        "or automatic mask generation (AMG). In addition to the arguments explained below,\n"
+        description="Run automatic segmentation or tracking for 2d, 3d or timeseries data.\n"
+        "Either a single input file or multiple input files are supported. You can specify multiple files "
+        "by either providing multiple filepaths to the '--i/--input_paths' argument, or by providing an argument "
+        "to '--pattern' to use a wildcard pattern ('*') for selecting multiple files.\n"
+        "NOTE: for automatic 3d segmentation or tracking the data has to be stored as volume / timeseries, "
+        "stacking individual tif images is not supported.\n"
+        "Segmentation is performed using one of the two modes supported by micro_sam: \n"
+        "automatic instance segmentation (AIS) or automatic mask generation (AMG).\n"
+        "In addition to the options listed below, "
         "you can also passed additional arguments for these two segmentation modes:\n"
         "For AIS: '--center_distance_threshold', '--boundary_distance_threshold' and other arguments of `InstanceSegmentationWithDecoder.generate`."  # noqa
         "For AMG: '--pred_iou_thresh', '--stability_score_thresh' and other arguments of `AutomaticMaskGenerator.generate`."  # noqa
     )
     parser.add_argument(
         "-i", "--input_path", required=True, type=str, nargs="+",
-        help="The filepath to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "
+        help="The filepath(s) to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "  # noqa
         "or elf.io.open_file (e.g. hdf5, zarr, mrc). For the latter you also need to pass the 'key' parameter."
     )
     parser.add_argument(
         "-o", "--output_path", required=True, type=str,
-        help="The filepath to store the instance segmentation. The current support stores segmentation in a 'tif' file."
+        help="The filepath to store the results. If multiple inputs are provied, "
+        "this should be a folder. For a single image, you should provide the path to a tif file for the output segmentation."  # noqa
+        "NOTE: Segmentation results are stored as tif files, tracking results in the CTC fil format ."
     )
     parser.add_argument(
-        "-e", "--embedding_path", default=None, type=str, help="The path where the embeddings will be saved."
+        "-e", "--embedding_path", default=None, type=str,
+        help="An optional path where the embeddings will be saved. If multiple inputs are provided, "
+        "this should be a folder. Otherwise you can store embeddings in single zarr file."
     )
     parser.add_argument(
         "--pattern", type=str, help="Pattern / wildcard for selecting files in a folder. To select all files use '*'."
@@ -411,8 +416,8 @@ def main():
         "By default, computes the image embeddings for one tile / z-plane at a time."
     )
     parser.add_argument(
-        "--tracking", action="store_true", help="Run tracking instead of instance segmentation. "
-        "Only supported for timeseries inputs.."
+        "--tracking", action="store_true", help="Run automatic tracking instead of instance segmentation. "
+        "NOTE: It is only supported for timeseries inputs."
     )
     parser.add_argument(
         "-v", "--verbose", action="store_true", help="Whether to allow verbosity of outputs."
@@ -473,34 +478,51 @@ def _convert_argval(value):
     )
 
     # Run automatic segmentation per image.
-    for path in tqdm(input_paths, desc="Run automatic segmentation"):
-        if has_one_input:  # if we have one image only.
-            _output_fpath = str(Path(output_path).with_suffix(".tif"))
-            _embedding_fpath = embedding_path
-
-        else:  # if we have multiple image, we need to make the other target filepaths compatible.
-            # Let's check for 'embedding_path'.
-            _embedding_fpath = embedding_path
-            if embedding_path:
-                if _has_extension(embedding_path):  # in this case, use filename as addl. suffix to provided path.
-                    _embedding_fpath = str(Path(embedding_path).with_suffix(".zarr"))
-                    _embedding_fpath = _embedding_fpath.replace(".zarr", f"_{Path(path).stem}.zarr")
-                else:   # otherwise, for directory, use image filename for multiple images.
-                    os.makedirs(embedding_path, exist_ok=True)
-                    _embedding_fpath = os.path.join(embedding_path, Path(os.path.basename(path)).with_suffix(".zarr"))
-
-            # Next, let's check for output file to store segmentation.
-            if _has_extension(output_path):  # in this case, use filename as addl. suffix to provided path.
-                _output_fpath = str(Path(output_path).with_suffix(".tif"))
-                _output_fpath = _output_fpath.replace(".tif", f"_{Path(path).stem}.tif")
-            else:  # otherwise, for directory, use image filename for multiple images.
-                os.makedirs(output_path, exist_ok=True)
-                _output_fpath = os.path.join(output_path, Path(os.path.basename(path)).with_suffix(".tif"))
+    for input_path in tqdm(input_paths, desc="Run automatic " + ("tracking" if args.tracking else "segmentation")):
+        if has_one_input:  # When we have only one image / volume.
+            _embedding_fpath = embedding_path  # Either folder or zarr file, would work for both.
+
+            output_fdir = os.path.splitext(output_path)[0]
+            os.makedirs(output_fdir, exist_ok=True)
+
+            # For tracking, we ensure that the output path is a folder,
+            # i.e. does not have an extension. We throw a warning if the user provided an extension.
+            if args.tracking:
+                if os.path.splitext(output_path)[-1]:
+                    warnings.warn(
+                        f"The output folder has an extension '{os.path.splitext(output_path)[-1]}'. "
+                        "We remove it and treat it as a folder to store tracking outputs in CTC format."
+                    )
+                _output_fpath = output_fdir
+            else:  # Otherwise, we can store outputs for user directly in the provided filepath, ensuring extension .tif
+                _output_fpath = f"{output_fdir}.tif"
+
+        else:  # When we have multiple images.
+            # Get the input filename, without the extension.
+            input_name = str(Path(input_path).stem)
+
+            # Let's check the 'embedding_path'.
+            if embedding_path is None:  # For computing embeddings on-the-fly, we don't care about the path logic.
+                _embedding_fpath = embedding_path
+            else:  # Otherwise, store each embeddings inside a folder.
+                embedding_folder = os.path.splitext(embedding_path)[0]  # Treat the provided embedding path as folder.
+                os.makedirs(embedding_folder, exist_ok=True)
+                _embedding_fpath = os.path.join(embedding_folder, f"{input_name}.zarr")  # Create each embedding file.
+
+            # Get the output folder name.
+            output_folder = os.path.splitext(output_path)[0]
+            os.makedirs(output_folder, exist_ok=True)
+
+            # Next, let's check for output file to store segmentation (or tracks).
+            if args.tracking:  # For tracking, we store CTC outputs in subfolders, with input_name as folder.
+                _output_fpath = os.path.join(output_folder, input_name)
+            else:  # Otherwise, store each result inside a folder.
+                _output_fpath = os.path.join(output_folder, f"{input_name}.tif")
 
         instance_seg_function(
             predictor=predictor,
             segmenter=segmenter,
-            input_path=path,
+            input_path=input_path,
             output_path=_output_fpath,
             embedding_path=_embedding_fpath,
             key=args.key,
diff --git a/micro_sam/evaluation/evaluation.py b/micro_sam/evaluation/evaluation.py
@@ -17,7 +17,6 @@
 from elf.evaluation import mean_segmentation_accuracy
 
 from ..util import load_image_data
-from ..automatic_segmentation import _has_extension
 
 
 def _run_evaluation(gt_paths, prediction_paths, verbose=True, thresholds=None):
@@ -206,7 +205,7 @@ def main():
     def _get_inputs_from_paths(paths, key):
         fpaths = []
         for path in paths:
-            if _has_extension(path):  # it is just one filepath and we check whether we can access it via 'elf'.
+            if os.path.isfile(path):  # it is just one filepath and we check whether we can access it via 'elf'.
                 fpaths.append(path if key is None else load_image_data(path=path, key=key))
             else:  # otherwise, path is a directory, fetch all inputs provided with a pattern.
                 assert key is not None, \
@@ -222,7 +221,7 @@ def _get_inputs_from_paths(paths, key):
     # Check whether output path is a csv or not, if passed.
     output_path = args.output_path
     if output_path is not None:
-        if not _has_extension(output_path):  # If it is a directory, store this in "<OUTPUT_PATH>/results.csv"
+        if not os.path.isfile(output_path):  # If it is a directory, store this in "<OUTPUT_PATH>/results.csv"
             os.makedirs(output_path, exist_ok=True)
             output_path = os.path.join(output_path, "results.csv")
 
diff --git a/micro_sam/multi_dimensional_segmentation.py b/micro_sam/multi_dimensional_segmentation.py
@@ -29,7 +29,7 @@
 
 try:
     from trackastra.model import Trackastra
-    from trackastra.tracking import graph_to_napari_tracks
+    from trackastra.tracking import graph_to_ctc, graph_to_napari_tracks
 except ImportError:
     Trackastra = None
 
@@ -570,14 +570,17 @@ def _filter_lineages(lineages, tracking_result):
     return filtered_lineages
 
 
-def _tracking_impl(timeseries, segmentation, mode, min_time_extent):
+def _tracking_impl(timeseries, segmentation, mode, min_time_extent, output_folder=None):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = Trackastra.from_pretrained("general_2d", device=device)
     lineage_graph = model.track(timeseries, segmentation, mode=mode)
     track_data, parent_graph, _ = graph_to_napari_tracks(lineage_graph)
     node_to_track, lineages = _extract_tracks_and_lineages(segmentation, track_data, parent_graph)
     tracking_result = recolor_segmentation(segmentation, node_to_track)
 
+    if output_folder is not None:  # Store tracking results in CTC format.
+        graph_to_ctc(lineage_graph, segmentation, outdir=output_folder)
+
     # TODO
     # We should check if trackastra supports this already.
     # Filter out short tracks / lineages.
@@ -599,6 +602,7 @@ def track_across_frames(
     verbose: bool = True,
     pbar_init: Optional[callable] = None,
     pbar_update: Optional[callable] = None,
+    output_folder: Optional[Union[os.PathLike, str]] = None,
 ) -> Tuple[np.ndarray, List[Dict]]:
     """Track segmented objects over time.
 
@@ -615,6 +619,7 @@ def track_across_frames(
         verbose: Verbosity flag. By default, set to 'True'.
         pbar_init: Function to initialize the progress bar.
         pbar_update: Function to update the progress bar.
+        output_folder: The folder where the tracking results are stored in CTC format.
 
     Returns:
         The tracking result. Each object is colored by its track id.
@@ -628,7 +633,11 @@ def track_across_frames(
         segmentation = _preprocess_closing(segmentation, gap_closing, pbar_update)
 
     segmentation, lineage = _tracking_impl(
-        np.asarray(timeseries), segmentation, mode="greedy", min_time_extent=min_time_extent
+        timeseries=np.asarray(timeseries),
+        segmentation=segmentation,
+        mode="greedy",
+        min_time_extent=min_time_extent,
+        output_folder=output_folder,
     )
     return segmentation, lineage
 
@@ -645,6 +654,7 @@ def automatic_tracking_implementation(
     verbose: bool = True,
     return_embeddings: bool = False,
     batch_size: int = 1,
+    output_folder: Optional[Union[os.PathLike, str]] = None,
     **kwargs,
 ) -> Tuple[np.ndarray, List[Dict]]:
     """Automatically track objects in a timesries based on per-frame automatic segmentation.
@@ -665,6 +675,7 @@ def automatic_tracking_implementation(
         verbose: Verbosity flag. By default, set to 'True'.
         return_embeddings: Whether to return the precomputed image embeddings. By default, set to 'False'.
         batch_size: The batch size to compute image embeddings over planes. By default, set to '1'.
+        output_folder: The folder where the tracking results are stored in CTC format.
         kwargs: Keyword arguments for the 'generate' method of the 'segmentor'.
 
     Returns:
@@ -685,7 +696,12 @@ def automatic_tracking_implementation(
     )
 
     segmentation, lineage = track_across_frames(
-        timeseries, segmentation, gap_closing=gap_closing, min_time_extent=min_time_extent, verbose=verbose,
+        timeseries=timeseries,
+        segmentation=segmentation,
+        gap_closing=gap_closing,
+        min_time_extent=min_time_extent,
+        verbose=verbose,
+        output_folder=output_folder,
     )
 
     if return_embeddings:
