micro_sam.sam_annotator.object_classifier

  1import os
  2from joblib import dump
  3from multiprocessing import cpu_count
  4from pathlib import Path
  5from typing import List, Optional, Tuple, Union
  6
  7import imageio.v3 as imageio
  8import napari
  9import numpy as np
 10import torch
 11
 12from magicgui import magic_factory, magicgui
 13from magicgui.widgets import Widget, Container, FunctionGui, create_widget
 14from qtpy import QtWidgets
 15
 16from skimage.measure import regionprops_table
 17from sklearn.ensemble import RandomForestClassifier
 18
 19from .. import util
 20from ..object_classification import compute_object_features, project_prediction_to_segmentation
 21from ._state import AnnotatorState
 22from . import _widgets as widgets
 23from .util import _sync_embedding_widget
 24
 25#
 26# Utility functionality.
 27# Some of this could be refactored to general purpose functionality that can also
 28# be used for inference with the trained classifier.
 29#
 30
 31
 32def _accumulate_labels(segmentation, annotations):
 33
 34    def majority_label(mask, annotation):
 35        ids, counts = np.unique(annotation[mask], return_counts=True)
 36        if len(ids) == 1 and ids[0] == 0:
 37            return 0
 38        if ids[0] == 0:
 39            ids, counts = ids[1:], counts[1:]
 40        return ids[np.argmax(counts)]
 41
 42    all_features = regionprops_table(
 43        segmentation, intensity_image=annotations, properties=("label",),
 44        extra_properties=[majority_label],
 45    )
 46    return all_features["majority_label"].astype("int")
 47
 48
 49def _train_rf(features, labels, previous_features=None, previous_labels=None, **rf_kwargs):
 50    assert len(features) == len(labels)
 51    valid = labels != 0
 52    X, y = features[valid], labels[valid]
 53
 54    if previous_features is not None:
 55        assert previous_labels is not None and len(previous_features) == len(previous_labels)
 56        X = np.concatenate([previous_features, X], axis=0)
 57        y = np.concatenate([previous_labels, y], axis=0)
 58
 59    rf = RandomForestClassifier(**rf_kwargs)
 60    rf.fit(X, y)
 61    return rf
 62
 63
 64# TODO do we add a shortcut?
 65@magic_factory(call_button="Train and predict")
 66def _train_and_predict_rf_widget(viewer: "napari.viewer.Viewer") -> None:
 67    # Get the object features and the annotations.
 68    state = AnnotatorState()
 69    state.annotator._require_layers()
 70    annotations = viewer.layers["annotations"].data
 71    segmentation = state.segmentation_selection.get_value().data
 72
 73    if state.object_features is None:
 74        if widgets._validate_embeddings(viewer):
 75            return None
 76        image_embeddings = state.image_embeddings
 77        seg_ids, features = compute_object_features(image_embeddings, segmentation)
 78        state.seg_ids = seg_ids
 79        state.object_features = features
 80    else:
 81        features, seg_ids = state.object_features, state.seg_ids
 82
 83    previous_features, previous_labels = state.previous_features, state.previous_labels
 84    labels = _accumulate_labels(segmentation, annotations)
 85    if (labels == 0).all() and (previous_labels is None):
 86        return widgets._generate_message("error", "You have not provided any annotations.")
 87
 88    # Run RF training and store it in the state.
 89    rf = _train_rf(
 90        features, labels, previous_features=previous_features, previous_labels=previous_labels,
 91        n_estimators=200, max_depth=10, n_jobs=cpu_count(),
 92    )
 93    state.object_rf = rf
 94
 95    # Run and set the prediction.
 96    pred = rf.predict(features)
 97    prediction_data = project_prediction_to_segmentation(segmentation, pred, seg_ids)
 98    viewer.layers["prediction"].data = prediction_data
 99
100    state.annotator._refresh_label_widget()
101
102
103@magic_factory(call_button="Export Classifier")
104def _create_export_rf_widget(export_path: Optional[Path] = None) -> None:
105    state = AnnotatorState()
106    rf = state.object_rf
107    if rf is None:
108        return widgets._generate_message("error", "You have not run training yet.")
109    if export_path is None or export_path == "":
110        return widgets._generate_message("error", "You have to provide an export path.")
111    # Do we add an extension? .joblib?
112    dump(rf, export_path)
113    # TODO show an info method about the export
114
115#
116# Object classifier implementation.
117#
118
119
120# TODO add a gui element that shows the current label ids, how many objects are labeled, and that
121# enables naming them so that the user can keep track of what has been labeled
122class ObjectClassifier(QtWidgets.QScrollArea):
123
124    def _require_layers(self, layer_choices: Optional[List[str]] = None):
125        # Check whether the image is initialized already. And use the image shape and scale for the layers.
126        state = AnnotatorState()
127        shape = self._shape if state.image_shape is None else state.image_shape
128
129        # Add the label layers for the current object, the automatic segmentation and the committed segmentation.
130        dummy_data = np.zeros(shape, dtype="uint32")
131        image_scale = state.image_scale
132
133        # Before adding new layers, we always check whether a layer with this name already exists or not.
134        if "annotations" not in self._viewer.layers:
135            if layer_choices and "annotations" in layer_choices:
136                widgets._validation_window_for_missing_layer("annotations")
137            annotation_layer = self._viewer.add_labels(data=dummy_data, name="annotations")
138            if image_scale is not None:
139                self._viewer.layers["annotations"].scale = image_scale
140            # Reduce the brush size and set the default mode to "paint" brush mode.
141            annotation_layer.brush_size = 3
142            annotation_layer.mode = "paint"
143
144        if "prediction" not in self._viewer.layers:
145            if layer_choices and "prediction" in layer_choices:
146                widgets._validation_window_for_missing_layer("prediction")
147            self._viewer.add_labels(data=dummy_data, name="prediction")
148            if image_scale is not None:
149                self._viewer.layers["prediction"].scale = image_scale
150
151    def _create_segmentation_layer_section(self):
152        segmentation_selection = QtWidgets.QVBoxLayout()
153        segmentation_layer_widget = QtWidgets.QLabel("Segmentation:")
154        segmentation_selection.addWidget(segmentation_layer_widget)
155        self.segmentation_selection = create_widget(annotation=napari.layers.Labels)
156        state = AnnotatorState()
157        state.segmentation_selection = self.segmentation_selection
158        segmentation_selection.addWidget(self.segmentation_selection.native)
159        return segmentation_selection
160
161    def _create_label_widget(self):
162        self._label_form = QtWidgets.QFormLayout()
163        scroll_area = QtWidgets.QScrollArea()
164        inner = QtWidgets.QWidget()
165        inner.setLayout(self._label_form)
166        scroll_area.setWidget(inner)
167        scroll_area.setWidgetResizable(True)
168
169        layout = QtWidgets.QVBoxLayout()
170        layout.addWidget(QtWidgets.QLabel("Object label names:"))
171        layout.addWidget(scroll_area)
172
173        return layout
174
175    def _refresh_label_widget(self):
176        state = AnnotatorState()
177
178        # Get the current label ids.
179        ids = np.unique(self._viewer.layers["annotations"].data)[1:]
180        if state.previous_labels is not None:
181            ids = np.union1d(ids, np.unique(state.previous_labels))
182
183        # Add new rows.
184        for lbl in ids:
185            if lbl in self._label_names:
186                continue
187            line = QtWidgets.QLineEdit(self._label_names.get(lbl, ""))
188            self._label_names[lbl] = ""
189            self._label_form.addRow(f"ID {lbl}", line)
190            line.textChanged.connect(lambda txt, lbl=lbl: self._label_names.__setitem__(lbl, txt))
191
192        # Remove rows whose label vanished.
193        for row in reversed(range(self._label_form.rowCount())):
194            lbl_text = self._label_form.itemAt(row, QtWidgets.QFormLayout.LabelRole).widget().text()
195            lbl_id = int(lbl_text.split()[1])
196            if lbl_id not in ids:
197                # Remove label+field widgets completely.
198                w_label = self._label_form.itemAt(row, QtWidgets.QFormLayout.LabelRole).widget()
199                w_edit = self._label_form.itemAt(row, QtWidgets.QFormLayout.FieldRole).widget()
200                self._label_form.removeRow(row)
201                w_label.deleteLater()
202                w_edit.deleteLater()
203                self.names.pop(lbl_id, None)
204
205    def _create_widgets(self):
206        # Create the embedding widget and connect all events related to it.
207        self._embedding_widget = widgets.EmbeddingWidget()
208        # Connect events for the image selection box.
209        self._viewer.layers.events.inserted.connect(self._embedding_widget.image_selection.reset_choices)
210        self._viewer.layers.events.removed.connect(self._embedding_widget.image_selection.reset_choices)
211        # Connect the run button with the function to update the image.
212        self._embedding_widget.run_button.clicked.connect(self._update_image)
213
214        # Create the widget for training and prediction of the classifier.
215        self._train_and_predict_widget = _train_and_predict_rf_widget()
216
217        # Create the widget for segmentation selection.
218        self._seg_selection_widget = self._create_segmentation_layer_section()
219
220        # Create the widget for displaying the current label state.
221        self._label_widget = self._create_label_widget()
222
223        # Cretate the widget for exporting the RF.
224        self._export_rf_widget = _create_export_rf_widget()
225
226        self._widgets = {
227            "embeddings": self._embedding_widget,
228            "segmentation_selection": self._seg_selection_widget,
229            "train_and_predict": self._train_and_predict_widget,
230            "label_widget": self._label_widget,
231            "export_rf": self._export_rf_widget,
232        }
233
234    def __init__(self, viewer: "napari.viewer.Viewer") -> None:
235        """Create the GUI for the object classifier.
236
237        Args:
238            viewer: The napari viewer.
239        """
240        super().__init__()
241        self._viewer = viewer
242        self._annotator_widget = QtWidgets.QWidget()
243        self._annotator_widget.setLayout(QtWidgets.QVBoxLayout())
244
245        # Add the layers for prompts and segmented obejcts.
246        # Initialize with a dummy shape, which is reset to the correct shape once an image is set.
247        self._shape = (256, 256)
248        self._require_layers()
249        self._ndim = len(self._shape)
250
251        # Create all the widgets and add them to the layout.
252        self._label_names = {}  # The names for the object labels.
253        self._create_widgets()
254
255        # We could refactor this.
256        for widget_name, widget in self._widgets.items():
257            widget_frame = QtWidgets.QGroupBox()
258            widget_layout = QtWidgets.QVBoxLayout()
259            if isinstance(widget, (Container, FunctionGui, Widget)):
260                # This is a magicgui type and we need to get the native qt widget.
261                widget_layout.addWidget(widget.native)
262            elif isinstance(widget, QtWidgets.QLayout):
263                widget_layout.addLayout(widget)
264            else:
265                # This is a qt type and we add the widget directly.
266                widget_layout.addWidget(widget)
267            widget_frame.setLayout(widget_layout)
268            self._annotator_widget.layout().addWidget(widget_frame)
269
270        # Connect the label layer and the refresh function.
271        self._refresh_label_widget()
272
273        # Set the expected annotator class to the state.
274        state = AnnotatorState()
275        state.annotator = self
276
277        # Add the widgets to the state.
278        state.widgets = self._widgets
279
280        # Add the widget to the scroll area.
281        self.setWidgetResizable(True)  # Allow widget to resize within scroll area.
282        self.setWidget(self._annotator_widget)
283
284    # napari 0.9 caps the height of dock widget content at its size hint (napari/napari#9393).
285    def showEvent(self, event):
286        size_policy = self.sizePolicy()
287        size_policy.setVerticalPolicy(QtWidgets.QSizePolicy.Policy.Expanding)
288        self.setSizePolicy(size_policy)
289        super().showEvent(event)
290
291    def _update_image(self, segmentation_result=None):
292        state = AnnotatorState()
293
294        # Whether embeddings already exist and avoid clearing objects in layers.
295        if state.skip_recomputing_embeddings:
296            return
297
298        if state.image_shape is None:
299            return
300
301        # Update the dimension and image shape if it has changed.
302        if state.image_shape != self._shape:
303            self._ndim = len(state.image_shape)
304            self._shape = state.image_shape
305
306        # Before we reset the layers, we ensure all expected layers exist.
307        self._require_layers()
308
309        # Update the image scale.
310        scale = state.image_scale
311
312        # Reset all layers.
313        self._viewer.layers["annotations"].data = np.zeros(self._shape, dtype="uint32")
314        self._viewer.layers["annotations"].scale = scale
315        self._viewer.layers["prediction"].data = np.zeros(self._shape, dtype="uint32")
316        self._viewer.layers["prediction"].scale = scale
317
318
319def object_classifier(
320    image: np.ndarray,
321    segmentation: np.ndarray,
322    embedding_path: Optional[Union[str, util.ImageEmbeddings]] = None,
323    model_type: str = util._DEFAULT_MODEL,
324    tile_shape: Optional[Tuple[int, int]] = None,
325    halo: Optional[Tuple[int, int]] = None,
326    return_viewer: bool = False,
327    viewer: Optional["napari.viewer.Viewer"] = None,
328    checkpoint_path: Optional[str] = None,
329    device: Optional[Union[str, torch.device]] = None,
330    ndim: Optional[int] = None,
331) -> Optional["napari.viewer.Viewer"]:
332    """Start the object classifier for a given image and segmentation.
333
334    Args:
335        image: The image data.
336        segmentation: The segmentation data.
337        embedding_path: Filepath where to save the embeddings
338            or the precompted image embeddings computed by `precompute_image_embeddings`.
339        model_type: The Segment Anything model to use. For details on the available models check out
340            https://computational-cell-analytics.github.io/micro-sam/micro_sam.html#finetuned-models.
341        tile_shape: Shape of tiles for tiled embedding prediction.
342            If `None` then the whole image is passed to Segment Anything.
343        halo: Shape of the overlap between tiles, which is needed to segment objects on tile borders.
344        return_viewer: Whether to return the napari viewer to further modify it before starting the tool.
345            By default, does not return the napari viewer.
346        viewer: The viewer to which the Segment Anything functionality should be added.
347            This enables using a pre-initialized viewer.
348        checkpoint_path: Path to a custom checkpoint from which to load the SAM model.
349        device: The computational device to use for the SAM model.
350            By default, automatically chooses the best available device.
351        ndim: The dimensionality of the data. If not given will be derived from the data.
352
353    Returns:
354        The napari viewer, only returned if `return_viewer=True`.
355    """
356    if ndim is None:
357        ndim = image.ndim - 1 if image.shape[-1] == 3 and image.ndim in (3, 4) else image.ndim
358
359    state = AnnotatorState()
360    state.image_shape = image.shape[:ndim]
361
362    state.initialize_predictor(
363        image, model_type=model_type, save_path=embedding_path,
364        halo=halo, tile_shape=tile_shape, precompute_amg_state=False,
365        ndim=ndim, checkpoint_path=checkpoint_path, device=device,
366        skip_load=False, use_cli=True,
367    )
368
369    if viewer is None:
370        viewer = napari.Viewer()
371
372    viewer.add_image(image, name="image")
373    viewer.add_labels(segmentation, name="segmentation")
374
375    annotator = ObjectClassifier(viewer)
376
377    # Trigger layer update of the annotator so that layers have the correct shape.
378    # And initialize the 'committed_objects' with the segmentation result if it was given.
379    annotator._update_image()
380
381    # Add the annotator widget to the viewer and sync widgets.
382    viewer.window.add_dock_widget(annotator)
383    _sync_embedding_widget(
384        widget=state.widgets["embeddings"],
385        model_type=model_type if checkpoint_path is None else state.predictor.model_type,
386        save_path=embedding_path,
387        checkpoint_path=checkpoint_path,
388        device=device,
389        tile_shape=tile_shape,
390        halo=halo,
391    )
392
393    if return_viewer:
394        return viewer
395
396    napari.run()
397
398
399def image_series_object_classifier(
400    images: List[np.ndarray],
401    segmentations: List[np.ndarray],
402    output_folder: str,
403    embedding_paths: Optional[List[Union[str, util.ImageEmbeddings]]] = None,
404    model_type: str = util._DEFAULT_MODEL,
405    tile_shape: Optional[Tuple[int, int]] = None,
406    halo: Optional[Tuple[int, int]] = None,
407    checkpoint_path: Optional[str] = None,
408    device: Optional[Union[str, torch.device]] = None,
409    ndim: Optional[int] = None,
410) -> None:
411    """Start the object classifier for a list of images and segmentations.
412
413    This function will save the all features and labels for annotated objects,
414    to enable training a random forest on multiple images.
415
416    Args:
417        images: The input images.
418        segmentations: The input segmentations.
419        output_folder: The folder where segmentation results, trained random forest
420            and the features, labels aggregated during training will be saved.
421        embedding_paths: Filepaths where to save the embeddings
422            or the precompted image embeddings computed by `precompute_image_embeddings`.
423        model_type: The Segment Anything model to use. For details on the available models check out
424            https://computational-cell-analytics.github.io/micro-sam/micro_sam.html#finetuned-models.
425        tile_shape: Shape of tiles for tiled embedding prediction.
426            If `None` then the whole image is passed to Segment Anything.
427        halo: Shape of the overlap between tiles, which is needed to segment objects on tile borders.
428        checkpoint_path: Path to a custom checkpoint from which to load the SAM model.
429        device: The computational device to use for the SAM model.
430            By default, automatically chooses the best available device.
431        ndim: The dimensionality of the data. If not given will be derived from the data.
432    """
433    # TODO precompute the embeddings if not computed, can re-use 'precompute' from image series annotator.
434    # TODO support file paths as inputs
435    # TODO option to skip segmented
436    if len(images) != len(segmentations):
437        raise ValueError(
438            f"Expect the same number of images and segmentations, got {len(images)}, {len(segmentations)}."
439        )
440
441    end_msg = "You have annotated the last image. Do you wish to close napari?"
442
443    # Initialize the object classifier on the fist image / segmentation.
444    viewer = object_classifier(
445        image=images[0], segmentation=segmentations[0],
446        embedding_path=None if embedding_paths is None else embedding_paths[0],
447        model_type=model_type, tile_shape=tile_shape, halo=halo,
448        return_viewer=True, checkpoint_path=checkpoint_path,
449        device=device, ndim=ndim,
450    )
451
452    os.makedirs(output_folder, exist_ok=True)
453    next_image_id = 0
454
455    def _save_prediction(image, pred, image_id):
456        fname = f"{Path(image).stem}_prediction.tif" if isinstance(image, str) else f"prediction_{image_id}.tif"
457        save_path = os.path.join(output_folder, fname)
458        imageio.imwrite(save_path, pred, compression="zlib")
459
460    # TODO handle cases where rf for the image was not trained, raise a message, enable contnuing
461    # Add functionality for going to the next image.
462    @magicgui(call_button="Next Image [N]")
463    def next_image(*args):
464        nonlocal next_image_id
465
466        # Get the state and the current segmentation (note that next image id has not yet been increased)
467        state = AnnotatorState()
468        segmentation = segmentations[next_image_id]
469
470        # Keep track of the previous features and labels.
471        labels = _accumulate_labels(segmentation, viewer.layers["annotations"].data)
472        valid = labels != 0
473        if valid.sum() > 0:
474            features, labels = state.object_features[valid], labels[valid]
475            if state.previous_features is None:
476                state.previous_features, state.previous_labels = features, labels
477            else:
478                state.previous_features = np.concatenate([state.previous_features, features], axis=0)
479                state.previous_labels = np.concatenate([state.previous_labels, labels], axis=0)
480            # Save the accumulated features and labels.
481            np.save(os.path.join(output_folder, "features.npy"), state.previous_features)
482            np.save(os.path.join(output_folder, "labels.npy"), state.previous_labels)
483
484        # Save the current prediction and RF.
485        _save_prediction(images[next_image_id], viewer.layers["prediction"].data, next_image_id)
486        dump(state.object_rf, os.path.join(output_folder, "rf.joblib"))
487
488        # Go to the next image.
489        next_image_id += 1
490
491        # Check if we are done.
492        if next_image_id == len(images):
493            # Inform the user via dialog.
494            abort = widgets._generate_message("info", end_msg)
495            if not abort:
496                viewer.close()
497            return
498
499        # Get the next image, segmentation and embedding_path.
500        image = images[next_image_id]
501        segmentation = segmentations[next_image_id]
502        embedding_path = None if embedding_paths is None else embedding_paths[next_image_id]
503
504        # Set the new image in the viewer, state and annotator.
505        viewer.layers["image"].data = image
506        viewer.layers["segmentation"].data = segmentation
507
508        state.initialize_predictor(
509            image, model_type=model_type, ndim=ndim,
510            save_path=embedding_path,
511            tile_shape=tile_shape, halo=halo,
512            predictor=state.predictor, device=device,
513        )
514        state.image_shape = image.shape if image.ndim == ndim else image.shape[:-1]
515        state.annotator._update_image()
516
517        # Clear the object features and seg-ids from the state.
518        state.object_features = None
519        state.seg_ids = None
520
521    viewer.window.add_dock_widget(next_image)
522
523    @viewer.bind_key("n", overwrite=True)
524    def _next_image(viewer):
525        next_image(viewer)
526
527    napari.run()
528
529
530# TODO: folder annotator
531# TODO: main function