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..e33323e64 100644 --- a/micro_sam/sam_annotator/_state.py +++ b/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 diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py index c92f4499a..99c9fa5eb 100644 --- a/micro_sam/sam_annotator/_tooltips.py +++ b/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.", diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py index 9baca7e59..b0e9662a6 100644 --- a/micro_sam/sam_annotator/_widgets.py +++ b/micro_sam/sam_annotator/_widgets.py @@ -644,6 +644,49 @@ 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 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 = seed_widget() + 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() + 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 +697,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 +707,10 @@ def _reset_tracking_state(viewer): state.annotator._tracking_widget[2].value = "1" state.annotator._tracking_widget[2].choices = ["1"] + panel = seed_widget() + if panel is not None: + panel.refresh_status() + # # Widgets implemented with magicgui. @@ -748,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. @@ -765,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 @@ -1149,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. @@ -1183,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" @@ -1198,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() @@ -1413,10 +1498,12 @@ def _validate_layers( state.annotator._require_layers() if not automatic_segmentation: - # Check prompts layer. + # 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 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) @@ -1542,7 +1629,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): @@ -2743,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. @@ -3079,6 +3192,16 @@ 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.""" + # 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 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) + 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 +3450,23 @@ def propagate_track(track_id, division_frame): ) z_scribbles = vutil.get_scribble_slices(box_layer, track_id=track_id) have_positive_cue = False + + # 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(): # 'add_mask_prompts' skips an empty mask as well. + continue + 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) + 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 +3491,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]) @@ -3842,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 @@ -3874,6 +4021,11 @@ def clear(self, viewer=None): else: i = int(self._viewer.dims.point[0]) vutil.clear_annotations_slice(self._viewer, i=i) + # Seeds are prompts too, so clearing a frame drops them. + _drop_seed_masks(frame=i) + panel = seed_widget() + if panel is not None: + panel.refresh_status() state = AnnotatorState() if state.interactive_segmenter is not None: @@ -3881,6 +4033,237 @@ def clear(self, viewer=None): gc.collect() +class SeedTrackWidget(_WidgetBase): + """Seed the current track from an object of an existing segmentation. + + 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. + parent: The parent Qt widget. + """ + + 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): + # Label left, dropdown right, so the row does not take two lines. + 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")) + 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 Mask") + self.seed_button.setToolTip(get_tooltip("seed_track", "seed_button")) + self.seed_button.clicked.connect(self.seed) + 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) + 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) + button_row.addWidget(self.commit_button) + self.layout().addLayout(button_row) + + self.status = QtWidgets.QLabel() + 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() + 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 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(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.""" + 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] + + # '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, + ) + 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.""" + 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]) + 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 = [] + 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} from object {ids} on {len(seeded)} frame(s).") + + def drop(self): + """Drop the current track's seeds, so its drawn prompts count again.""" + state = AnnotatorState() + track_id = state.current_track_id + 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() + + 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 4da71de01..7397457d5 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.""" + panel = widgets.seed_widget() + if panel is not None: + panel.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 @@ -418,7 +432,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 +441,43 @@ 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. + + Kept separate from 'committed_objects' so it stays a reference while refined tracks commit. + """ + 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 Mask' reads the loaded result. + panel = widgets.seed_widget() + if panel is not None: + panel.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 +493,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 +544,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")) 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)