Utility to draw Semantic Segmentation Masks

[oke-aditya]

🚀 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.

[datumbox]

@oke-aditya +1 on the proposal. I like the fact of providing a similar API as for draw_bounding_boxes().

[oke-aditya]

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.

[datumbox]

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.

[oke-aditya]

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 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?

[datumbox]

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:

vision/test/test_utils.py

Line 88 in a0e5bbe

colors = ["green", "#FF00FF", (0, 255, 0), "red"]

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.

[oke-aditya]

Great ! I will send a PR let's investigate few options 😀

This is what I have in my mind.

Taken from FCN paper

[datumbox]

@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. :)

[oke-aditya]

I did have a go at this 😃 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")

[oke-aditya]

Some points, I'm unclear about

1. Is the user interested in plotting any text here ? If yes, then where and how ?
   If not, then we can forgo labels, font, font size parameters.

2. This is not very slow, it takes under 1 second for a decent sized image.

3. How do we assign default colors ? Otherwise everything will be colored with one color !

[datumbox]

@oke-aditya Thanks for providing the prototype implementation. This is really helpful because it raises lots of interesting questions about the proposed API.

image = torch.ones((3, 224, 224), dtype=torch.uint8)

From the above, I understand that you create a new image after making the prediction and before start drawing. Do you intentionally propose to create a new picture that contains only the masks OR you plan to apply the mask on top of the original image (with transparency) as we do on the bounding box util?

Concerning your questions, here are some thoughts:

1. Good point about the labels. I can't think of a way of including them without completely cluttering the picture. In most visualizations of semantic segmentation, the labels are omitted, so I think we could do the same here. The only cases I've seen of using the labels was to build a legend, which I'm not sure is useful here. If you have other ideas, let me know!
2. That's probably because we write one pixel at a time. Vectorizing the operation can probably speed up the execution. See example.
3. I think in this util the colours might need to be a mandatory field. Most-likely I would require the user to provide a list of colours equal to the number of classes. Alternative ways include defining a palette, choosing random colours or repeating input colours but all these have been discussed and turned down on the earlier PR versions of the bounding box util because they require a large number of assumptions.

@fmassa if you have any thoughts on the above, pitch in.

[oke-aditya]

I think make_grid avoids in-place modification, and we should avoid in-place operations for all plotting utils.

Maybe people want to have a look at the predictions, and do further post-processing to the image, then redraw the predictions. I guess we should just do new_image = image.clone() and avoid in place plotting for all utils.

I am not sure mask on the original image with transparency will be nice here. Maybe that is better for Instance segmentation than semantic?

1. I too think we can drop labels, Even the paper outputs attached above do not use labels. Plus we cannot create a consistent plotting of labels without overlap.

2. Vectorizing is fine, though this isn't very slow. I'm not sure how vectorizing will work with PIL.

3. Colors I guess should be mandatory, coz semantic segmentation is highly dependent. One way of color palette is here

# create a color pallette, selecting a color for each class
palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
colors = (colors % 255).numpy().astype("uint8")

# plot the semantic segmentation predictions of 21 classes in each color
r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)
r.putpalette(colors)

This is quite fast and good, but we lose control over colors. which can be easily supplied by a list or tuple.

Taken from Colab Notebook for instance segmentation. This is share to PyTorch hub,