Added utility to draw segmentation masks

docs/source/utils.rst

@@ -7,4 +7,6 @@ torchvision.utils
 
 .. autofunction:: save_image
 
-.. autofunction:: draw_bounding_boxes
+.. autofunction:: draw_bounding_boxes
+
+.. autofunction:: draw_segmentation_masks

test/test_utils.py

@@ -96,6 +96,39 @@ def test_draw_boxes(self):
         expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
         self.assertTrue(torch.equal(result, expected))
 
+    def test_draw_segmentation_masks_colors(self):
+        img = torch.full((3, 20, 20), 255, dtype=torch.uint8)
+        colors = ["green", "#FF00FF", (0, 255, 0), "red"]
+        # TODO
+        masks = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+        result = utils.draw_segmentation_masks(img, masks, colors=colors)
+
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets",
+                            "fakedata", "draw_segm_masks_colors_util.png")
+
+        if not os.path.exists(path):
+            res = Image.fromarray(result.permute(1, 2, 0).contiguous().numpy())
+            res.save(path)
+
+        expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
+        self.assertTrue(torch.equal(result, expected))
+
+    def test_draw_segmentation_masks_no_colors(self):
+        img = torch.full((3, 20, 20), 255, dtype=torch.uint8)
+        # TODO
+        masks = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+        result = utils.draw_segmentation_masks(img, masks, colors=None)
+
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets",
+                            "fakedata", "draw_segm_masks_no_colors_util.png")
+
+        if not os.path.exists(path):
+            res = Image.fromarray(result.permute(1, 2, 0).contiguous().numpy())
+            res.save(path)
+
+        expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
+        self.assertTrue(torch.equal(result, expected))
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/utils.py

@@ -6,7 +6,7 @@
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont, ImageColor
 
-__all__ = ["make_grid", "save_image", "draw_bounding_boxes"]
+__all__ = ["make_grid", "save_image", "draw_bounding_boxes", "draw_segmentation_masks"]
 
 
 @torch.no_grad()
@@ -210,3 +210,59 @@ def draw_bounding_boxes(
             draw.text((bbox[0], bbox[1]), labels[i], fill=color, font=txt_font)
 
     return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1)
+
+
+@torch.no_grad()
+def draw_segmentation_masks(
+    image: torch.Tensor,
+    masks: torch.Tensor,
+    colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None,
+) -> torch.Tensor:
+
+    """
+    Draws segmentation masks on given image.
+    The values of the input image should be uint8 between 0 and 255.
+
+    Args:
+        image (Tensor): Tensor of shape (C x H x W)
+        masks (Tensor): Tensor of shape (H, W). Each containing predicted class.
+        labels (List[str]): List containing the labels of masks.
+        colors (List[Union[str, Tuple[int, int, int]]]): List containing the colors of masks. The colors can
+            be represented as `str` or `Tuple[int, int, int]`.
+    """
+
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"Tensor expected, got {type(image)}")
+    elif image.dtype != torch.uint8:
+        raise ValueError(f"Tensor uint8 expected, got {image.dtype}")
+    elif image.dim() != 3:
+        raise ValueError("Pass individual images, not batches")
+
+    img_to_draw = Image.fromarray(masks.byte().cpu().numpy()).resize(image.size()[1:])
+
+    if colors is None:
+        # Default color palette assumes 21 classes.
+        palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
+        colors_t = torch.as_tensor([i for i in range(21)])[:, None] * palette
+        color_arr = (colors_t % 255).numpy().astype("uint8")
+
+    else:
+        color_list = []
+        for color in colors:
+            if isinstance(color, str):
+                # This will automatically raise Error if rgb cannot be parsed.
+                fill_color = ImageColor.getrgb(color)  # + (100,)
+                color_list.append(fill_color)
+            elif isinstance(color, tuple):
+                # fill_color = color + (100,)
+                # Use the given colors list and create ndarray of colors.
+                color_list.append(color)
+
+        color_arr = np.array(color_list).astype("uint8")
+
+    img_to_draw.putpalette(color_arr)
+    img_to_draw = torch.from_numpy(np.array(img_to_draw))
+
+    # Project the drawn image to orignal one
+    image[: 1] = img_to_draw
+    return image