From c218637e65486d5004fb56bb7403886deaf5d890 Mon Sep 17 00:00:00 2001
From: anwai98 <anwai.archit@gmail.com>
Date: Thu, 30 Jul 2026 10:20:42 +0200
Subject: [PATCH 1/4] Support tracking with existing segmentations in tracking
 annotator

---
 micro_sam/_cli.py                             |  12 +
 micro_sam/sam_annotator/_state.py             |   6 +
 micro_sam/sam_annotator/_tooltips.py          |   5 +
 micro_sam/sam_annotator/_widgets.py           | 234 +++++++++++++++++-
 micro_sam/sam_annotator/annotator_tracking.py |  58 ++++-
 5 files changed, 311 insertions(+), 4 deletions(-)

diff --git a/micro_sam/_cli.py b/micro_sam/_cli.py
index fe487eace..0d45f4656 100644
--- a/micro_sam/_cli.py
+++ b/micro_sam/_cli.py
@@ -150,9 +150,19 @@ 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
@@ -160,9 +170,11 @@ def 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,
diff --git a/micro_sam/sam_annotator/_state.py b/micro_sam/sam_annotator/_state.py
index 88b266518..a69a63e93 100644
--- a/micro_sam/sam_annotator/_state.py
+++ b/micro_sam/sam_annotator/_state.py
@@ -75,6 +75,11 @@ class AnnotatorState(metaclass=Singleton):
     lineage: Optional[Dict] = None
     committed_lineages: Optional[List[Dict]] = None
 
+    # Mask prompts seeded from an existing segmentation (e.g. a trackastra result) in the tracking
+    # annotator, as '{track_id: {frame: boolean mask}}'. A seeded frame conditions SAM2 on that mask
+    # instead of on the prompts drawn for the track on that frame.
+    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 +435,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
diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py
index c92f4499a..c230b48b9 100644
--- a/micro_sam/sam_annotator/_tooltips.py
+++ b/micro_sam/sam_annotator/_tooltips.py
@@ -86,6 +86,11 @@
         "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": {
+        "mask_layer": "Select the label layer holding the existing masks to seed a track from, e.g. a tracking result that was loaded into the tool.",  # noqa
+        "seed_button": "Register the selected object of the mask layer as the prompt for the current track on the current frame. Pick the object with the label picker of the mask layer, or place a positive point prompt on it. Then run 'Segment Object' to refine it on this frame or to propagate it through the timeseries.",  # noqa
+        "drop_button": "Drop the seed of the current track on the current frame, so that the prompts drawn for it count again.",  # noqa
+    },
     "batch_annotator": {
         "folder": "Select the folder with the images to annotate.",
         "output_folder": "Select the folder for saving the segmentation results.",
diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py
index 9baca7e59..81a0144f5 100644
--- a/micro_sam/sam_annotator/_widgets.py
+++ b/micro_sam/sam_annotator/_widgets.py
@@ -644,6 +644,37 @@ def _create_pbar_for_threadworker(initial_description=None):
     return pbar, pbar_signals
 
 
+def _mask_to_box(mask):
+    """The (x0, y0, x1, y1) bounding box of a boolean mask, in the order the SAM2 paths expect."""
+    ys, xs = np.nonzero(mask)
+    return np.array([xs.min(), ys.min(), xs.max(), ys.max()], dtype="float32")
+
+
+def _seed_mask(track_id, frame):
+    """The mask seeded for 'track_id' on 'frame', or None if the frame is not seeded."""
+    return AnnotatorState().seed_masks.get(track_id, {}).get(frame)
+
+
+def _set_seed_mask(track_id, frame, mask):
+    """Register 'mask' as the seed of 'track_id' on 'frame'."""
+    AnnotatorState().seed_masks.setdefault(track_id, {})[frame] = mask
+
+
+def _seed_frames(track_id):
+    """The frames seeded for 'track_id', in ascending order."""
+    return sorted(AnnotatorState().seed_masks.get(track_id, {}))
+
+
+def _drop_seed_masks(frame=None):
+    """Drop the seeded masks, either on one frame or (with 'frame=None') on all frames."""
+    state = AnnotatorState()
+    if frame is None:
+        state.seed_masks = {}
+        return
+    for seeds in state.seed_masks.values():
+        seeds.pop(frame, None)
+
+
 def _reset_tracking_state(viewer):
     """Reset the tracking state.
 
@@ -654,6 +685,7 @@ def _reset_tracking_state(viewer):
     # Reset the lineage and track id.
     state.current_track_id = 1
     state.lineage = {1: []}
+    _drop_seed_masks()
 
     # Reset the layer properties.
     viewer.layers["point_prompts"].property_choices["track_id"] = ["1"]
@@ -663,6 +695,10 @@ def _reset_tracking_state(viewer):
     state.annotator._tracking_widget[2].value = "1"
     state.annotator._tracking_widget[2].choices = ["1"]
 
+    seed_widget = state.widgets.get("seed")
+    if seed_widget is not None:
+        seed_widget.refresh_status()
+
 
 #
 # Widgets implemented with magicgui.
@@ -1413,10 +1449,13 @@ def _validate_layers(
     state.annotator._require_layers()
 
     if not automatic_segmentation:
-        # Check prompts layer.
+        # Check prompts layer. A mask seeded from an existing segmentation is a prompt too, even
+        # though it lives in the state rather than in a layer (it is always empty outside tracking).
+        have_seeds = any(seeds for seeds in state.seed_masks.values())
         if (
             len(viewer.layers["prompts"].data) == 0
             and len(viewer.layers["point_prompts"].data) == 0
+            and not have_seeds
         ):
             msg = "No prompts were given. Please provide prompts to run interactive segmentation."
             return _generate_message("error", msg)
@@ -3079,6 +3118,18 @@ def _segment_slice_batched(self, z, points, labels, boxes, masks, shape):
 
     def _segment_track_on_frame(self, state, t, track_id, shape):
         """Segment a single track's object on frame 't'. Returns the binary mask or None."""
+        # A seeded frame is conditioned on its mask (refined into the object by the mask decoder,
+        # like a drawn polygon), so the prompts drawn for this track on this frame are not read.
+        # SAM2 conditions a frame on either a mask or points / boxes, never on both.
+        seed = _seed_mask(track_id, t)
+        if seed is not None:
+            if not seed.any():  # An empty seed is no cue, and has no bounding box to prompt with.
+                return None
+            seg = state.interactive_segmenter.segment_slice(
+                frame_idx=t, boxes=[_mask_to_box(seed)], masks=[seed],
+            )
+            return None if seg is None else (np.asarray(seg).squeeze() > 0)
+
         prompt_layer = self._viewer.layers["prompts"]
         prompts = vutil.collect_frame_prompts(
             self._viewer.layers["point_prompts"], prompt_layer, shape, i=t, track_id=track_id,
@@ -3327,7 +3378,23 @@ def propagate_track(track_id, division_frame):
             )
             z_scribbles = vutil.get_scribble_slices(box_layer, track_id=track_id)
             have_positive_cue = False
+
+            # Frames seeded from an existing segmentation are conditioned on their mask, refined into
+            # the object by the mask decoder first (see 'add_mask_prompts'). SAM2 conditions a frame on
+            # either a mask or points / boxes, so a seeded frame's drawn prompts are skipped below.
+            seeded_frames = set()
+            for t in _seed_frames(track_id):
+                seed = _seed_mask(track_id, t)
+                seeded_frames.add(t)
+                if not seed.any():  # An empty seed is no cue: 'add_mask_prompts' skips it as well.
+                    continue
+                state.interactive_segmenter.add_mask_prompts(frame_ids=t, masks=[seed])
+                have_positive_cue = True
+                prompted_frames.append(t)
+
             for t in sorted({int(t) for t in z_points} | {int(t) for t in z_scribbles}):
+                if t in seeded_frames:
+                    continue
                 # Exclude division markers: they signal a lineage event and bound propagation
                 # (see below), but must not be fed to SAM2 as conditioning prompts - doing so
                 # adds a second conditioning frame that corrupts the mother track's propagation.
@@ -3352,6 +3419,8 @@ def propagate_track(track_id, division_frame):
                 if box_layer.data else np.zeros(0, dtype=int)
             )
             for t in z_boxes:
+                if int(t) in seeded_frames:
+                    continue
                 boxes, _ = vutil.shape_layer_to_prompts(box_layer, shape=shape[1:], i=int(t), track_id=track_id)
                 for box in boxes:
                     state.interactive_segmenter.add_box_prompts(frame_ids=int(t), boxes=[box])
@@ -3874,6 +3943,11 @@ def clear(self, viewer=None):
         else:
             i = int(self._viewer.dims.point[0])
             vutil.clear_annotations_slice(self._viewer, i=i)
+            # The seeded masks are prompts too, so clearing a frame's annotations drops its seeds.
+            _drop_seed_masks(frame=i)
+            seed_widget = AnnotatorState().widgets.get("seed")
+            if seed_widget is not None:
+                seed_widget.refresh_status()
 
         state = AnnotatorState()
         if state.interactive_segmenter is not None:
@@ -3881,6 +3955,164 @@ def clear(self, viewer=None):
         gc.collect()
 
 
+class SeedTrackWidget(_WidgetBase):
+    """Seed the current track from an object of an existing segmentation.
+
+    Lets the user pick one object out of a label layer that was loaded into the tool (e.g. a
+    trackastra tracking result) and register its mask as the SAM2 prompt for the current track on the
+    current frame. Seeding is bookkeeping only: the mask is pushed to the video predictor when
+    'Segment Object' runs, so 'Segment Frame' refines the object on this frame and 'Apply to All
+    Frames' propagates it through the timeseries.
+
+    The object is read from the mask layer's selected label, or - when nothing is selected there -
+    from the label under the positive point prompts of the current track on the current frame. So the
+    usual workflow (click the object, hit segment) also works, without having to use the label picker.
+
+    A seeded frame is conditioned on its mask alone: SAM2 conditions a frame on either a mask or
+    points / boxes, so the prompts drawn for the track on a seeded frame are skipped while the seed
+    is registered. Clearing the annotations drops the seeds again.
+
+    Args:
+        viewer: The napari viewer.
+        parent: The parent Qt widget.
+    """
+
+    def __init__(self, viewer, parent=None):
+        super().__init__(parent=parent)
+        self._viewer = viewer
+        self._create_widget()
+
+    def _create_widget(self):
+        self.layout().addWidget(QtWidgets.QLabel("Mask Layer:"))
+        self.mask_selection = create_widget(annotation=napari.layers.Labels)
+        self.mask_selection.native.setToolTip(get_tooltip("seed_track", "mask_layer"))
+        self.layout().addWidget(self.mask_selection.native)
+
+        self.seed_button = QtWidgets.QPushButton("Seed Track From Mask")
+        self.seed_button.setToolTip(get_tooltip("seed_track", "seed_button"))
+        self.seed_button.clicked.connect(self.seed)
+        self.drop_button = QtWidgets.QPushButton("Drop Seeds")
+        self.drop_button.setToolTip(get_tooltip("seed_track", "drop_button"))
+        self.drop_button.clicked.connect(self.drop)
+        button_row = QtWidgets.QHBoxLayout()
+        button_row.addWidget(self.seed_button)
+        button_row.addWidget(self.drop_button)
+        self.layout().addLayout(button_row)
+
+        self.status = QtWidgets.QLabel()
+        self.layout().addWidget(self.status)
+        self.refresh_status()
+
+    def set_mask_layer(self, layer):
+        """Select 'layer' in the mask layer dropdown, e.g. after a tracking result was loaded."""
+        self.mask_selection.reset_choices()
+        try:
+            self.mask_selection.value = layer
+        except ValueError:  # The layer is not among the choices (it is not a label layer).
+            pass
+
+    def refresh_status(self):
+        """Update the line reporting which frames the current track is seeded on."""
+        state = AnnotatorState()
+        track_id = state.current_track_id
+        frames = [] if track_id is None else _seed_frames(track_id)
+        if frames:
+            self.status.setText(f"Track {track_id} is seeded on frame(s): {', '.join(map(str, frames))}")
+        else:
+            self.status.setText("No seeded frames for the current track.")
+
+    def _selected_mask_layer(self):
+        """The chosen label layer, or None (with a message) if it cannot be used to seed."""
+        layer = self.mask_selection.value
+        if layer is None:
+            _generate_message("error", "There is no label layer to seed the track from.")
+            return None
+        state = AnnotatorState()
+        if state.image_shape is not None and tuple(layer.data.shape) != tuple(state.image_shape):
+            _generate_message(
+                "error",
+                f"The layer '{layer.name}' has shape {tuple(layer.data.shape)}, which does not match "
+                f"the timeseries shape {tuple(state.image_shape)}."
+            )
+            return None
+        return layer
+
+    def _object_ids(self, frame, layer, track_id):
+        """The ids to seed from: the layer's selected label, else the labels under the point prompts."""
+        selected = int(layer.selected_label)
+        if selected != 0:
+            if not np.any(frame == selected):
+                _generate_message(
+                    "error",
+                    f"The selected label {selected} of '{layer.name}' is not present on this frame. "
+                    "Pick an object with the label picker, or place a positive point prompt on it."
+                )
+                return []
+            return [selected]
+
+        # Fall back to the objects under this track's positive point prompts on this frame. A lone
+        # negative point is read as a normal prompt here, not as a stop annotation, so that it is
+        # simply ignored below instead of hiding the frame's positive points.
+        points, labels = vutil.point_layer_to_prompts(
+            self._viewer.layers["point_prompts"], i=int(self._viewer.dims.point[0]),
+            track_id=track_id, with_stop_annotation=False,
+        )
+        ids = {
+            int(frame[y, x]) for y, x in np.round(points[labels == 1]).astype(int)
+            if 0 <= y < frame.shape[0] and 0 <= x < frame.shape[1] and frame[y, x] != 0
+        }
+        if not ids:
+            _generate_message(
+                "error",
+                "No object was selected. Either pick a label of the mask layer with the label picker, "
+                "or place a positive point prompt on the object you want to seed the track from."
+            )
+        return sorted(ids)
+
+    def seed(self):
+        """Register the selected object of the mask layer as the current track's seed on this frame."""
+        state = AnnotatorState()
+        if state.image_shape is None:
+            _generate_message("error", "There is no timeseries loaded yet.")
+            return
+        layer = self._selected_mask_layer()
+        if layer is None:
+            return
+
+        t = int(self._viewer.dims.point[0])
+        track_id = state.current_track_id
+        frame = np.asarray(layer.data[t])
+        object_ids = self._object_ids(frame, layer, track_id)
+        if not object_ids:
+            return
+
+        _set_seed_mask(track_id, t, np.isin(frame, object_ids))
+
+        # The seed is a prompt, so it is deliberately not painted into 'current_object': that layer
+        # holds segmentation results, and the tracking annotator reads it to decide whether a
+        # dividing track still has to be propagated. The selected object is visible in the mask
+        # layer anyway, and the status line below records the seeded frames.
+        self.refresh_status()
+        ids = ", ".join(map(str, object_ids))
+        show_info(
+            f"Seeded track {track_id} on frame {t} from object(s) {ids} of '{layer.name}'. "
+            "Run 'Segment Object' to refine or propagate it."
+        )
+
+    def drop(self):
+        """Drop the seed of the current track on this frame, so its drawn prompts count again."""
+        state = AnnotatorState()
+        t = int(self._viewer.dims.point[0])
+        track_id = state.current_track_id
+        if _seed_mask(track_id, t) is None:
+            show_info(f"Track {track_id} is not seeded on frame {t}.")
+            return
+
+        state.seed_masks[track_id].pop(t)
+        self.refresh_status()
+        show_info(f"Dropped the seed of track {track_id} on frame {t}.")
+
+
 class AutoSegmentV1Widget(_WidgetBase):
     """Automatic segmentation widget for the SAM (v1) AMG/AIS generators.
 
diff --git a/micro_sam/sam_annotator/annotator_tracking.py b/micro_sam/sam_annotator/annotator_tracking.py
index 4da71de01..99fa9e22c 100644
--- a/micro_sam/sam_annotator/annotator_tracking.py
+++ b/micro_sam/sam_annotator/annotator_tracking.py
@@ -26,6 +26,17 @@
 """@private"""
 
 
+TRACKING_RESULT_LAYER = "tracking_result"
+"""@private"""
+
+
+def _refresh_seed_status():
+    """Re-render the seed widget's status line, e.g. after the current track changed."""
+    seed_widget = AnnotatorState().widgets.get("seed")
+    if seed_widget is not None:
+        seed_widget.refresh_status()
+
+
 def _validate_tracking_model_type(model_type):
     if not model_type.startswith("hvit_"):
         raise ValueError(
@@ -100,6 +111,7 @@ def update_track_id(event):
             if new_id != track_id_menu.value:
                 track_id_menu.value = new_id
                 state.current_track_id = int(new_id)
+                _refresh_seed_status()
 
     # def update_state_boxes(event):
     #     new_state = str(box_layer.current_properties["state"][0])
@@ -112,6 +124,7 @@ def update_track_id_boxes(event):
             if new_id != track_id_menu.value:
                 track_id_menu.value = new_id
                 state.current_track_id = int(new_id)
+                _refresh_seed_status()
 
     points_layer.events.current_properties.connect(update_state)
     points_layer.events.current_properties.connect(update_track_id)
@@ -134,6 +147,7 @@ def track_id_changed(new_track_id):
         except KeyError:
             pass
         state.current_track_id = int(new_track_id)
+        _refresh_seed_status()
 
     # def state_changed_boxes(new_state):
     #     current_properties = box_layer.current_properties
@@ -341,6 +355,7 @@ def _get_widgets(self):
         )
         return {
             "interactive": interactive,
+            "seed": widgets.SeedTrackWidget(self._viewer),
             "autosegment": autotrack,
             "commit": widgets.commit_track(),
             "export": widgets.export_track(),
@@ -418,7 +433,7 @@ def _init_track_state(self):
         state.lineage = {1: []}
         state.committed_lineages = []
 
-    def _update_image(self):
+    def _update_image(self, tracking_result=None):
         super()._update_image()
         self._init_track_state()
         state = AnnotatorState()
@@ -427,11 +442,44 @@ def _update_image(self):
         else:
             state.autoseg_state = vutil._load_amg_state(state.embedding_path)
 
+        if tracking_result is not None:
+            self._set_tracking_result(tracking_result)
+
+    def _set_tracking_result(self, tracking_result):
+        """Show an existing tracking result in its own layer, to seed tracks from its masks.
+
+        It is kept separate from 'committed_objects' so that the loaded result stays available as a
+        reference and as a source of seed masks while the refined tracks are committed next to it.
+        """
+        state = AnnotatorState()
+        if state.image_shape is None:
+            raise RuntimeError("A tracking result can only be loaded once a timeseries is selected.")
+
+        tracking_result = np.asarray(tracking_result)
+        if tuple(tracking_result.shape) != tuple(state.image_shape):
+            raise ValueError(
+                f"The tracking result of shape {tracking_result.shape} does not match "
+                f"the timeseries shape {tuple(state.image_shape)}."
+            )
+
+        if TRACKING_RESULT_LAYER in self._viewer.layers:
+            layer = self._viewer.layers[TRACKING_RESULT_LAYER]
+            layer.data = tracking_result
+        else:
+            layer = self._viewer.add_labels(data=tracking_result, name=TRACKING_RESULT_LAYER)
+        if state.image_scale is not None:
+            layer.scale = state.image_scale
+
+        # Preselect it in the seed widget, so 'Seed Track From Mask' reads the loaded result.
+        seed_widget = state.widgets.get("seed")
+        if seed_widget is not None:
+            seed_widget.set_mask_layer(layer)
+
 
 def annotator_tracking(
     image: np.ndarray,
     embedding_path: Optional[str] = None,
-    # tracking_result: Optional[str] = None,
+    tracking_result: Optional[np.ndarray] = None,
     model_type: str = DEFAULT_MODEL,
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
@@ -447,6 +495,9 @@ def annotator_tracking(
     Args:
         image: The image data.
         embedding_path: Filepath for saving the precomputed embeddings.
+        tracking_result: An existing tracking result (a TYX label volume, e.g. from trackastra) to
+            load into the tool. It is shown in the 'tracking_result' layer and its objects can be
+            used to seed a track, so SAM2 refines and propagates an existing mask.
         model_type: The Segment Anything model to use. For details on the available models check out
             https://computational-cell-analytics.github.io/micro-sam/micro_sam.html#finetuned-models.
         tile_shape: Shape of tiles for tiled embedding prediction.
@@ -495,7 +546,8 @@ def annotator_tracking(
     annotator = AnnotatorTracking(viewer, reset_state=False)
 
     # Trigger layer update of the annotator so that layers have the correct shape.
-    annotator._update_image()
+    # And load the tracking result into its own layer if one was given.
+    annotator._update_image(tracking_result=tracking_result)
 
     # Add the annotator widget to the viewer and sync widgets.
     viewer.window.add_dock_widget(annotator, name=get_dock_title("tracking"))

From 259dded65a58d4f6cfafe4a768fce6cc6f01673c Mon Sep 17 00:00:00 2001
From: anwai98 <anwai.archit@gmail.com>
Date: Thu, 30 Jul 2026 10:36:12 +0200
Subject: [PATCH 2/4] Update docstrings

---
 micro_sam/sam_annotator/_state.py             |  4 +-
 micro_sam/sam_annotator/_widgets.py           | 46 +++++--------------
 micro_sam/sam_annotator/annotator_tracking.py |  3 +-
 3 files changed, 14 insertions(+), 39 deletions(-)

diff --git a/micro_sam/sam_annotator/_state.py b/micro_sam/sam_annotator/_state.py
index a69a63e93..e33323e64 100644
--- a/micro_sam/sam_annotator/_state.py
+++ b/micro_sam/sam_annotator/_state.py
@@ -75,9 +75,7 @@ class AnnotatorState(metaclass=Singleton):
     lineage: Optional[Dict] = None
     committed_lineages: Optional[List[Dict]] = None
 
-    # Mask prompts seeded from an existing segmentation (e.g. a trackastra result) in the tracking
-    # annotator, as '{track_id: {frame: boolean mask}}'. A seeded frame conditions SAM2 on that mask
-    # instead of on the prompts drawn for the track on that frame.
+    # 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.
diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py
index 81a0144f5..f46b5725f 100644
--- a/micro_sam/sam_annotator/_widgets.py
+++ b/micro_sam/sam_annotator/_widgets.py
@@ -1449,8 +1449,7 @@ def _validate_layers(
     state.annotator._require_layers()
 
     if not automatic_segmentation:
-        # Check prompts layer. A mask seeded from an existing segmentation is a prompt too, even
-        # though it lives in the state rather than in a layer (it is always empty outside tracking).
+        # A seeded mask is a prompt too, but lives in the state rather than in a layer.
         have_seeds = any(seeds for seeds in state.seed_masks.values())
         if (
             len(viewer.layers["prompts"].data) == 0
@@ -3118,16 +3117,12 @@ def _segment_slice_batched(self, z, points, labels, boxes, masks, shape):
 
     def _segment_track_on_frame(self, state, t, track_id, shape):
         """Segment a single track's object on frame 't'. Returns the binary mask or None."""
-        # A seeded frame is conditioned on its mask (refined into the object by the mask decoder,
-        # like a drawn polygon), so the prompts drawn for this track on this frame are not read.
-        # SAM2 conditions a frame on either a mask or points / boxes, never on both.
+        # SAM2 conditions a frame on either a mask or points / boxes, so a seed replaces the prompts.
         seed = _seed_mask(track_id, t)
         if seed is not None:
-            if not seed.any():  # An empty seed is no cue, and has no bounding box to prompt with.
+            if not seed.any():  # An empty seed has no bounding box to prompt with.
                 return None
-            seg = state.interactive_segmenter.segment_slice(
-                frame_idx=t, boxes=[_mask_to_box(seed)], masks=[seed],
-            )
+            seg = state.interactive_segmenter.segment_slice(frame_idx=t, boxes=[_mask_to_box(seed)], masks=[seed])
             return None if seg is None else (np.asarray(seg).squeeze() > 0)
 
         prompt_layer = self._viewer.layers["prompts"]
@@ -3379,14 +3374,12 @@ def propagate_track(track_id, division_frame):
             z_scribbles = vutil.get_scribble_slices(box_layer, track_id=track_id)
             have_positive_cue = False
 
-            # Frames seeded from an existing segmentation are conditioned on their mask, refined into
-            # the object by the mask decoder first (see 'add_mask_prompts'). SAM2 conditions a frame on
-            # either a mask or points / boxes, so a seeded frame's drawn prompts are skipped below.
+            # A seeded frame is conditioned on its mask, so its drawn prompts are skipped below.
             seeded_frames = set()
             for t in _seed_frames(track_id):
                 seed = _seed_mask(track_id, t)
                 seeded_frames.add(t)
-                if not seed.any():  # An empty seed is no cue: 'add_mask_prompts' skips it as well.
+                if not seed.any():  # 'add_mask_prompts' skips an empty mask as well.
                     continue
                 state.interactive_segmenter.add_mask_prompts(frame_ids=t, masks=[seed])
                 have_positive_cue = True
@@ -3943,7 +3936,7 @@ def clear(self, viewer=None):
         else:
             i = int(self._viewer.dims.point[0])
             vutil.clear_annotations_slice(self._viewer, i=i)
-            # The seeded masks are prompts too, so clearing a frame's annotations drops its seeds.
+            # Seeds are prompts too, so clearing a frame drops them.
             _drop_seed_masks(frame=i)
             seed_widget = AnnotatorState().widgets.get("seed")
             if seed_widget is not None:
@@ -3958,19 +3951,9 @@ def clear(self, viewer=None):
 class SeedTrackWidget(_WidgetBase):
     """Seed the current track from an object of an existing segmentation.
 
-    Lets the user pick one object out of a label layer that was loaded into the tool (e.g. a
-    trackastra tracking result) and register its mask as the SAM2 prompt for the current track on the
-    current frame. Seeding is bookkeeping only: the mask is pushed to the video predictor when
-    'Segment Object' runs, so 'Segment Frame' refines the object on this frame and 'Apply to All
-    Frames' propagates it through the timeseries.
-
-    The object is read from the mask layer's selected label, or - when nothing is selected there -
-    from the label under the positive point prompts of the current track on the current frame. So the
-    usual workflow (click the object, hit segment) also works, without having to use the label picker.
-
-    A seeded frame is conditioned on its mask alone: SAM2 conditions a frame on either a mask or
-    points / boxes, so the prompts drawn for the track on a seeded frame are skipped while the seed
-    is registered. Clearing the annotations drops the seeds again.
+    Registers the mask of one object of a label layer as the SAM2 prompt for the current track on the
+    current frame. The object is read from the layer's selected label, or from the label under the
+    track's positive point prompts. The mask is pushed to the predictor when 'Segment Object' runs.
 
     Args:
         viewer: The napari viewer.
@@ -4050,9 +4033,7 @@ def _object_ids(self, frame, layer, track_id):
                 return []
             return [selected]
 
-        # Fall back to the objects under this track's positive point prompts on this frame. A lone
-        # negative point is read as a normal prompt here, not as a stop annotation, so that it is
-        # simply ignored below instead of hiding the frame's positive points.
+        # 'with_stop_annotation=False' so a lone negative point is ignored, not read as a stop.
         points, labels = vutil.point_layer_to_prompts(
             self._viewer.layers["point_prompts"], i=int(self._viewer.dims.point[0]),
             track_id=track_id, with_stop_annotation=False,
@@ -4088,10 +4069,7 @@ def seed(self):
 
         _set_seed_mask(track_id, t, np.isin(frame, object_ids))
 
-        # The seed is a prompt, so it is deliberately not painted into 'current_object': that layer
-        # holds segmentation results, and the tracking annotator reads it to decide whether a
-        # dividing track still has to be propagated. The selected object is visible in the mask
-        # layer anyway, and the status line below records the seeded frames.
+        # Not painted into 'current_object': the division logic reads that layer for existing results.
         self.refresh_status()
         ids = ", ".join(map(str, object_ids))
         show_info(
diff --git a/micro_sam/sam_annotator/annotator_tracking.py b/micro_sam/sam_annotator/annotator_tracking.py
index 99fa9e22c..fb59ec41d 100644
--- a/micro_sam/sam_annotator/annotator_tracking.py
+++ b/micro_sam/sam_annotator/annotator_tracking.py
@@ -448,8 +448,7 @@ def _update_image(self, tracking_result=None):
     def _set_tracking_result(self, tracking_result):
         """Show an existing tracking result in its own layer, to seed tracks from its masks.
 
-        It is kept separate from 'committed_objects' so that the loaded result stays available as a
-        reference and as a source of seed masks while the refined tracks are committed next to it.
+        Kept separate from 'committed_objects' so it stays a reference while refined tracks commit.
         """
         state = AnnotatorState()
         if state.image_shape is None:

From bbb4a4e1db6882624c16f45657db038c6f5d16d1 Mon Sep 17 00:00:00 2001
From: anwai98 <anwai.archit@gmail.com>
Date: Fri, 31 Jul 2026 11:23:59 +0200
Subject: [PATCH 3/4] Add seeding options and mask refinement control to the
 tracking annotator

---
 micro_sam/sam_annotator/_tooltips.py          |   8 +-
 micro_sam/sam_annotator/_widgets.py           | 111 +++++++++++++-----
 micro_sam/v2/prompt_based_segmentation.py     |  39 +++---
 .../test_scribble_prompts.py                  |   3 +-
 .../test_volume_and_tracking_prompts.py       |   2 +-
 5 files changed, 114 insertions(+), 49 deletions(-)

diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py
index c230b48b9..f7074a689 100644
--- a/micro_sam/sam_annotator/_tooltips.py
+++ b/micro_sam/sam_annotator/_tooltips.py
@@ -87,9 +87,11 @@
         "export_button": "Export the committed tracking result in the chosen format (CTC, GEFF or TrackMate XML).",  # noqa
     },
     "seed_track": {
-        "mask_layer": "Select the label layer holding the existing masks to seed a track from, e.g. a tracking result that was loaded into the tool.",  # noqa
-        "seed_button": "Register the selected object of the mask layer as the prompt for the current track on the current frame. Pick the object with the label picker of the mask layer, or place a positive point prompt on it. Then run 'Segment Object' to refine it on this frame or to propagate it through the timeseries.",  # noqa
-        "drop_button": "Drop the seed of the current track on the current frame, so that the prompts drawn for it count again.",  # noqa
+        "mask_layer": "Select the label layer with the existing masks to seed a track from.",
+        "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.",
     },
     "batch_annotator": {
         "folder": "Select the folder with the images to annotate.",
diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py
index f46b5725f..5f0fc0335 100644
--- a/micro_sam/sam_annotator/_widgets.py
+++ b/micro_sam/sam_annotator/_widgets.py
@@ -665,6 +665,12 @@ def _seed_frames(track_id):
     return sorted(AnnotatorState().seed_masks.get(track_id, {}))
 
 
+def _seed_refine_enabled():
+    """Whether the seed widget is set to refine the seeded masks before they condition propagation."""
+    widget = AnnotatorState().widgets.get("seed")
+    return True if widget is None else bool(widget.refine_masks)
+
+
 def _drop_seed_masks(frame=None):
     """Drop the seeded masks, either on one frame or (with 'frame=None') on all frames."""
     state = AnnotatorState()
@@ -1580,7 +1586,8 @@ def _push_volume_prompts(segmenter, plan):
         elif kind == "box":
             segmenter.add_box_prompts(frame_ids=frame_id, boxes=payload, object_id=object_ids)
         else:
-            segmenter.add_mask_prompts(frame_ids=frame_id, masks=payload, object_id=object_ids)
+            # A drawn polygon or ellipse is an outline, so it is refined into the object.
+            segmenter.add_mask_prompts(frame_ids=frame_id, masks=payload, object_id=object_ids, refine=True)
 
 
 def _segment_object_2d(viewer, batched=False):
@@ -3122,6 +3129,8 @@ def _segment_track_on_frame(self, state, t, track_id, shape):
         if seed is not None:
             if not seed.any():  # An empty seed has no bounding box to prompt with.
                 return None
+            if not _seed_refine_enabled():  # Propagation conditions on it unchanged, so show that.
+                return seed
             seg = state.interactive_segmenter.segment_slice(frame_idx=t, boxes=[_mask_to_box(seed)], masks=[seed])
             return None if seg is None else (np.asarray(seg).squeeze() > 0)
 
@@ -3381,7 +3390,9 @@ def propagate_track(track_id, division_frame):
                 seeded_frames.add(t)
                 if not seed.any():  # 'add_mask_prompts' skips an empty mask as well.
                     continue
-                state.interactive_segmenter.add_mask_prompts(frame_ids=t, masks=[seed])
+                state.interactive_segmenter.add_mask_prompts(
+                    frame_ids=t, masks=[seed], object_id=None, refine=_seed_refine_enabled(),
+                )
                 have_positive_cue = True
                 prompted_frames.append(t)
 
@@ -3951,9 +3962,9 @@ def clear(self, viewer=None):
 class SeedTrackWidget(_WidgetBase):
     """Seed the current track from an object of an existing segmentation.
 
-    Registers the mask of one object of a label layer as the SAM2 prompt for the current track on the
-    current frame. The object is read from the layer's selected label, or from the label under the
-    track's positive point prompts. The mask is pushed to the predictor when 'Segment Object' runs.
+    The object is read from the layer's selected label, or from the label under the track's positive
+    point prompts, and its masks are pushed to the predictor when 'Segment Object' runs. 'All Frames'
+    seeds every frame it appears on, so propagation only fills in the frames the result is missing.
 
     Args:
         viewer: The napari viewer.
@@ -3963,20 +3974,45 @@ class SeedTrackWidget(_WidgetBase):
     def __init__(self, viewer, parent=None):
         super().__init__(parent=parent)
         self._viewer = viewer
+        self.all_frames = True
+        self.refine_masks = True
         self._create_widget()
 
     def _create_widget(self):
-        self.layout().addWidget(QtWidgets.QLabel("Mask Layer:"))
+        # Label left, dropdown right, so the row does not take two lines.
+        layer_label = QtWidgets.QLabel("Mask Layer:")
+        layer_label.setToolTip(get_tooltip("seed_track", "mask_layer"))
         self.mask_selection = create_widget(annotation=napari.layers.Labels)
         self.mask_selection.native.setToolTip(get_tooltip("seed_track", "mask_layer"))
-        self.layout().addWidget(self.mask_selection.native)
+        layer_row = QtWidgets.QHBoxLayout()
+        layer_row.addWidget(layer_label)
+        layer_row.addWidget(self.mask_selection.native, 1)
+        self.layout().addLayout(layer_row)
+
+        self.all_frames_checkbox = self._add_boolean_param(
+            "all_frames", self.all_frames, title="All Frames",
+            tooltip=get_tooltip("seed_track", "all_frames"),
+        )
+        self.all_frames_checkbox.stateChanged.connect(self._on_all_frames_changed)
+        self.refine_masks_checkbox = self._add_boolean_param(
+            "refine_masks", self.refine_masks, title="Refine Masks",
+            tooltip=get_tooltip("seed_track", "refine_masks"),
+        )
+        self.refine_masks_checkbox.stateChanged.connect(self._on_refine_masks_changed)
 
-        self.seed_button = QtWidgets.QPushButton("Seed Track From Mask")
+        self.seed_button = QtWidgets.QPushButton("Seed Mask")
         self.seed_button.setToolTip(get_tooltip("seed_track", "seed_button"))
         self.seed_button.clicked.connect(self.seed)
-        self.drop_button = QtWidgets.QPushButton("Drop Seeds")
+        self.drop_button = QtWidgets.QPushButton("Drop Seed")
         self.drop_button.setToolTip(get_tooltip("seed_track", "drop_button"))
         self.drop_button.clicked.connect(self.drop)
+
+        checkbox_row = QtWidgets.QHBoxLayout()
+        checkbox_row.addWidget(self.all_frames_checkbox)
+        checkbox_row.addStretch()
+        checkbox_row.addWidget(self.refine_masks_checkbox)
+        self.layout().addLayout(checkbox_row)
+
         button_row = QtWidgets.QHBoxLayout()
         button_row.addWidget(self.seed_button)
         button_row.addWidget(self.drop_button)
@@ -3986,6 +4022,12 @@ def _create_widget(self):
         self.layout().addWidget(self.status)
         self.refresh_status()
 
+    def _on_all_frames_changed(self, state):
+        self.all_frames = bool(state)
+
+    def _on_refine_masks_changed(self, state):
+        self.refine_masks = bool(state)
+
     def set_mask_layer(self, layer):
         """Select 'layer' in the mask layer dropdown, e.g. after a tracking result was loaded."""
         self.mask_selection.reset_choices()
@@ -3999,10 +4041,12 @@ def refresh_status(self):
         state = AnnotatorState()
         track_id = state.current_track_id
         frames = [] if track_id is None else _seed_frames(track_id)
-        if frames:
-            self.status.setText(f"Track {track_id} is seeded on frame(s): {', '.join(map(str, frames))}")
+        if not frames:
+            self.status.setText(f"Track {track_id}: no seeds")
+        elif len(frames) > 6:  # A dense seeding covers too many frames to list.
+            self.status.setText(f"Track {track_id}: {len(frames)} seeds, frames {frames[0]}-{frames[-1]}")
         else:
-            self.status.setText("No seeded frames for the current track.")
+            self.status.setText(f"Track {track_id}: seeds on frame {', '.join(map(str, frames))}")
 
     def _selected_mask_layer(self):
         """The chosen label layer, or None (with a message) if it cannot be used to seed."""
@@ -4051,7 +4095,7 @@ def _object_ids(self, frame, layer, track_id):
         return sorted(ids)
 
     def seed(self):
-        """Register the selected object of the mask layer as the current track's seed on this frame."""
+        """Register the selected object of the mask layer as the current track's seed."""
         state = AnnotatorState()
         if state.image_shape is None:
             _generate_message("error", "There is no timeseries loaded yet.")
@@ -4062,33 +4106,44 @@ def seed(self):
 
         t = int(self._viewer.dims.point[0])
         track_id = state.current_track_id
-        frame = np.asarray(layer.data[t])
-        object_ids = self._object_ids(frame, layer, track_id)
+        object_ids = self._object_ids(np.asarray(layer.data[t]), layer, track_id)
         if not object_ids:
             return
 
-        _set_seed_mask(track_id, t, np.isin(frame, object_ids))
+        # The object is identified on the current frame, then seeded on every frame it appears on.
+        frames = range(layer.data.shape[0]) if self.all_frames else [t]
+        seeded = []
+        for frame_id in frames:
+            mask = np.isin(np.asarray(layer.data[frame_id]), object_ids)
+            if mask.any():
+                _set_seed_mask(track_id, int(frame_id), mask)
+                seeded.append(int(frame_id))
 
         # Not painted into 'current_object': the division logic reads that layer for existing results.
         self.refresh_status()
         ids = ", ".join(map(str, object_ids))
-        show_info(
-            f"Seeded track {track_id} on frame {t} from object(s) {ids} of '{layer.name}'. "
-            "Run 'Segment Object' to refine or propagate it."
-        )
+        show_info(f"Seeded track {track_id} from object {ids} on {len(seeded)} frame(s).")
 
     def drop(self):
-        """Drop the seed of the current track on this frame, so its drawn prompts count again."""
+        """Drop the current track's seeds, so its drawn prompts count again."""
         state = AnnotatorState()
-        t = int(self._viewer.dims.point[0])
         track_id = state.current_track_id
-        if _seed_mask(track_id, t) is None:
-            show_info(f"Track {track_id} is not seeded on frame {t}.")
-            return
-
-        state.seed_masks[track_id].pop(t)
+        seeds = state.seed_masks.get(track_id, {})
+        if not self.all_frames:
+            t = int(self._viewer.dims.point[0])
+            if t not in seeds:
+                show_info(f"Track {track_id}: no seed on frame {t}.")
+                return
+            seeds.pop(t)
+            show_info(f"Dropped seed of track {track_id} on frame {t}.")
+        else:
+            if not seeds:
+                show_info(f"Track {track_id}: no seeds to drop.")
+                return
+            n = len(seeds)
+            seeds.clear()
+            show_info(f"Dropped {n} seed(s) of track {track_id}.")
         self.refresh_status()
-        show_info(f"Dropped the seed of track {track_id} on frame {t}.")
 
 
 class AutoSegmentV1Widget(_WidgetBase):
diff --git a/micro_sam/v2/prompt_based_segmentation.py b/micro_sam/v2/prompt_based_segmentation.py
index 73eb6c83e..27afa2d65 100644
--- a/micro_sam/v2/prompt_based_segmentation.py
+++ b/micro_sam/v2/prompt_based_segmentation.py
@@ -645,18 +645,24 @@ def _prepare_mask(self, mask):
     def add_mask_prompts(
         self,
         frame_ids: Union[int, List[int]],
-        masks: Optional[List[np.ndarray]] = None,
-        object_id: Optional[Union[int, List[int]]] = None,
+        masks: Optional[List[np.ndarray]],
+        object_id: Optional[Union[int, List[int]]],
+        refine: bool,
     ):
         """Add mask prompts (full-resolution 2d boolean masks) to the persistent SAM2 state.
 
-        A napari polygon or ellipse is filled into a mask (from 'shape_layer_to_prompts'). We first
-        refine the drawn shape into the object on its seed frame (box + soft mask-logit cue, as in
-        the per-slice path), then seed propagation with the refined mask - so the seed slice matches
-        the per-slice result instead of reproducing the raw outline. SAM2's video predictor conditions
-        a frame on either a mask or points/box (not both), so a mask prompt does not combine with
-        points on the same object/frame. A mask already pushed (same object, frame, content) is
-        skipped so re-runs only add newly drawn masks.
+        SAM2 conditions a frame on either a mask or points / box, so a mask does not combine with
+        points on the same object and frame. A mask already pushed is skipped, so a re-run only adds
+        the new ones.
+
+        Args:
+            frame_ids: The frame(s) to add the mask(s) to.
+            masks: The full-resolution 2d boolean masks.
+            object_id: The object id(s) the masks belong to.
+            refine: Whether to refine the mask into the object, cued by its bounding box and mask
+                logits, before it conditions propagation. Pass False for a mask that is already a
+                segmentation, so that it conditions propagation as it is. A refined mask matches the
+                per-slice result, which suits a drawn outline.
         """
         if masks is None or len(masks) == 0:
             return
@@ -675,19 +681,20 @@ def add_mask_prompts(
                 continue
             seen.add(signature)
 
-            # Refine the drawn shape into the object on the seed frame, then seed propagation with the
-            # refined mask. The box is the shape's bounding box (nonzero extent of the filled mask).
             ys, xs = np.nonzero(mask)
             if len(ys) == 0:
                 continue
-            box = np.array([xs.min(), ys.min(), xs.max(), ys.max()], dtype="float32")  # (x0, y0, x1, y1)
-            refined = self._image_style_predict(frame_id, box=box, mask=mask)
+
+            if refine:
+                # Refine the drawn shape into the object, cued by its bounding box and mask logits.
+                box = np.array([xs.min(), ys.min(), xs.max(), ys.max()], dtype="float32")  # (x0, y0, x1, y1)
+                mask = self._image_style_predict(frame_id, box=box, mask=mask)
 
             self.predictor.add_new_mask(
                 inference_state=self.inference_state,
                 frame_idx=frame_id,
                 obj_id=obj_id,
-                mask=self._prepare_mask(refined),
+                mask=self._prepare_mask(mask),
             )
 
     def _propagate_in_direction(
@@ -1236,7 +1243,7 @@ def add_box_prompts(self, frame_ids, boxes=None, object_id=None):
                 frame_ids=frame_ids, boxes=np.array(local_boxes), object_id=tile_ids[tile_id],
             )
 
-    def add_mask_prompts(self, frame_ids, masks=None, object_id=None):
+    def add_mask_prompts(self, frame_ids, masks, object_id, refine):
         """Add mask prompts. Each mask is routed to the tiles its filled region overlaps, cropped to
         each tile's outer block, and added there (so a mask spanning tiles is added on both sides)."""
         if masks is None or len(masks) == 0:
@@ -1256,7 +1263,7 @@ def add_mask_prompts(self, frame_ids, masks=None, object_id=None):
             for tid in _box_to_tiles(self.tiling, self.halo, box_yx):
                 self._get_segmenter(tid).add_mask_prompts(
                     frame_ids=frame_ids, masks=[_crop_mask_to_tile(self.tiling, self.halo, tid, mask)],
-                    object_id=obj_id,
+                    object_id=obj_id, refine=refine,
                 )
 
     def predict(self, update_progress=None, early_stop_patience=None, z_range=None):
diff --git a/test/test_sam_annotator/test_scribble_prompts.py b/test/test_sam_annotator/test_scribble_prompts.py
index 168274cd1..bfa15c305 100644
--- a/test/test_sam_annotator/test_scribble_prompts.py
+++ b/test/test_sam_annotator/test_scribble_prompts.py
@@ -467,7 +467,8 @@ def test_volume_scribbles_pass_layer_validation(monkeypatch):
         "point_prompts": SimpleNamespace(data=[]),
     })
     annotator = SimpleNamespace(_require_layers=lambda: None)
-    monkeypatch.setattr(_widgets, "AnnotatorState", lambda: SimpleNamespace(annotator=annotator))
+    state = SimpleNamespace(annotator=annotator, seed_masks={})
+    monkeypatch.setattr(_widgets, "AnnotatorState", lambda: state)
     result = _widgets._validate_layers(viewer)
 
     assert result is False
diff --git a/test/test_sam_annotator/test_volume_and_tracking_prompts.py b/test/test_sam_annotator/test_volume_and_tracking_prompts.py
index 0d673fa41..61c8f57fd 100644
--- a/test/test_sam_annotator/test_volume_and_tracking_prompts.py
+++ b/test/test_sam_annotator/test_volume_and_tracking_prompts.py
@@ -408,7 +408,7 @@ def run_tracking(monkeypatch, point_layer, prompt_layer, shape=(6, 32, 32)):
     })
     segmenter = Segmenter(shape)
     state = SimpleNamespace(
-        is_sam2=True, image_shape=shape, current_track_id=1, lineage={1: {}},
+        is_sam2=True, image_shape=shape, current_track_id=1, lineage={1: {}}, seed_masks={},
         image_embeddings={"input_size": (1024, 1024)}, interactive_segmenter=segmenter,
     )
     widget = SimpleNamespace(_viewer=viewer)

From 0ac49bc153395e4ade1858b1f7926cd1737c6383 Mon Sep 17 00:00:00 2001
From: anwai98 <anwai.archit@gmail.com>
Date: Mon, 3 Aug 2026 13:21:46 +0200
Subject: [PATCH 4/4] Add a track correction panel with per-track commit to the
 tracking annotator

---
 micro_sam/sam_annotator/_tooltips.py          |   4 +-
 micro_sam/sam_annotator/_widgets.py           | 166 +++++++++++++++---
 micro_sam/sam_annotator/annotator_tracking.py |  15 +-
 3 files changed, 152 insertions(+), 33 deletions(-)

diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py
index f7074a689..99c9fa5eb 100644
--- a/micro_sam/sam_annotator/_tooltips.py
+++ b/micro_sam/sam_annotator/_tooltips.py
@@ -87,11 +87,13 @@
         "export_button": "Export the committed tracking result in the chosen format (CTC, GEFF or TrackMate XML).",  # noqa
     },
     "seed_track": {
-        "mask_layer": "Select the label layer with the existing masks to seed a track from.",
+        "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.",
diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py
index 5f0fc0335..b0e9662a6 100644
--- a/micro_sam/sam_annotator/_widgets.py
+++ b/micro_sam/sam_annotator/_widgets.py
@@ -665,9 +665,15 @@ def _seed_frames(track_id):
     return sorted(AnnotatorState().seed_masks.get(track_id, {}))
 
 
+def seed_widget():
+    """The seed panel, which is nested in the interactive tracking widget."""
+    interactive = AnnotatorState().widgets.get("interactive")
+    return getattr(interactive, "seed_panel", None)
+
+
 def _seed_refine_enabled():
     """Whether the seed widget is set to refine the seeded masks before they condition propagation."""
-    widget = AnnotatorState().widgets.get("seed")
+    widget = seed_widget()
     return True if widget is None else bool(widget.refine_masks)
 
 
@@ -701,9 +707,9 @@ def _reset_tracking_state(viewer):
     state.annotator._tracking_widget[2].value = "1"
     state.annotator._tracking_widget[2].choices = ["1"]
 
-    seed_widget = state.widgets.get("seed")
-    if seed_widget is not None:
-        seed_widget.refresh_status()
+    panel = seed_widget()
+    if panel is not None:
+        panel.refresh_status()
 
 
 #
@@ -790,7 +796,7 @@ def _mask_matched_objects(seg, prev_seg, preservation_threshold):
     return preserve_mask
 
 
-def _commit_impl(viewer, layer, preserve_mode, preservation_threshold):
+def _commit_impl(viewer, layer, preserve_mode, preservation_threshold, track_ids=None, preserve_ids=False):
     state = AnnotatorState()
 
     # Check whether all layers exist as expected or create new ones automatically.
@@ -807,19 +813,24 @@ def _commit_impl(viewer, layer, preserve_mode, preservation_threshold):
     # Otherwise we run into type conversion errors later.
     dtype = viewer.layers["committed_objects"].data.dtype
     seg = viewer.layers[layer].data[bb].astype(dtype)
+    if track_ids is not None:  # Commit only the selected objects, leaving the others to work on.
+        seg = np.where(np.isin(seg, track_ids), seg, 0)
     shape = seg.shape
 
     # We parallelize these operations because they take quite long for large volumes.
 
-    # Compute the max id in the commited objects.
-    # id_offset = int(viewer.layers["committed_objects"].data.max())
-    full_shape = viewer.layers["committed_objects"].data.shape
-    id_offset = int(
-        elf.parallel.max(
-            viewer.layers["committed_objects"].data,
-            block_shape=util.get_block_shape(full_shape),
+    # Compute the max id in the commited objects. Tracking keeps its ids instead, so that a committed
+    # track is still recognisable by the id it was annotated with.
+    if preserve_ids:
+        id_offset = 0
+    else:
+        full_shape = viewer.layers["committed_objects"].data.shape
+        id_offset = int(
+            elf.parallel.max(
+                viewer.layers["committed_objects"].data,
+                block_shape=util.get_block_shape(full_shape),
+            )
         )
-    )
 
     # Compute the mask for the current object.
     # mask = seg != 0
@@ -1191,9 +1202,33 @@ def commit_track(
         commit_path: Select a file path where the committed results and prompts will be saved.
             This feature is still experimental.
     """
-    # Commit the segmentation layer.
+    commit_tracks(viewer, layer, preserve_mode, preservation_threshold, commit_path, track_ids=None)
+
+
+def commit_tracks(viewer, layer, preserve_mode, preservation_threshold, commit_path, track_ids):
+    """Commit tracks to the committed-objects layer.
+
+    Args:
+        viewer: The napari viewer.
+        layer: The layer to commit.
+        preserve_mode: How to preserve already committed objects. See `commit_track`.
+        preservation_threshold: The overlap threshold for preserving objects.
+        commit_path: Optional filepath for saving the committed results and prompts.
+        track_ids: The track ids to commit, or None for all of them. The tracks that are not
+            committed stay in the layer, and the tracking state is kept so they can be worked on.
+    """
+    selected_ids = track_ids
+
+    # Warn instead of silently merging a track into a committed object of the same id.
+    if selected_ids is not None:
+        committed = viewer.layers["committed_objects"].data
+        clashing = [i for i in selected_ids if np.any(committed == i)]
+        if clashing:
+            show_info(f"Track(s) {clashing} are already committed. They will be merged.")
+
+    # Commit the segmentation layer. Track ids are kept, see '_commit_impl'.
     id_offset, seg, mask, bb = _commit_impl(
-        viewer, layer, preserve_mode, preservation_threshold
+        viewer, layer, preserve_mode, preservation_threshold, track_ids=selected_ids, preserve_ids=True
     )
 
     # Update the lineages.
@@ -1225,7 +1260,13 @@ def commit_track(
         )
 
     if layer == "current_object":
-        vutil.clear_annotations(viewer)
+        if selected_ids is None:
+            vutil.clear_annotations(viewer)
+        else:
+            # Remove only the committed tracks, so the others stay in the layer to work on.
+            data = viewer.layers["current_object"].data
+            data[np.isin(data, selected_ids)] = 0
+            viewer.layers["current_object"].refresh()
 
     # Create / update the tracking layer.
     layer_name = "tracks"
@@ -1240,8 +1281,10 @@ def commit_track(
     else:
         viewer.add_tracks(track_data, name=layer_name, graph=parent_graph)
 
-    # Reset the tracking state.
-    _reset_tracking_state(viewer)
+    # Reset the tracking state. A partial commit keeps it, since the tracks that were not committed
+    # still need their track ids and seeds.
+    if selected_ids is None:
+        _reset_tracking_state(viewer)
 
     # Perform garbage collection.
     gc.collect()
@@ -2788,6 +2831,31 @@ def _mother_division_frame(point_layer, lineage, track_id):
     return None
 
 
+def _retarget_track_id(viewer, track_id):
+    """Make 'track_id' the current track, registering it in the lineage and the menus if it is new.
+
+    Seeding from an existing result adopts the id of the object it was seeded from, so the committed
+    track carries the id it had in that result instead of the value the menu happened to hold.
+    """
+    state = AnnotatorState()
+    track_id = int(track_id)
+    if state.current_track_id == track_id and track_id in state.lineage:
+        return
+
+    state.lineage.setdefault(track_id, [])
+    state.current_track_id = track_id
+
+    # (index 2: prompt, track_state, track_id). Setting the value fires 'track_id_changed', which
+    # writes the id back to the state and to the prompt layers.
+    track_ids = sorted(state.lineage.keys())
+    menu = state.annotator._tracking_widget[2]
+    menu.choices = list(map(str, track_ids))
+    menu.value = str(track_id)
+
+    viewer.layers["point_prompts"].property_choices["track_id"] = list(map(str, track_ids))
+    viewer.layers["prompts"].property_choices["track_id"] = list(map(str, track_ids))
+
+
 def _update_lineage(viewer, mother=None):
     """Record a division for 'mother' by seeding two daughter track ids and refreshing the menus.
 
@@ -3915,6 +3983,12 @@ def _create_widget(self):
         button_row.addWidget(self.clear_button)
         self.layout().addLayout(button_row)
 
+        # Seeding a track from an existing segmentation, collapsed by default.
+        self.seed_panel = SeedTrackWidget(self._viewer)
+        self.layout().addWidget(
+            _make_collapsible(self.seed_panel, title="Track Correction", tooltip=get_tooltip("seed_track", "panel"))
+        )
+
     def _align_menu_rows(self):
         # Each menu row is a QHBoxLayout of [QLabel, QComboBox]. Insert a stretch between them so the
         # label stays left and the (fixed-width) combo box is right-aligned. Idempotent, since the
@@ -3949,9 +4023,9 @@ def clear(self, viewer=None):
             vutil.clear_annotations_slice(self._viewer, i=i)
             # Seeds are prompts too, so clearing a frame drops them.
             _drop_seed_masks(frame=i)
-            seed_widget = AnnotatorState().widgets.get("seed")
-            if seed_widget is not None:
-                seed_widget.refresh_status()
+            panel = seed_widget()
+            if panel is not None:
+                panel.refresh_status()
 
         state = AnnotatorState()
         if state.interactive_segmenter is not None:
@@ -3980,7 +4054,7 @@ def __init__(self, viewer, parent=None):
 
     def _create_widget(self):
         # Label left, dropdown right, so the row does not take two lines.
-        layer_label = QtWidgets.QLabel("Mask Layer:")
+        layer_label = QtWidgets.QLabel("Tracking Result:")
         layer_label.setToolTip(get_tooltip("seed_track", "mask_layer"))
         self.mask_selection = create_widget(annotation=napari.layers.Labels)
         self.mask_selection.native.setToolTip(get_tooltip("seed_track", "mask_layer"))
@@ -4006,6 +4080,9 @@ def _create_widget(self):
         self.drop_button = QtWidgets.QPushButton("Drop Seed")
         self.drop_button.setToolTip(get_tooltip("seed_track", "drop_button"))
         self.drop_button.clicked.connect(self.drop)
+        self.commit_button = QtWidgets.QPushButton("Commit Track")
+        self.commit_button.setToolTip(get_tooltip("seed_track", "commit_button"))
+        self.commit_button.clicked.connect(self.commit)
 
         checkbox_row = QtWidgets.QHBoxLayout()
         checkbox_row.addWidget(self.all_frames_checkbox)
@@ -4016,6 +4093,7 @@ def _create_widget(self):
         button_row = QtWidgets.QHBoxLayout()
         button_row.addWidget(self.seed_button)
         button_row.addWidget(self.drop_button)
+        button_row.addWidget(self.commit_button)
         self.layout().addLayout(button_row)
 
         self.status = QtWidgets.QLabel()
@@ -4105,11 +4183,14 @@ def seed(self):
             return
 
         t = int(self._viewer.dims.point[0])
-        track_id = state.current_track_id
-        object_ids = self._object_ids(np.asarray(layer.data[t]), layer, track_id)
+        object_ids = self._object_ids(np.asarray(layer.data[t]), layer, state.current_track_id)
         if not object_ids:
             return
 
+        # Adopt the object's id as the track id, so the committed track keeps the id it has here.
+        _retarget_track_id(self._viewer, object_ids[0])
+        track_id = state.current_track_id
+
         # The object is identified on the current frame, then seeded on every frame it appears on.
         frames = range(layer.data.shape[0]) if self.all_frames else [t]
         seeded = []
@@ -4145,6 +4226,43 @@ def drop(self):
             show_info(f"Dropped {n} seed(s) of track {track_id}.")
         self.refresh_status()
 
+    def commit(self):
+        """Commit the current track, keeping the other tracks and the tracking state.
+
+        A track that is already correct in the mask layer is copied over as it is, so it does not
+        have to be segmented first. Otherwise the segmentation of the current track is committed.
+        """
+        state = AnnotatorState()
+        track_id = state.current_track_id
+        current_object = self._viewer.layers["current_object"]
+        source = "segmentation"
+
+        if not np.any(current_object.data == track_id):
+            # Nothing segmented for this track: copy the selected object of the mask layer instead.
+            layer = self._selected_mask_layer()
+            if layer is None:
+                return
+            t = int(self._viewer.dims.point[0])
+            object_ids = self._object_ids(np.asarray(layer.data[t]), layer, track_id)
+            if not object_ids:
+                return
+            # Adopt the object's id, so it keeps the id it has in the mask layer.
+            _retarget_track_id(self._viewer, object_ids[0])
+            track_id = state.current_track_id
+            for frame_id in range(layer.data.shape[0]):
+                mask = np.isin(np.asarray(layer.data[frame_id]), object_ids)
+                if mask.any():
+                    current_object.data[frame_id][mask] = track_id
+            current_object.refresh()
+            source = f"'{layer.name}'"
+
+        commit_tracks(
+            self._viewer, "current_object", preserve_mode="pixels", preservation_threshold=0.75,
+            commit_path=None, track_ids=[track_id],
+        )
+        self.refresh_status()
+        show_info(f"Committed track {track_id} from {source}.")
+
 
 class AutoSegmentV1Widget(_WidgetBase):
     """Automatic segmentation widget for the SAM (v1) AMG/AIS generators.
diff --git a/micro_sam/sam_annotator/annotator_tracking.py b/micro_sam/sam_annotator/annotator_tracking.py
index fb59ec41d..7397457d5 100644
--- a/micro_sam/sam_annotator/annotator_tracking.py
+++ b/micro_sam/sam_annotator/annotator_tracking.py
@@ -32,9 +32,9 @@
 
 def _refresh_seed_status():
     """Re-render the seed widget's status line, e.g. after the current track changed."""
-    seed_widget = AnnotatorState().widgets.get("seed")
-    if seed_widget is not None:
-        seed_widget.refresh_status()
+    panel = widgets.seed_widget()
+    if panel is not None:
+        panel.refresh_status()
 
 
 def _validate_tracking_model_type(model_type):
@@ -355,7 +355,6 @@ def _get_widgets(self):
         )
         return {
             "interactive": interactive,
-            "seed": widgets.SeedTrackWidget(self._viewer),
             "autosegment": autotrack,
             "commit": widgets.commit_track(),
             "export": widgets.export_track(),
@@ -469,10 +468,10 @@ def _set_tracking_result(self, tracking_result):
         if state.image_scale is not None:
             layer.scale = state.image_scale
 
-        # Preselect it in the seed widget, so 'Seed Track From Mask' reads the loaded result.
-        seed_widget = state.widgets.get("seed")
-        if seed_widget is not None:
-            seed_widget.set_mask_layer(layer)
+        # Preselect it in the seed widget, so 'Seed Mask' reads the loaded result.
+        panel = widgets.seed_widget()
+        if panel is not None:
+            panel.set_mask_layer(layer)
 
 
 def annotator_tracking(