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,58 @@
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+from magicgui import magic_factory
+from magicgui.tqdm import tqdm
+from superqt.utils import thread_worker
+
+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(call_button="Compute image embeddings",
+               device = {"choices": available_devices_list})
+def embedding_widget(
+    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)
+    custom_model: Optional[str] = None,  # A filepath or URL to custom model weights.
+) -> ImageEmbeddings:
+    """Image embedding widget."""
+    # for access to the predictor and the image embeddings in the widgets
+    global PREDICTOR, IMAGE_EMBEDDINGS
+    # Initialize the model
+    PREDICTOR = get_sam_model(model_type=model.name)
+
+    # Get image dimensions
+    if not image.rgb:
+        ndim = image.data.ndim
+    else:
+        # assumes RGB channels are the last dimension
+        ndim = image.data.ndim
+
+    # Compute the embeddings for the image data
+    with tqdm() as pbar:
+        @thread_worker(connect={"finished": lambda: pbar.progressbar.hide()})
+        def _compute_image_embedding(PREDICTOR, image, save_path, ndim=None):
+            IMAGE_EMBEDDINGS = precompute_image_embeddings(
+                predictor = PREDICTOR,
+                input_ = image,
+                save_path = str(save_path),
+                ndim=ndim,
+            )
+
+        _compute_image_embedding(PREDICTOR, image.data, save_path, ndim=ndim)

micro_sam/util.py

@@ -127,10 +127,7 @@ def _get_checkpoint(model_type, checkpoint_path=None):
     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 @@ def _get_device(device):
     return device
 
 
+def _get_device(device=None):
+    if device is None or device == "auto":
+        device = _get_default_device()
+    else:
+        if device == "cuda":
+            if not torch.cuda.is_available():
+                raise RuntimeError("PyTorch CUDA backend is not available.")
+        elif device == "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 == "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(
     device: Optional[str] = None,
     model_type: str = _DEFAULT_MODEL,

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