Embedding widget for napari plugin

micro_sam/napari.yaml

@@ -23,6 +23,9 @@ contributions:
     - id: micro-sam.sample_data_segmentation
       python_name: micro_sam.sample_data:sample_data_segmentation
       title: Load segmentation sample data from micro-sam plugin
+    - id: micro-sam.embedding_widget
+      python_name: micro_sam.sam_annotator._widgets:embedding_widget
+      title: Embedding widget
   sample_data:
     - command: micro-sam.sample_data_image_series
       display_name: Image series example data
@@ -45,3 +48,6 @@ contributions:
     - command: micro-sam.sample_data_segmentation
       display_name: Segmentation sample dataset
       key: micro-sam-segmentation
+  widgets:
+    - command: micro-sam.embedding_widget
+      display_name: Embedding widget

micro_sam/sam_annotator/_state.py

@@ -46,10 +46,10 @@
         init_sum = sum((have_image_embeddings, have_predictor, have_image_shape))
         if init_sum == 3:
             return True
         elif init_sum == 0:
             return False
         else:
             raise RuntimeError(
                 f"Invalid AnnotatorState: {init_sum} / 3 parts of the state "
                 "needed for interactive segmentation are initialized."
             )
@@ -60,10 +60,20 @@
         init_sum = sum((have_current_track_id, have_lineage))
         if init_sum == 2:
             return True
         elif init_sum == 0:
             return False
         else:
             raise RuntimeError(
                 f"Invalid AnnotatorState: {init_sum} / 2 parts of the state "
                 "needed for tracking are initialized."
             )
+
+    def reset_state(self):
+        """Reset state, clear all attributes."""
+        self.image_embeddings = None
+        self.predictor = None
+        self.image_shape = None
+        self.amg = None
+        self.amg_state = None
+        self.current_track_id = None
+        self.lineage = None

micro_sam/sam_annotator/_widgets.py

@@ -0,0 +1,84 @@
+from enum import Enum
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+from magicgui import magic_factory, widgets
+from napari.qt.threading import thread_worker
+import zarr
+from zarr.errors import PathNotFoundError
+
+from micro_sam.sam_annotator._state import AnnotatorState
+from micro_sam.util import (
+    ImageEmbeddings,
+    get_sam_model,
+    precompute_image_embeddings,
+    _MODEL_URLS,
+    _DEFAULT_MODEL,
+    _available_devices,
+)
+
+if TYPE_CHECKING:
+    import napari
+
+Model = Enum("Model", _MODEL_URLS)
+available_devices_list = ["auto"] + _available_devices()
+
+
+@magic_factory(
+    pbar={'visible': False, 'max': 0, 'value': 0, 'label': 'working...'},
+    call_button="Compute image embeddings",
+    device = {"choices": available_devices_list},
+    save_path={"mode": "d"},  # choose a directory
+    )
+def embedding_widget(
+    pbar: widgets.ProgressBar,
+    image: "napari.layers.Image",
+    model: Model = Model.__getitem__(_DEFAULT_MODEL),
+    device = "auto",
+    save_path: Optional[Path] = None,  # where embeddings for this image are cached (optional)
+    optional_custom_weights: Optional[Path] = None,  # A filepath or URL to custom model weights.
+) -> ImageEmbeddings:
+    """Image embedding widget."""
+    state = AnnotatorState()
+    state.reset_state()
+    # Get image dimensions
+    if image.rgb:
+        ndim = image.data.ndim - 1
+        state.image_shape = image.data.shape[:-1]
+    else:
+        ndim = image.data.ndim
+        state.image_shape = image.data.shape
+
+    @thread_worker(connect={'started': pbar.show, 'finished': pbar.hide})
+    def _compute_image_embedding(state, image_data, save_path, ndim=None,
+                                 device="auto", model=Model.__getitem__(_DEFAULT_MODEL),
+                                 optional_custom_weights=None):
+        # Make sure save directory exists and is an empty directory
+        if save_path is not None:
+            os.makedirs(save_path, exist_ok=True)
+            if not save_path.is_dir():
+                raise NotADirectoryError(
+                    f"The user selected 'save_path' is not a direcotry: {save_path}"
+                )
+            if len(os.listdir(save_path)) > 0:
+                try:
+                    zarr.open(save_path, "r")
+                except PathNotFoundError:
+                    raise RuntimeError(
+                        "The user selected 'save_path' is not a zarr array "
+                        f"or empty directory: {save_path}"
+                    )
+        # Initialize the model
+        state.predictor  = get_sam_model(device=device, model_type=model.name,
+                                         checkpoint_path=optional_custom_weights)
+        # Compute the image embeddings
+        state.image_embeddings = precompute_image_embeddings(
+            predictor = state.predictor,
+            input_ = image_data,
+            save_path = str(save_path),
+            ndim=ndim,
+        )
+        return state  # returns napari._qt.qthreading.FunctionWorker
+
+    return _compute_image_embedding(state, image.data, save_path, ndim=ndim, device=device, model=model, optional_custom_weights=optional_custom_weights)

micro_sam/util.py

@@ -127,10 +127,7 @@
     return checkpoint_path
 
 
-def _get_device(device):
-    if device is not None:
-        return device
-
+def _get_default_device():
     # Use cuda enabled gpu if it's available.
     if torch.cuda.is_available():
         device = "cuda"
@@ -145,6 +142,36 @@
     return device
 
 
+def _get_device(device=None):
+    if device is None or device == "auto":
+        device = _get_default_device()
+    else:
+        if device.lower() == "cuda":
+            if not torch.cuda.is_available():
+                raise RuntimeError("PyTorch CUDA backend is not available.")
+        elif device.lower() == "mps":
+            if not (torch.backends.mps.is_available() and torch.backends.mps.is_built()):
+                raise RuntimeError("PyTorch MPS backend is not available or is not built correctly.")
+        elif device.lower() == "cpu":
+            pass  # cpu is always available
+        else:
+            raise RuntimeError(f"Unsupported device: {device}\n"
+                               "Please choose from 'cpu', 'cuda', or 'mps'.")
+    return device
+
+
+def _available_devices():
+    available_devices = []
+    for i in ["cuda", "mps", "cpu"]:
+        try:
+            device = _get_device(i)
+        except RuntimeError:
+            pass
+        else:
+            available_devices.append(device)
+    return available_devices
+
+
 def get_sam_model(
     model_type: str = _DEFAULT_MODEL,
     device: Optional[str] = None,
@@ -269,7 +296,7 @@
         save_path: Where to save the exported model.
     """
     _, state = get_custom_sam_model(
-        checkpoint_path, model_type=model_type, return_state=True, device=torch.device("cpu"),
+        checkpoint_path, model_type=model_type, return_state=True, device="cpu",
     )
     model_state = state["model_state"]
     prefix = "sam."
@@ -552,7 +579,7 @@
                     continue
                 # check whether the key signature does not match or is not in the file
                 if key not in f.attrs or f.attrs[key] != val:
-                    warnings.warn(
+                    raise RuntimeError(
                         f"Embeddings file {save_path} is invalid due to unmatching {key}: "
                         f"{f.attrs.get(key)} != {val}.Please recompute embeddings in a new file."
                     )

micro_sam/visualization.py

@@ -3,8 +3,6 @@
 """
 from typing import Tuple
 
-from typing import Tuple
-
 import numpy as np
 
 from elf.segmentation.embeddings import embedding_pca

test/test_sam_annotator/test_widgets.py

@@ -0,0 +1,46 @@
+import json
+import os
+
+from mobile_sam.predictor import SamPredictor as MobileSamPredictor
+from segment_anything.predictor import SamPredictor
+import torch
+import zarr
+
+from micro_sam.sam_annotator._state import AnnotatorState
+from micro_sam.sam_annotator._widgets import embedding_widget, Model
+from micro_sam.util import _compute_data_signature
+
+
+# make_napari_viewer is a pytest fixture that returns a napari viewer object
+# you don't need to import it, as long as napari is installed
+# in your testing environment.
+# tmp_path is a regular pytest fixture.
+def test_embedding_widget(make_napari_viewer, tmp_path):
+    """Test embedding widget for micro-sam napari plugin."""
+    # setup
+    viewer = make_napari_viewer()
+    layer = viewer.open_sample('napari', 'camera')[0]
+    my_widget = embedding_widget()
+    # run image embedding widget
+    worker = my_widget(image=layer, model=Model.vit_t, device="cpu", save_path=tmp_path)
+    worker.await_workers()  # blocks until thread worker is finished the embedding
+    # Check in-memory state - predictor
+    assert isinstance(AnnotatorState().predictor, (SamPredictor, MobileSamPredictor))
+    # Check in-memory state - image embeddings
+    assert AnnotatorState().image_embeddings is not None
+    assert 'features' in AnnotatorState().image_embeddings.keys()
+    assert 'input_size' in AnnotatorState().image_embeddings.keys()
+    assert 'original_size' in AnnotatorState().image_embeddings.keys()
+    assert isinstance(AnnotatorState().image_embeddings["features"], torch.Tensor)
+    assert AnnotatorState().image_embeddings["original_size"] == layer.data.shape
+    # Check saved embedding results are what we expect to have
+    temp_path_files = os.listdir(tmp_path)
+    temp_path_files.sort()
+    assert temp_path_files == ['.zattrs', '.zgroup', 'features']
+    with open(os.path.join(tmp_path, ".zattrs")) as f:
+        content = f.read()
+    zarr_dict = json.loads(content)
+    assert zarr_dict.get("original_size") == list(layer.data.shape)
+    assert zarr_dict.get("data_signature") == _compute_data_signature(layer.data)
+    assert zarr.open(os.path.join(tmp_path, "features")).shape == (1, 256, 64, 64)
+    viewer.close()  # must close the viewer at the end of tests