Repository: pytorch/vision

Issue:
Utility to draw Semantic Segmentation Masks
## 🚀 Feature

We recently added utility to draw bounding boxes, which works really well with detection models. #2785 #2556 
It might be nice to draw segmentation masks which we obtain from instance segmentation models.

## Motivation

Same as bounding box utils. It is very useful to have these. It reduces the dependence of users over other plotting libraries.

## Pitch

Our API should be compatible with segmentation models, so we should probably use Tensors.
I think most params remain as same as the previous util. This keeps consistency too.

```
@torch.no_grad()
def draw_segmentation_masks(image: torch.Tensor,
    masks: torch.Tensor,
    labels: Optional[List[str]] = None,
    colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None,
    width: int = 1,
    font: Optional[str] = None,
    font_size: int = 10)
```

We might need to see a method using which we can draw with PIL.
We used `draw.rectangle()` to draw a box in utils, maybe there is some functionality in PIL that can help us draw shapes.

## Alternatives

Let's discuss further, how this API would work with our models.
Also, it would be nice if this works directly for instance segmentation model Mask RCNN.



Discussion:
@oke-aditya +1 on the proposal. I like the fact of providing a similar API as for `draw_bounding_boxes()`.
I would also suggest adding an example notebook/script / tutorial where we can show how both the utils work.
I will give a try at this using PIL and post a sample result here, let's see if it is satisfactory.
> I will give a try at this using PIL and post a sample result here, let's see if it is satisfactory.

Perfect! I think using PIL should be fine given that's how the bounding box version works. Looking forward to it.
Hey @datumbox. For Instance segmentation model we can do something like #3280.

I'm looking into semantic segmentation models,
The output of model seems to be the probability of each pixel in the image into 1 / 21 (pretrained classes, if it is COCO)

I guess then the predicted label is `argmax()` of the n classes on each pixel.

I think colors remain the same.
Maybe a user simply passes a `Tuple(int, int, int)` each containing mapping to ith label?
Or the user passes a `List[str]` each mapping to ith label?

I think we can use [ImageDraw.point](https://pillow.readthedocs.io/en/stable/reference/ImageDraw.html#PIL.ImageDraw.ImageDraw.point) to draw each point? 
This would be very slow I think for bigger images. Since we fill color on every pixel of tensor.

The width makes little sense here since if we will be drawing each pixel on image?


> I guess then the predicted label is argmax() of the n classes on each pixel.

Agreed.

> I think colors remain the same.

Yes they should remain the same. We currently support both tuples of RGB ints and strings. See an example we have on the tests:
https://github.com/pytorch/vision/blob/a0e5bbeab4b39ed1078384d10f742f757f74df00/test/test_utils.py#L88

> Or the user passes a List[str] each mapping to ith label?

We should use `colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None,` like on boxes.

> This would be very slow I think for bigger images. Since we fill color on every pixel of tensor.

Worth investigating a couple of options and measure them. If you send a PR and tag me, I'm happy to share feedback.

> The width makes little sense here since if we will be drawing each pixel on image?

Yes, I think you are right.
Great ! I will send a PR let's investigate few options 😀

This is what I have in my mind.

![image](https://user-images.githubusercontent.com/47158509/105824638-a1ffa500-5fe4-11eb-995d-44f20c5ccf6e.png)


Taken from FCN [paper](https://arxiv.org/pdf/1411.4038.pdf)


@oke-aditya Awesome, looking forward to it. 

I'm actually doing work around segmentation, so I'm definitely going to use your prototypes to validate the models. :)
I did have a go at this :smiley:  Here is code to reproduce a simple mask using `draw.point`. It is not that slow

```
image = torch.ones(1, 3, 224, 224)
fcn = fcn_resnet50(pretrained=False, num_classes=21)
fcn = fcn.eval()

with torch.no_grad():
   fcn_out = fcn(image)['out'][0]
    # print(fcn_out)

output_predictions = fcn_out.argmax(0)
# print(output_predictions)
# print(output_predictions.shape)

img_preds = output_predictions.to(torch.int64).tolist()

image = torch.ones((3, 224, 224), dtype=torch.uint8)
ndarr = image.permute(1, 2, 0).numpy()
img_to_draw = Image.fromarray(ndarr)

labels = ["blue", "red", "purple", "orange", "red", "purple", "orange", "red",
              "purple", "black", "red", "yellow", "black", "red", "yellow", "orange", "red",
              "yellow", "aqua", "red", "yellow", "aqua"]

draw = ImageDraw.Draw(img_to_draw)

for i in range(len(output_predictions)):
    for j in range(len(output_predictions)):
        draw.point((i, j), fill=labels[output_predictions[i][j]])

img_to_draw.save("temp.png")
```
![temp](https://user-images.githubusercontent.com/47158509/106366498-4b86c380-6362-11eb-9faf-9c307a3ff8ae.png)



Resolving patch:
diff --git a/torchvision/utils.py b/torchvision/utils.py
--- a/torchvision/utils.py
+++ b/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()
@@ -153,7 +153,7 @@ def draw_bounding_boxes(
     If filled, Resulting Tensor should be saved as PNG image.
 
     Args:
-        image (Tensor): Tensor of shape (C x H x W)
+        image (Tensor): Tensor of shape (C x H x W) and dtype uint8.
         boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that
             the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and
             `0 <= ymin < ymax < H`.
@@ -210,3 +210,61 @@ 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,
+    alpha: float = 0.2,
+    colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None,
+) -> torch.Tensor:
+
+    """
+    Draws segmentation masks on given RGB image.
+    The values of the input image should be uint8 between 0 and 255.
+
+    Args:
+        image (Tensor): Tensor of shape (3 x H x W) and dtype uint8.
+        masks (Tensor): Tensor of shape (num_masks, H, W). Each containing probability of predicted class.
+        alpha (float): Float number between 0 and 1 denoting factor of transpaerency 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")
+    elif image.size()[0] != 3:
+        raise ValueError("Pass an RGB image. Other Image formats are not supported")
+
+    num_masks = masks.size()[0]
+    masks = masks.argmax(0)
+
+    if colors is None:
+        palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
+        colors_t = torch.as_tensor([i for i in range(num_masks)])[:, 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)
+                color_list.append(fill_color)
+            elif isinstance(color, tuple):
+                color_list.append(color)
+
+        color_arr = np.array(color_list).astype("uint8")
+
+    _, h, w = image.size()
+    img_to_draw = Image.fromarray(masks.byte().cpu().numpy()).resize((w, h))
+    img_to_draw.putpalette(color_arr)
+
+    img_to_draw = torch.from_numpy(np.array(img_to_draw.convert('RGB')))
+    img_to_draw = img_to_draw.permute((2, 0, 1))
+
+    return (image.float() * alpha + img_to_draw.float() * (1.0 - alpha)).to(dtype=torch.uint8)


Test patch:
diff --git a/test/assets/fakedata/draw_segm_masks_colors_util.png b/test/assets/fakedata/draw_segm_masks_colors_util.png
new file mode 100644
Binary files /dev/null and b/test/assets/fakedata/draw_segm_masks_colors_util.png differ
diff --git a/test/assets/fakedata/draw_segm_masks_no_colors_util.png b/test/assets/fakedata/draw_segm_masks_no_colors_util.png
new file mode 100644
Binary files /dev/null and b/test/assets/fakedata/draw_segm_masks_no_colors_util.png differ
diff --git a/test/test_utils.py b/test/test_utils.py
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -9,6 +9,30 @@
 import torchvision.transforms.functional as F
 from PIL import Image
 
+masks = torch.tensor([
+    [
+        [-2.2799, -2.2799, -2.2799, -2.2799, -2.2799],
+        [5.0914, 5.0914, 5.0914, 5.0914, 5.0914],
+        [-2.2799, -2.2799, -2.2799, -2.2799, -2.2799],
+        [-2.2799, -2.2799, -2.2799, -2.2799, -2.2799],
+        [-2.2799, -2.2799, -2.2799, -2.2799, -2.2799]
+    ],
+    [
+        [5.0914, 5.0914, 5.0914, 5.0914, 5.0914],
+        [-2.2799, -2.2799, -2.2799, -2.2799, -2.2799],
+        [5.0914, 5.0914, 5.0914, 5.0914, 5.0914],
+        [5.0914, 5.0914, 5.0914, 5.0914, 5.0914],
+        [-1.4541, -1.4541, -1.4541, -1.4541, -1.4541]
+    ],
+    [
+        [-1.4541, -1.4541, -1.4541, -1.4541, -1.4541],
+        [-1.4541, -1.4541, -1.4541, -1.4541, -1.4541],
+        [-1.4541, -1.4541, -1.4541, -1.4541, -1.4541],
+        [-1.4541, -1.4541, -1.4541, -1.4541, -1.4541],
+        [5.0914, 5.0914, 5.0914, 5.0914, 5.0914],
+    ]
+], dtype=torch.float)
+
 
 class Tester(unittest.TestCase):
 
@@ -96,6 +120,35 @@ 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, 5, 5), 255, dtype=torch.uint8)
+        colors = ["#FF00FF", (0, 255, 0), "red"]
+        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)
+        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()