Support tracking with existing segmentations in tracking annotator

micro_sam/_cli.py

@@ -150,19 +150,31 @@ def annotator_segmentation(
 
 
 @annotator_group.command("tracking")
+@click.option(
+    "-t", "--tracking_result",
+    help="Optional filepath to an existing tracking result (a TYX label volume, e.g. from trackastra). "
+    "Its objects can be used to seed a track, so that SAM2 refines and propagates an existing mask."
+)
+@click.option(
+    "-tk", "--tracking_key", default=None,
+    help="The key for opening the tracking result. Same rules as '--key'."
+)
 @_interactive_options
 def annotator_tracking(
     input_, key, embedding_path, model_type, checkpoint_path, decoder_path, device, tile_shape, halo,
+    tracking_result, tracking_key,
 ):
     """Interactively track cells in a timeseries."""
     from .util import load_image_data
     from .sam_annotator.annotator_tracking import annotator_tracking as run_annotator_tracking
     from .v2.util import DEFAULT_MODEL
 
     image = load_image_data(input_, key=key)
+    tracks = None if tracking_result is None else load_image_data(tracking_result, key=tracking_key)
     run_annotator_tracking(
         image,
         embedding_path=embedding_path,
+        tracking_result=tracks,
         model_type=model_type or DEFAULT_MODEL,
         tile_shape=tile_shape or None,
         halo=halo or None,

micro_sam/sam_annotator/_state.py

@@ -75,6 +75,9 @@ class AnnotatorState(metaclass=Singleton):
     lineage: Optional[Dict] = None
     committed_lineages: Optional[List[Dict]] = None
 
+    # Masks seeded from an existing segmentation, as '{track_id: {frame: boolean mask}}'.
+    seed_masks: Dict[int, Dict[int, np.ndarray]] = field(default_factory=dict)
+
     # Dict to keep track of all widgets, so that we can update their states.
     widgets: Dict[str, QWidget] = field(default_factory=dict)
 
@@ -430,6 +433,7 @@ def reset_state(self):
         self.current_track_id = None
         self.lineage = None
         self.committed_lineages = None
+        self.seed_masks = {}
         self.z_range = None
         self.data_signature = None
         self.interactive_segmenter = None

micro_sam/sam_annotator/_tooltips.py

@@ -86,6 +86,15 @@
         "track_state": "Select the state of the current annotation. Choose 'division' if the object divides in the current frame.",  # noqa
         "export_button": "Export the committed tracking result in the chosen format (CTC, GEFF or TrackMate XML).",  # noqa
     },
+    "seed_track": {
+        "panel": "Fix an existing tracking result: seed a track from one of its objects, or commit a track that is already correct.",  # noqa
+        "mask_layer": "Select the label layer with the tracking result to fix.",
+        "seed_button": "Use the selected object of the mask layer as the prompt for the current track.",
+        "drop_button": "Drop the seeds of the current track, so that the prompts drawn for it count again.",
+        "all_frames": "Seed every frame the object appears on, instead of only the current frame.",
+        "refine_masks": "Fit the seeded masks to the image before tracking. Uncheck to keep them as they are.",
+        "commit_button": "Commit only the current track. A track that needs no fixing is copied straight from the mask layer.",  # noqa
+    },
     "batch_annotator": {
         "folder": "Select the folder with the images to annotate.",
         "output_folder": "Select the folder for saving the segmentation results.",