Added utility to draw segmentation masks

[oke-aditya]

Closes #3272.

Initial Implementation

• Adds Code
• Adds Tests
• Adds Docs

[oke-aditya]

I need small help here as this projects to a black background.

My guess is we need an alpha channel which will make masks transparent ?

[datumbox]

I supposed there are many ways to do it. One that I had in mind originally (which might not be the optimal) is to convert the img_to_draw from palette to RGBA, replace the background colour with transparent and then combine it with image to achieve the "projection". Worth experimenting with the approach because it's likely there is a better way.

[oke-aditya]

Yep, DETR has very nice visualization, but they use matplotlib. Unsure how to reproduce them though.
As you pointed out before Mask RCNN utils have nice way to apply mask too.

[datumbox]

I just sent it for reference not necessarily for reproduction. :)

[datumbox]

@oke-aditya It's going towards the right direction. I left a few comments to get your thoughts.

BTW @fmassa brought to my attention a nice guide they have to DETR with some visualization utils we might want to look into for inspiration.

[datumbox]

I supposed there are many ways to do it. One that I had in mind originally (which might not be the optimal) is to convert the img_to_draw from palette to RGBA, replace the background colour with transparent and then combine it with image to achieve the "projection". Worth experimenting with the approach because it's likely there is a better way.

[oke-aditya]

Hi @datumbox I just resolved the num_classes hardcoding by changing params to probabilities.

[oke-aditya]

I'm bit unsure now how to proceed, one way I think is to create another util to apply weighted mask just like MaskRCNN matterpott benchmark

We can probably use this and apply the mask here.

Unsure of tests, it might be too lengthy and calculative to write 20 x 20 x 3 tensor with probabilities. I thought to just run FCN on torch.full and use its outputs.

Let me know how to proceed ! This part seems tricky

This does double check, but couples tests with models.

[datumbox]

@oke-aditya Great changes, I think we are almost there!

As I commented above, there are multiple ways to do this. Here is a hacky, quick and dirty approach to get an idea. I'm sure you can do it in a much better way:

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

    classifications = masks.argmax(0).byte()
    img_to_draw = Image.fromarray(classifications.cpu().numpy())

    if colors is None:
        num_classes = masks.size(0)
        palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1, 0])
        colors_t = torch.as_tensor([i for i in range(num_classes)])[:, None] * palette
        color_arr = (colors_t % 255).numpy().astype("uint8")
        color_arr[1:, 3] = 255
    else:
        color_list = []
        for color in colors:
            if isinstance(color, str):
                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")

    img_to_draw.putpalette(color_arr)

    img_to_draw = torch.from_numpy(np.array(img_to_draw.convert('RGBA')))
    img_to_draw = img_to_draw.permute((2, 0, 1))

    alpha = 0.6
    return (torch.cat([image, torch.full(image.shape[1:], 255).unsqueeze(0)]).float()*alpha+img_to_draw.float()*(1.0-alpha)).to(dtype=torch.uint8)

Ping me when you have a demo you are comfortable with to discuss the last details of the API. :)

[oke-aditya]

Sorry for the delay @datumbox . Here are few ouputs for different values of alpha

Alpha = 0.2

Alpha = 0.3

Alpha = 0.6

Alpha = 0.7

Also I made alpha as a paramter as it is really useful, to remove or keep background.

Another thought I had was to make util apply_mask which can be used to project mask.
This might be useful in some other frequently use alpha blending, alpha masking cases.
So let me know !

def apply_mask(image, mask, alpha):

torch.cat([image, torch.full(image.shape[1:], 255).unsqueeze(0)]).float()
            * alpha + mask.float() * (1.0 - alpha)).to(dtype=torch.uint8)

[fmassa]

Hi,

The PR looks very nice, thanks a lot!

I have one suggestion which I think would make the function more generic, and wouldn't involve too many changes. Let me know what you think

[fmassa]

Instead of being specific for semantic segmentation and not supporting instance segmentation nor panoptic segmentation, I think we could make this slightly more generic while supporting all the use-cases I mentioned. The idea would be to accept a mask as a [num_masks, H, W] boolean Tensor.
This way, the user can get the semantic segmentation masks to pass to this function as follows

out.argmax(0) == torch.arange(out.shape[0])[:, None, None]

[fmassa]

If we add support for instance segmentation and panoptic segmentation, I think it would be a good idea to add an example from using the output of a semantic segmentation model and an instance segmentation model (for example from those from torchvision)

[oke-aditya]

Sure, this is a TODO. I will add GitHub gist and other minor documentation improvements for both the utilities.

Can I do this in a follow-up PR which will address all the issues as mentioned in #3364 ?

[fmassa]

Sure, adding documentation improvements on a follow-up PR is ok with me. What do you think about the other comment as well? Because it would be a breaking change in functionality if we support it, so better do it once (specially that the branch cut is happening very soon so if we merge it now it can get integrated in the release, in which case breaking backwards-compatibility is more annoying)

[oke-aditya]

Definetely, I will refactor the other comment ASAP in this PR 😄 I understand how bad it would be with BC change.

[oke-aditya]

Slight problem while implementing.
Seems that argmax is not supported for bool tensor

I am pushing my latest changes, can someone please have look ?

A simple code to reproduce bug

masks =    [
        [False, False, False, False, False],
        [True, True, True, True, True],
        [False, False, False, False, False],
        [False, False, False, False, False],
        [False, False, False, False, False]
    ],
    [
        [True, True, True, True, True],
        [False, False, False, False, False],
        [True, True, True, True, True],
        [True, True, True, True, True],
        [False, False, False, False, False]
    ],
    [
        [False, False, False, False, False],
        [False, False, False, False, False],
        [False, False, False, False, False],
        [False, False, False, False, False],
        [True, True, True, True, True],
    ]
], dtype=torch.bool)

masks = masks.argmax(0)
print(masks)

RuntimeError: "argmax_cpu" not implemented for 'Bool'

Maybe there is some workaround? I guess this was the only change, to support bool Tensor instead of float Tensor

The error would occur when we try to get [H, W] from [num_masks, H, W] tensor.

[fmassa]

@oke-aditya maybe I'm missing something, but if we pass a bool tensor we don't need to compute the argmax anymore, because the independent masks have already been computed?

EDIT: oh I see, you don't perform a for loop over each one of the masks as of now. Computing the argmax could be done after casting the mask to float for example, but note that in instance segmentation each pixel can be covered by multiple masks, so the argmax wouldn't be enough to handle those. But it can be a first approximation

[oke-aditya]

Computing the argmax could be done after casting the mask to float for example, but note that in instance segmentation each pixel can be covered by multiple masks, so the argmax wouldn't be enough to handle those. But it can be a first approximation

That's possible. But may I know why we decided to pass a Bool Tensor and not a Float mask Tensor ?
Sorry if the question is slightly dumb.

Edit: My idea was to make initial implementation compatible with single channel masks. That would make it compatible with Mask RCNN. I thought in the same previous function, we could handle single channel case differently.

[oke-aditya]

I think with release cut coming quite soon, we could wait for and add in next release ? (I mean after 0.9.0)

[oke-aditya]

Hey @fmassa and @datumbox. Any thoughts on how to proceed further with this. I'm willing to incorporate any changes requested 😃

[fmassa]

Sorry for the delay in reviewing.

I've made a comment to unblock you, but I think in the future we might need to change the implementation to let it handle overlapping masks (which is the whole purpose of the function accepting a tensor with a per-object map).

The input mask doesn't need to be a boolean by the way, but it can be a floating point representing probabilities for the given instance.

But to unblock for now let's just make this small change I proposed, and we can improve this in a follow-up PR

[fmassa]

In order to unblock you can do masks.to(torch.int64).argmax(0) or cast it to float if you want.
This won't handle overlapping instances very well thought, and we will need to remove this in the future and probably replace it with a for loop so that overlapping masks are taken into account.

Plus, by using for loops and letting the mask be a floating point if the user wants, we can allow the user to have heatmaps being passed (instead of only binary maps), which would be very nice

[oke-aditya]

Extremely sorry for the delay (my health let me down 😢)

I think that boolean tensor leads to some limitations and again we re-cast it to int/float.

I refactored to use floating-point masks. Each point represents the probability of class.
Floating-point masks make more sense as we can either take argmax() and plot the best masks.
Or we could take topk to plot top masks. Or as you said we can cover up the overlapping case.

Currently, this code accepts a floating-point tensor of (num_masks, H, W)

Let me know what we need to do in further PRs / this PR.

[fmassa]

Thanks!