Add masks to boundaries

docs/source/ops.rst

@@ -22,6 +22,7 @@ The below operators perform pre-processing as well as post-processing required i
 
     batched_nms
     masks_to_boxes
+    masks_to_boundaries
     nms
     roi_align
     roi_pool

test/test_ops.py

@@ -11,7 +11,7 @@
 import torch.fx
 import torch.nn.functional as F
 from common_utils import assert_equal, cpu_and_cuda, needs_cuda
-from PIL import Image
+from PIL import Image, ImageDraw
 from torch import nn, Tensor
 from torch.autograd import gradcheck
 from torch.nn.modules.utils import _pair
@@ -621,6 +621,32 @@ def test_is_leaf_node(self, device):
         assert len(graph_node_names[0]) == len(graph_node_names[1])
         assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
 
+class TestMasksToBoundaries(ABC):
+
+    @pytest.mark.parametrize("device", ['cpu', 'cuda'])  
+    def test_masks_to_boundaries(self, device):
+        # Create masks
+        mask = torch.zeros(4, 32, 32, dtype=torch.bool)
+        mask[0, 1:10, 1:10] = True
+        mask[0, 12:20, 12:20] = True
+        mask[0, 15:18, 20:32] = True
+        mask[1, 15:23, 15:23] = True
+        mask[1, 22:33, 22:33] = True
+        mask[2, 1:5, 22:30] = True
+        mask[2, 5:14, 25:27] = True
+        pil_img = Image.new("L", (32, 32))
+        draw = ImageDraw.Draw(pil_img)
+        draw.ellipse([2, 7, 26, 26], fill=1, outline=1, width=1)
+        mask[3, ...] = torch.from_numpy(np.asarray(pil_img))
+        mask = mask.to(device)
+        dilation_ratio = 0.02
+        boundaries = ops.masks_to_boundaries(mask, dilation_ratio)
+        # Generate expected output
+        # TODO: How we generate handle the expected output?
+        # replace with actual code to generate expected output
+        expected_boundaries = torch.zeros_like(mask)  
+        torch.testing.assert_close(expected_boundaries, boundaries)
+
 
 class TestNMS:
     def _reference_nms(self, boxes, scores, iou_threshold):

torchvision/ops/__init__.py

@@ -9,6 +9,7 @@
     distance_box_iou,
     generalized_box_iou,
     masks_to_boxes,
+    masks_to_boundaries,
     nms,
     remove_small_boxes,
 )
@@ -32,6 +33,7 @@
 
 __all__ = [
     "masks_to_boxes",
+    "masks_to_boundaries",
     "deform_conv2d",
     "DeformConv2d",
     "nms",

torchvision/ops/boxes.py

@@ -382,7 +382,40 @@ def _box_diou_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tuple[Te
     # distance between boxes' centers squared.
     return iou - (centers_distance_squared / diagonal_distance_squared), iou
 
+def masks_to_boundaries(masks: torch.Tensor, dilation_ratio: float = 0.02) -> torch.Tensor:
+    """
+    Compute the boundaries around the provided masks using morphological operations.
+
+    Returns a tensor of the same shape as the input masks containing the boundaries of each mask.
+
+    Args:
+    masks (Tensor[N, H, W]): masks to transform where N is the number of masks
+    and (H, W) are the spatial dimensions.
+    dilation_ratio (float, optional): ratio used for the dilation operation. Default: 0.02
+
+    Returns:
+    Tensor[N, H, W]: boundaries
+    """
+    # If no masks are provided, return an empty tensor
+    if masks.numel() == 0:
+        return torch.zeros_like(masks)
+
+    n, h, w = masks.shape
+    img_diag = math.sqrt(h ** 2 + w ** 2)
+    dilation = int(round(dilation_ratio * img_diag))
+    selem_size = dilation * 2 + 1
+    selem = torch.ones((n, 1, selem_size, selem_size), device=masks.device)
 
+    # Compute the boundaries for each mask
+    masks = masks.float().unsqueeze(1)
+    eroded_masks = F.conv2d(masks, selem, padding=dilation)
+    # Make the output binary
+    eroded_masks = (eroded_masks == selem.view(n, -1).sum(-1).view(n, 1, 1, 1)).byte()
+
+    contours = masks.byte() - eroded_masks
+
+    return contours.squeeze(1)
+
 def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor:
     """
     Compute the bounding boxes around the provided masks.
@@ -415,3 +448,4 @@ def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor:
         bounding_boxes[index, 3] = torch.max(y)
 
     return bounding_boxes
+