micro_sam.sam_annotator.annotator_tracking

  1from typing import Optional, Tuple, Union, List
  2
  3import numpy as np
  4
  5import torch
  6
  7import napari
  8from magicgui.widgets import ComboBox, Container
  9
 10from .. import util
 11from . import util as vutil
 12from . import _widgets as widgets
 13from ._tooltips import get_tooltip
 14from ._state import AnnotatorState
 15from ._annotator import _AnnotatorBase
 16
 17
 18# Cyan (track) and Magenta (division)
 19STATE_COLOR_CYCLE = ["#00FFFF", "#FF00FF", ]
 20"""@private"""
 21
 22
 23# This solution is a bit hacky, so I won't move it to _widgets.py yet.
 24def create_tracking_menu(points_layer, box_layer, states, track_ids, tracking_widget=None):
 25    """@private"""
 26    state = AnnotatorState()
 27
 28    def _get_widget_menu(container, label):
 29        for w in container:
 30            if isinstance(w, ComboBox) and w.label == label:
 31                return w
 32        raise ValueError(f"ComboBox with label '{label}' not found.")
 33
 34    if tracking_widget is None:
 35        state_menu = ComboBox(
 36            label="track_state", choices=states, tooltip=get_tooltip("annotator_tracking", "track_state")
 37        )
 38        track_id_menu = ComboBox(
 39            label="track_id", choices=list(map(str, track_ids)), tooltip=get_tooltip("annotator_tracking", "track_id")
 40        )
 41        tracking_widget = Container(widgets=[state_menu, track_id_menu])
 42    else:
 43        state_menu = _get_widget_menu(tracking_widget, "track_state")
 44        track_id_menu = _get_widget_menu(tracking_widget, "track_id")
 45
 46    def update_state(event):
 47        if "state" in points_layer.current_properties:
 48            new_state = str(points_layer.current_properties["state"][0])
 49            if new_state != state_menu.value:
 50                state_menu.value = new_state
 51
 52    def update_track_id(event):
 53        if "track_id" in points_layer.current_properties:
 54            new_id = str(points_layer.current_properties["track_id"][0])
 55            if new_id != track_id_menu.value:
 56                track_id_menu.value = new_id
 57                state.current_track_id = int(new_id)
 58
 59    # def update_state_boxes(event):
 60    #     new_state = str(box_layer.current_properties["state"][0])
 61    #     if new_state != state_menu.value:
 62    #         state_menu.value = new_state
 63
 64    # napari 0.9 no longer syncs 'current_properties' with the selection on shape layers (napari/napari#9221).
 65    def update_track_id_boxes(*args):
 66        if "track_id" not in box_layer.properties:
 67            return
 68        selected = list(box_layer.selected_data)
 69        track_ids = set(box_layer.properties["track_id"][selected])
 70        if len(track_ids) != 1:  # A mixed selection would relabel every selected box.
 71            return
 72        new_id = str(track_ids.pop())
 73        if new_id != track_id_menu.value:
 74            track_id_menu.value = new_id
 75            state.current_track_id = int(new_id)
 76
 77    points_layer.events.current_properties.connect(update_state)
 78    points_layer.events.current_properties.connect(update_track_id)
 79    # box_layer.events.current_properties.connect(update_state_boxes)
 80    box_layer.selected_data.events.items_changed.connect(update_track_id_boxes)
 81
 82    def state_changed(new_state):
 83        current_properties = points_layer.current_properties
 84        current_properties["state"] = np.array([new_state])
 85        points_layer.current_properties = current_properties
 86        points_layer.refresh_colors()
 87
 88    def track_id_changed(new_track_id):
 89        current_properties = points_layer.current_properties
 90        current_properties["track_id"] = np.array([new_track_id])
 91        # Note: this fails with a key error after committing a lineage with multiple tracks.
 92        # I think this does not cause any further errors, so we just skip this.
 93        try:
 94            points_layer.current_properties = current_properties
 95        except KeyError:
 96            pass
 97        state.current_track_id = int(new_track_id)
 98
 99    # def state_changed_boxes(new_state):
100    #     current_properties = box_layer.current_properties
101    #     current_properties["state"] = np.array([new_state])
102    #     box_layer.current_properties = current_properties
103    #     box_layer.refresh_colors()
104
105    def track_id_changed_boxes(new_track_id):
106        current_properties = box_layer.current_properties
107        current_properties["track_id"] = np.array([new_track_id])
108        box_layer.current_properties = current_properties
109        state.current_track_id = int(new_track_id)
110
111    state_menu.changed.connect(state_changed)
112    track_id_menu.changed.connect(track_id_changed)
113    # state_menu.changed.connect(state_changed_boxes)
114    track_id_menu.changed.connect(track_id_changed_boxes)
115
116    state_menu.set_choice("track")
117    return tracking_widget
118
119
120class AnnotatorTracking(_AnnotatorBase):
121
122    # The tracking annotator needs different settings for the prompt layers
123    # to support the additional tracking state.
124    # That's why we over-ride this function.
125    def _require_layers(self, layer_choices: Optional[List[str]] = None):
126
127        # Check whether the image is initialized already. And use the image shape and scale for the layers.
128        state = AnnotatorState()
129        shape = self._shape if state.image_shape is None else state.image_shape
130
131        # Add the label layers for the current object, the automatic segmentation and the committed segmentation.
132        dummy_data = np.zeros(shape, dtype="uint32")
133        image_scale = state.image_scale
134
135        # Before adding new layers, we always check whether a layer with this name already exists or not.
136        if "current_object" not in self._viewer.layers:
137            if layer_choices and "current_object" in layer_choices:  # Check at 'commit' call button.
138                widgets._validation_window_for_missing_layer("current_object")
139            self._viewer.add_labels(data=dummy_data, name="current_object")
140            if image_scale is not None:
141                self._viewer.layers["current_object"].scale = image_scale
142
143        if "auto_segmentation" not in self._viewer.layers:
144            if layer_choices and "auto_segmentation" in layer_choices:  # Check at 'commit' call button.
145                widgets._validation_window_for_missing_layer("auto_segmentation")
146            self._viewer.add_labels(data=dummy_data, name="auto_segmentation")
147            if image_scale is not None:
148                self._viewer.layers["auto_segmentation"].scale = image_scale
149
150        if "committed_objects" not in self._viewer.layers:
151            if layer_choices and "committed_objects" in layer_choices:  # Check at 'commit' call button.
152                widgets._validation_window_for_missing_layer("committed_objects")
153            self._viewer.add_labels(data=dummy_data, name="committed_objects")
154            # Randomize colors so it is easy to see when object committed.
155            self._viewer.layers["committed_objects"].new_colormap()
156            if image_scale is not None:
157                self._viewer.layers["committed_objects"].scale = image_scale
158
159        # Add the point prompts layer.
160        self._point_labels = ["positive", "negative"]
161        self._track_state_labels = ["track", "division"]
162        _point_prompt_property_choices = {
163            "label": self._point_labels,
164            "state": self._track_state_labels,
165            "track_id": ["1"],  # we use string to avoid pandas warning
166        }
167
168        point_layer_mismatch = True
169        if "point_prompts" in self._viewer.layers:
170            # Check whether the 'property_choices' match or not.
171            curr_property_choices = self._viewer.layers["point_prompts"].property_choices
172            point_layer_mismatch = set(curr_property_choices.keys()) != set(_point_prompt_property_choices.keys())
173
174        if point_layer_mismatch and "point_prompts" not in self._viewer.layers:
175            self._point_prompt_layer = self._viewer.add_points(
176                name="point_prompts",
177                property_choices=_point_prompt_property_choices,
178                border_color="label",
179                border_color_cycle=vutil.LABEL_COLOR_CYCLE,
180                symbol="o",
181                face_color="state",
182                face_color_cycle=STATE_COLOR_CYCLE,
183                border_width=0.4,
184                size=12,
185                ndim=self._ndim,
186            )
187            self._point_prompt_layer.border_color_mode = "cycle"
188            self._point_prompt_layer.face_color_mode = "cycle"
189            _new_point_layer = True
190        else:
191            self._point_prompt_layer = self._viewer.layers["point_prompts"]
192            _new_point_layer = False
193
194        # Add the point prompts layer.
195        _box_prompt_property_choices = {"track_id": ["1"]}
196
197        box_layer_mismatch = True
198        if "prompts" in self._viewer.layers:
199            # Check whether the 'property_choices' match or not.
200            curr_property_choices = self._viewer.layers["prompts"].property_choices
201            box_layer_mismatch = set(curr_property_choices.keys()) != set(_box_prompt_property_choices.keys())
202
203        if box_layer_mismatch and "prompts" not in self._viewer.layers:
204            # Using the box layer to set divisions currently doesn't work.
205            # That's why some of the code below is commented out.
206            self._box_prompt_layer = self._viewer.add_shapes(
207                shape_type="rectangle",
208                edge_width=4,
209                ndim=self._ndim,
210                face_color="transparent",
211                name="prompts",
212                edge_color="green",
213                property_choices=_box_prompt_property_choices,
214                # property_choices={"track_id": ["1"], "state": self._track_state_labels},
215                # edge_color_cycle=STATE_COLOR_CYCLE,
216            )
217            # self._box_prompt_layer.edge_color_mode = "cycle"
218            _new_box_layer = True
219        else:
220            self._box_prompt_layer = self._viewer.layers["prompts"]
221            _new_box_layer = False
222
223        # Trigger a new connection for the tracking state menu only when a new layer is (re)created.
224        if _new_point_layer or _new_box_layer:
225            self._tracking_widget = create_tracking_menu(
226                points_layer=self._point_prompt_layer,
227                box_layer=self._box_prompt_layer,
228                states=self._track_state_labels,
229                track_ids=list(state.lineage.keys()),
230                tracking_widget=state.widgets.get("tracking"),
231            )
232            state.widgets["tracking"] = self._tracking_widget
233
234    def _get_widgets(self):
235        state = AnnotatorState()
236        self._require_layers()
237
238        # Create the tracking state menu.
239        # NOTE: Check whether it exists already from `_require_layers` or needs to be created.
240        if state.widgets.get("tracking") is None:
241            self._tracking_widget = create_tracking_menu(
242                points_layer=self._point_prompt_layer,
243                box_layer=self._box_prompt_layer,
244                states=self._track_state_labels,
245                track_ids=list(state.lineage.keys()),
246            )
247        else:
248            self._tracking_widget = state.widgets.get("tracking")
249
250        segment_nd = widgets.SegmentNDWidget(self._viewer, tracking=True)
251        autotrack = widgets.AutoTrackWidget(self._viewer, with_decoder=self._with_decoder, volumetric=True)
252        return {
253            "tracking": self._tracking_widget,
254            "segment": widgets.segment_frame(),
255            "segment_nd": segment_nd,
256            "autosegment": autotrack,
257            "commit": widgets.commit_track(),
258            "clear": widgets.clear_track(),
259        }
260
261    def __init__(self, viewer: "napari.viewer.Viewer", reset_state: bool = True) -> None:
262        # Initialize the state for tracking.
263        self._init_track_state()
264        self._with_decoder = AnnotatorState().decoder is not None
265        super().__init__(viewer=viewer, ndim=3)
266        # Go to t=0.
267        self._viewer.dims.current_step = (0, 0, 0) + tuple(sh // 2 for sh in self._shape[1:])
268
269        # Set the expected annotator class to the state.
270        state = AnnotatorState()
271
272        # Reset the state.
273        if reset_state:
274            state.reset_state()
275
276        state.annotator = self
277
278    def _init_track_state(self):
279        state = AnnotatorState()
280        state.current_track_id = 1
281        state.lineage = {1: []}
282        state.committed_lineages = []
283
284    def _update_image(self):
285        super()._update_image()
286        self._init_track_state()
287        state = AnnotatorState()
288        if self._with_decoder:
289            state.amg_state = vutil._load_is_state(state.embedding_path)
290        else:
291            state.amg_state = vutil._load_amg_state(state.embedding_path)
292
293
294def annotator_tracking(
295    image: np.ndarray,
296    embedding_path: Optional[str] = None,
297    # tracking_result: Optional[str] = None,
298    model_type: str = util._DEFAULT_MODEL,
299    tile_shape: Optional[Tuple[int, int]] = None,
300    halo: Optional[Tuple[int, int]] = None,
301    return_viewer: bool = False,
302    viewer: Optional["napari.viewer.Viewer"] = None,
303    precompute_amg_state: bool = False,
304    checkpoint_path: Optional[str] = None,
305    decoder_path: Optional[str] = None,
306    device: Optional[Union[str, torch.device]] = None,
307) -> Optional["napari.viewer.Viewer"]:
308    """Start the tracking annotation tool fora given timeseries.
309
310    Args:
311        image: The image data.
312        embedding_path: Filepath for saving the precomputed embeddings.
313        model_type: The Segment Anything model to use. For details on the available models check out
314            https://computational-cell-analytics.github.io/micro-sam/micro_sam.html#finetuned-models.
315        tile_shape: Shape of tiles for tiled embedding prediction.
316            If `None` then the whole image is passed to Segment Anything.
317        halo: Shape of the overlap between tiles, which is needed to segment objects on tile borders.
318        return_viewer: Whether to return the napari viewer to further modify it before starting the tool.
319            By default, does not return the napari viewer.
320        viewer: The viewer to which the Segment Anything functionality should be added.
321            This enables using a pre-initialized viewer.
322        precompute_amg_state: Whether to precompute the state for automatic mask generation.
323            This will take more time when precomputing embeddings, but will then make
324            automatic mask generation much faster. By default, set to 'False'.
325        checkpoint_path: Path to a custom checkpoint from which to load the SAM model.
326        decoder_path: Path to a custom decoder checkpoint from which to load the 'micro-sam` decoder.
327        device: The computational device to use for the SAM model.
328            By default, automatically chooses the best available device.
329
330    Returns:
331        The napari viewer, only returned if `return_viewer=True`.
332    """
333
334    # Initialize the predictor state.
335    state = AnnotatorState()
336    state.initialize_predictor(
337        image, model_type=model_type, save_path=embedding_path,
338        halo=halo, tile_shape=tile_shape, prefer_decoder=True,
339        ndim=3, checkpoint_path=checkpoint_path, device=device,
340        precompute_amg_state=precompute_amg_state, use_cli=True,
341        decoder_path=decoder_path,
342    )
343    state.image_shape = image.shape[:-1] if image.ndim == 4 else image.shape
344
345    if viewer is None:
346        viewer = napari.Viewer()
347
348    viewer.add_image(image, name="image")
349    annotator = AnnotatorTracking(viewer, reset_state=False)
350
351    # Trigger layer update of the annotator so that layers have the correct shape.
352    annotator._update_image()
353
354    # Add the annotator widget to the viewer and sync widgets.
355    viewer.window.add_dock_widget(annotator)
356    vutil._sync_embedding_widget(
357        widget=state.widgets["embeddings"],
358        model_type=model_type if checkpoint_path is None else state.predictor.model_type,
359        save_path=embedding_path,
360        checkpoint_path=checkpoint_path,
361        device=device,
362        tile_shape=tile_shape,
363        halo=halo,
364    )
365
366    if return_viewer:
367        return viewer
368
369    napari.run()
370
371
372def main():
373    """@private"""
374    parser = vutil._initialize_parser(
375        description="Run interactive segmentation for an image volume.",
376        with_segmentation_result=False,
377        with_instance_segmentation=False,
378    )
379
380    # Tracking result is not yet supported, we need to also deserialize the lineage.
381    # parser.add_argument(
382    #     "-t", "--tracking_result",
383    #     help="Optional filepath to a precomputed tracking result. If passed this will be used to initialize the "
384    #     "'committed_tracks' layer. This can be useful if you want to correct an existing tracking result or if you "
385    #     "have saved intermediate results from the annotator and want to continue. "
386    #     "Supports the same file formats as 'input'."
387    # )
388    # parser.add_argument(
389    #     "-tk", "--tracking_key",
390    #     help="The key for opening the tracking result. Same rules as for 'key' apply."
391    # )
392
393    args = parser.parse_args()
394    image = util.load_image_data(args.input, key=args.key)
395
396    annotator_tracking(
397        image, embedding_path=args.embedding_path, model_type=args.model_type,
398        tile_shape=args.tile_shape, halo=args.halo,
399        checkpoint_path=args.checkpoint, decoder_path=args.decoder_path, device=args.device,
400    )