Add support for MPS device

micro_sam/__init__.py

@@ -5,5 +5,8 @@
 .. include:: ../doc/python_library.md
 .. include:: ../doc/finetuned_models.md
 """
+import os
 
 from .__version__ import __version__
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

micro_sam/evaluation/inference.py

@@ -136,7 +136,7 @@ def _run_inference_with_prompts_for_image(
     prompts = deepcopy((input_point, input_label, input_box))
 
     # Transform the prompts into batches
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = util._get_device()
     input_point = torch.tensor(np.array(input_point)).to(device) if len(input_point) > 0 else None
     input_label = torch.tensor(np.array(input_label)).to(device) if len(input_label) > 0 else None
     input_box = torch.tensor(np.array(input_box)).to(device) if len(input_box) > 0 else None
@@ -524,7 +524,7 @@ def run_inference_with_iterative_prompting(
     """@private"""
     warnings.warn("The iterative prompting functionality is not working correctly yet.")
 
-    device = torch.device("cuda")
+    device = util._get_device()
     model = get_trainable_sam_model(model_type, checkpoint_path)
 
     # create all prediction folders

micro_sam/instance_segmentation.py

@@ -189,7 +189,7 @@ def _postprocess_small_regions(self, mask_data, min_area, nms_thresh):
             mask, changed = amg_utils.remove_small_regions(mask, min_area, mode="islands")
             unchanged = unchanged and not changed
 
-            new_masks.append(torch.as_tensor(mask).unsqueeze(0))
+            new_masks.append(torch.as_tensor(mask, dtype=torch.int).unsqueeze(0))
             # give score=0 to changed masks and score=1 to unchanged masks
             # so NMS will prefer ones that didn't need postprocessing
             scores.append(float(unchanged))
@@ -199,7 +199,7 @@ def _postprocess_small_regions(self, mask_data, min_area, nms_thresh):
         boxes = amg_utils.batched_mask_to_box(masks)
         keep_by_nms = batched_nms(
             boxes.float(),
-            torch.as_tensor(scores),
+            torch.as_tensor(scores, dtype=torch.float),
             torch.zeros_like(boxes[:, 0]),  # categories
             iou_threshold=nms_thresh,
         )
@@ -258,7 +258,7 @@ def _to_mask_data(self, masks, iou_preds, crop_box, original_size, points=None):
         # serialize predictions and store in MaskData
         data = amg_utils.MaskData(masks=masks.flatten(0, 1), iou_preds=iou_preds.flatten(0, 1))
         if points is not None:
-            data["points"] = torch.as_tensor(points.repeat(masks.shape[1], axis=0))
+            data["points"] = torch.as_tensor(points.repeat(masks.shape[1], axis=0), dtype=torch.float)
 
         del masks
 
@@ -269,6 +269,7 @@ def _to_mask_data(self, masks, iou_preds, crop_box, original_size, points=None):
 
         # threshold masks and calculate boxes
         data["masks"] = data["masks"] > self._predictor.model.mask_threshold
+        data["masks"] = data["masks"].type(torch.int)
         data["boxes"] = amg_utils.batched_mask_to_box(data["masks"])
 
         # compress to RLE
@@ -364,7 +365,7 @@ def __init__(
     def _process_batch(self, points, im_size, crop_box, original_size):
         # run model on this batch
         transformed_points = self._predictor.transform.apply_coords(points, im_size)
-        in_points = torch.as_tensor(transformed_points, device=self._predictor.device)
+        in_points = torch.as_tensor(transformed_points, device=self._predictor.device, dtype=torch.float)
         in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device)
         masks, iou_preds, _ = self._predictor.predict_torch(
             in_points[:, None, :],

micro_sam/training/util.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from ..prompt_generators import PointAndBoxPromptGenerator
-from ..util import get_centers_and_bounding_boxes, get_sam_model
+from ..util import get_centers_and_bounding_boxes, get_sam_model, _get_device
 from .trainable_sam import TrainableSAM
 
 
@@ -28,8 +28,7 @@ def get_trainable_sam_model(
         The trainable segment anything model.
     """
     # set the device here so that the correct one is passed to TrainableSAM below
-    if device is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = _get_device(device)
     _, sam = get_sam_model(device, model_type, checkpoint_path, return_sam=True)
 
     # freeze components of the model if freeze was passed

micro_sam/util.py

@@ -116,6 +116,24 @@ def _get_checkpoint(model_type, checkpoint_path=None):
     return checkpoint_path
 
 
+def _get_device(device):
+    if device is not None:
+        return device
+
+    # Use cuda enabled gpu if it's available.
+    if torch.cuda.is_available():
+        device = "cuda"
+    # As second priority use mps.
+    # See https://pytorch.org/docs/stable/notes/mps.html for details
+    elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
+        print("Using apple MPS device.")
+        device = "mps"
+    # Use the CPU as fallback.
+    else:
+        device = "cpu"
+    return device
+
+
 def get_sam_model(
     device: Optional[str] = None,
     model_type: str = _DEFAULT_MODEL,
@@ -138,8 +156,7 @@ def get_sam_model(
         The segment anything predictor.
     """
     checkpoint = _get_checkpoint(model_type, checkpoint_path)
-    if device is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = _get_device(device)
 
     # Our custom model types have a suffix "_...". This suffix needs to be stripped
     # before calling sam_model_registry.
@@ -196,8 +213,7 @@ def get_custom_sam_model(
     custom_pickle = pickle
     custom_pickle.Unpickler = _CustomUnpickler
 
-    if device is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = _get_device()
     sam = sam_model_registry[model_type]()
 
     # load the model state, ignoring any attributes that can't be found by pickle