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/_widgets.py

@@ -0,0 +1,76 @@
+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
+
+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."""
+    # Make sure save directory exists and is an empty directory
+    if save_path is not None:
+        if not save_path.exists():
+            os.makedirs(save_path)
+        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:
+            raise RuntimeError(
+                f"The user selected 'save_path' is not empty: {save_path}"
+            )
+
+    state = AnnotatorState()
+    # Initialize the model
+    state.predictor  = get_sam_model(device=device, model_type=model.name,
+                              checkpoint_path=optional_custom_weights)
+    # Get image dimensions
+    ndim = image.data.ndim
+    if image.rgb:
+        ndim -= 1
+
+    # Compute the image embeddings
+    @thread_worker(connect={'started': pbar.show, 'returned': pbar.hide})
+    def _compute_image_embedding(state, image_data, save_path, ndim=None):
+        if save_path is not None:
+            save_path = str(save_path)
+        state.image_embeddings = precompute_image_embeddings(
+            predictor = state.predictor,
+            input_ = image_data,
+            save_path = save_path,
+            ndim=ndim,
+        )
+        return state.image_embeddings  # returns napari._qt.qthreading.FunctionWorker
+
+    return _compute_image_embedding(state, image.data, save_path, ndim=ndim)

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,
@@ -258,6 +285,7 @@
     checkpoint_path: Union[str, os.PathLike],
     model_type: str,
     save_path: Union[str, os.PathLike],
+    device: str = None,
 ) -> None:
     """Export a finetuned segment anything model to the standard model format.
 
@@ -269,7 +297,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=device,
     )
     model_state = state["model_state"]
     prefix = "sam."

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

test/test_training.py

@@ -110,13 +110,14 @@ def _train_model(self, model_type, device):
         )
         trainer.fit(epochs=1)
 
-    def _export_model(self, checkpoint_path, export_path, model_type):
+    def _export_model(self, checkpoint_path, export_path, model_type, device):
         from micro_sam.util import export_custom_sam_model
 
         export_custom_sam_model(
             checkpoint_path=checkpoint_path,
             model_type=model_type,
             save_path=export_path,
+            device=device,
         )
 
     def _run_inference_and_check_results(
@@ -152,7 +153,7 @@ def test_training(self):
 
         # Export the model.
         export_path = os.path.join(self.tmp_folder, "exported_model.pth")
-        self._export_model(checkpoint_path, export_path, model_type)
+        self._export_model(checkpoint_path, export_path, model_type, device)
         self.assertTrue(os.path.exists(export_path))
 
         # Check the model with inference with a single point prompt.