Add MPS Pytorch support for batched_mask_to_box function

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/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
@@ -196,7 +196,7 @@ def _postprocess_small_regions(self, mask_data, min_area, nms_thresh):
 
         # 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, dtype=torch.float),
@@ -269,8 +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["masks"] = data["masks"].type(torch.int)
-        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)

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)