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

@@ -0,0 +1,61 @@
+"""
+Functions from other third party libraries.
+
+We can remove these functions once the bug affecting our code is fixed upstream.
+
+The license type of the thrid party software project must be compatible with
+the software license the micro-sam project is distributed under.
+"""
+import torch
+
+
+# segment_anything/util/amg.py
+# https://github.com/facebookresearch/segment-anything
+def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates boxes in XYXY format around masks. Return [0,0,0,0] for
+    an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.
+    """
+    assert masks.dtype == torch.bool
+
+    # torch.max below raises an error on empty inputs, just skip in this case
+    if torch.numel(masks) == 0:
+        return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
+
+    # Normalize shape to CxHxW
+    shape = masks.shape
+    h, w = shape[-2:]
+    if len(shape) > 2:
+        masks = masks.flatten(0, -3)
+    else:
+        masks = masks.unsqueeze(0)
+
+    # Get top and bottom edges
+    in_height, _ = torch.max(masks, dim=-1)
+    in_height_coords = in_height * torch.arange(h, dtype=torch.int, device=in_height.device)[None, :]
+    bottom_edges, _ = torch.max(in_height_coords, dim=-1)
+    in_height_coords = in_height_coords + h * (~in_height)
+    in_height_coords = in_height_coords.type(torch.int)
+    top_edges, _ = torch.min(in_height_coords, dim=-1)
+
+    # Get left and right edges
+    in_width, _ = torch.max(masks, dim=-2)
+    in_width_coords = in_width * torch.arange(w, dtype=torch.int, device=in_width.device)[None, :]
+    right_edges, _ = torch.max(in_width_coords, dim=-1)
+    in_width_coords = in_width_coords + w * (~in_width)
+    in_width_coords = in_width_coords.type(torch.int)
+    left_edges, _ = torch.min(in_width_coords, dim=-1)
+
+    # If the mask is empty the right edge will be to the left of the left edge.
+    # Replace these boxes with [0, 0, 0, 0]
+    empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
+    out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
+    out = out * (~empty_filter).unsqueeze(-1)
+
+    # Return to original shape
+    if len(shape) > 2:
+        out = out.reshape(*shape[:-2], 4)
+    else:
+        out = out[0]
+
+    return out

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

@@ -32,7 +32,7 @@
 
 from . import util
 from .prompt_based_segmentation import segment_from_mask
-
+from ._vendored import batched_mask_to_box
 
 #
 # Utility Functionality
@@ -189,17 +189,17 @@ 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))
 
         # recalculate boxes and remove any new duplicates
         masks = torch.cat(new_masks, dim=0)
-        boxes = amg_utils.batched_mask_to_box(masks)
+        boxes = 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,7 +269,8 @@ 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["boxes"] = amg_utils.batched_mask_to_box(data["masks"])
+        data["masks"] = data["masks"].type(torch.bool)
+        data["boxes"] = batched_mask_to_box(data["masks"])
 
         # compress to RLE
         data["masks"] = amg_utils.uncrop_masks(data["masks"], crop_box, orig_h, orig_w)
@@ -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

test/test_vendored.py

@@ -0,0 +1,43 @@
+import unittest
+
+import numpy as np
+import torch
+
+
+class TestVendored(unittest.TestCase):
+    def setUp(self):
+        mask_numpy = np.zeros((10,10)).astype(bool)
+        mask_numpy[7:9, 3:5] = True
+        self.mask = mask_numpy
+        self.expected_result = [3, 7, 4, 8]
+
+    def test_cpu_batched_mask_to_box(self):
+        from micro_sam._vendored import batched_mask_to_box
+
+        device = "cpu"
+        mask = torch.as_tensor(self.mask, dtype=torch.bool, device=device)
+        expected_result = torch.as_tensor(self.expected_result, dtype=torch.int, device=device)
+        result = batched_mask_to_box(mask)
+        assert all(result == expected_result)
+
+    @unittest.skipIf(not torch.cuda.is_available(),
+                     "CUDA Pytorch backend is not available")
+    def test_cuda_batched_mask_to_box(self):
+        from micro_sam._vendored import batched_mask_to_box
+
+        device = "cuda"
+        mask = torch.as_tensor(self.mask, dtype=torch.bool, device=device)
+        expected_result = torch.as_tensor(self.expected_result, dtype=torch.int, device=device)
+        result = batched_mask_to_box(mask)
+        assert all(result == expected_result)
+
+    @unittest.skipIf(not (torch.backends.mps.is_available() and torch.backends.mps.is_built()),
+                     "MPS Pytorch backend is not available")
+    def test_mps_batched_mask_to_box(self):
+        from micro_sam._vendored import batched_mask_to_box
+
+        device = "mps"
+        mask = torch.as_tensor(self.mask, dtype=torch.bool, device=device)
+        expected_result = torch.as_tensor(self.expected_result, dtype=torch.int, device=device)
+        result = batched_mask_to_box(mask)
+        assert all(result == expected_result)