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.

[oke-aditya]

I too think so, if we set the background with transparency channel we might be able to project it into the image. And additionally have a fill parameter (like we did for boxes). Let me have run with this and post the outputs 😃

[oke-aditya]

I had a brief thought about all points.

1. Yes, providing the image tensor is unnecessary. But let’s go for consistency in utils, it doesn’t hurt much. It would be surprising for a user to remember that we don’t pass image only in this util.

We can just do new_image = image.clone() and continue plotting in new_image. Also, probably we should do this for bounding boxes too and avoid in-place operation there.

1. The above example is vectorized ! and has good defaults. We could also ask the user to optionally provide a list of colors for each class.

So the new API proposal 😄

@torch.no_grad()
def draw_segmentation_masks(
    image: torch.Tensor,
    masks: torch.Tensor,
    colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None,
)

Args:
        image (Tensor): Tensor of shape (C x H x W)
        masks (Tensor): Tensor of shape (H, W). Each containing predicted class.
        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]`.
        It should be equal to the number of output classes.