``draw_segmentation_masks`` throws IndexError when given [0, H, W] in mask

[oxabz]

🐛 Describe the bug

draw_segmentation_masks throws an IndexError when it recieves an empty mask tensor as it tries to check preconditions on colors.

It'd be good to fix the bug to allow to draw empty mask tensors or to add the precondition in the doc and a clear assert.

As an asside it looks like draw_bounding_boxes has the same probleme.

from torchvision.utils import draw_segmentation_masks
import torch 
from torch import uint8

img = torch.rand((3,100,100)).type(uint8)
masks = torch.rand((0,100,100)).bool()

draw_segmentation_masks(img, masks)

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
/tmp/ipykernel_548012/3010128033.py in <module>
      6 masks = torch.rand((0,100,100)).bool()
      7 
----> 8 draw_segmentation_masks(img, masks)

/data/apps/conda/matthieu/envs/apriorics/lib/python3.9/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     26         def decorate_context(*args, **kwargs):
     27             with self.__class__():
---> 28                 return func(*args, **kwargs)
     29         return cast(F, decorate_context)
     30 

/data/apps/conda/matthieu/envs/apriorics/lib/python3.9/site-packages/torchvision/utils.py in draw_segmentation_masks(image, masks, alpha, colors)
    278     if not isinstance(colors, list):
    279         colors = [colors]
--> 280     if not isinstance(colors[0], (tuple, str)):
    281         raise ValueError("colors must be a tuple or a string, or a list thereof")
    282     if isinstance(colors[0], tuple) and len(colors[0]) != 3:

IndexError: list index out of range

Versions

PyTorch version: 1.10.2
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: CentOS Stream release 8 (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-10)
Clang version: Could not collect
CMake version: version 3.20.2
Libc version: glibc-2.28

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.18.0-365.el8.x86_64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla P100-PCIE-16GB
GPU 1: Tesla P100-PCIE-16GB

Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.8.3.2
/usr/lib64/libcudnn_adv_infer.so.8.3.2
/usr/lib64/libcudnn_adv_train.so.8.3.2
/usr/lib64/libcudnn_cnn_infer.so.8.3.2
/usr/lib64/libcudnn_cnn_train.so.8.3.2
/usr/lib64/libcudnn_ops_infer.so.8.3.2
/usr/lib64/libcudnn_ops_train.so.8.3.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.942
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.1
[pip3] pytorch-lightning==1.5.9
[pip3] torch==1.10.2
[pip3] torchaudio==0.10.2
[pip3] torchmetrics==0.7.1
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.11.3
[conda] blas 1.0 mkl
[conda] cudatoolkit 10.2.89 hfd86e86_1
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] mkl 2022.0.1 h06a4308_117
[conda] mypy 0.942 pypi_0 pypi
[conda] mypy-extensions 0.4.3 pypi_0 pypi
[conda] numpy 1.22.1 pypi_0 pypi
[conda] pytorch 1.10.2 py3.9_cuda10.2_cudnn7.6.5_0 pytorch
[conda] pytorch-lightning 1.5.9 pypi_0 pypi
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] torchaudio 0.10.2 py39_cu102 pytorch
[conda] torchmetrics 0.7.1 pypi_0 pypi
[conda] torchsummary 1.5.1 pypi_0 pypi
[conda] torchvision 0.11.3 py39_cu102 pytorch

[oke-aditya]

Can you try

masks = torch.rand((1,100,100)).bool()

instead? or even

masks = torch.rand((100,100)).bool()

In the doc it says

        masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool.

We cannot have num_masks as 0 as if a mask is present we start the count from 1.

[NicolasHug]

@oxabz in which scenario did you end up with a masks tensor with a zero first dimension? As @oke-aditya mentioned, the semantic of that means that there are no masks to plot.

I guess we could assert num_masks > 0, but I wonder if there is some misunderstanding in how to use this function. Also, could you please clarify what you mean by

allow to draw empty mask tensors

[oxabz]

The scenario in which the bug appeared is the following

masks = pred["masks"][idxs].cpu().squeeze(1) > 0.5
img_w_mask = draw_segmentation_masks(img, masks, alpha=0.5, colors=colors_preds)

The model gives us several mask for each detected object. When the model diverge it might gives us a N*W*H tensor full of zeroes wich can be squeezed down to 0*W*H.

In that case the bug doesnt cause that big of a problem since the training is already useless in that scenario. But putting a clear assert could save future head scratching.

could you please clarify what you mean by

allow to draw empty mask tensors

I think, that alternatively the behaviour of the function could be to simply return the unchanged image when it is given no mask.

[NicolasHug]

Thanks for the details @oxabz , that makes sense. We'd be happy to review a PR that raises a clean ValueError when num_masks == 0. Would you like to submit a PR, or should we ask for other contributors?

[vadimkantorov]

could it just do no-op if num_masks == 0 and exit early (without throwing)?

[NicolasHug]

From @oxabz 's comment:

In that case the bug doesnt cause that big of a problem since the training is already useless in that scenario. But putting a clear assert could save future head scratching.

num_masks is 0 in situations where something went wrong in the training, or in setting the threshold. It's probably best to raise a proper error in these situations rather than silently passing, to let the user know that there is a problem somewhere else.

But I don't feel too strongly about it, it's just a util. We can also raise a warning and return early. But we should avoid being silent about this.

[vadimkantorov]

Not necessarily - sometimes in datasets there are examples with no ground truth masks (e.g. only boxes or only crowd/difficult instances which are filtered out). And it's sometimes better to not throw rather than having to roll proper dataset filtering immediately.

Not failing with empty-sized dim tensors is a general theme of regular pytorch ops. So having it break in the vis utility is often a nuisance.

[oke-aditya]

I think we can just raise a Warning instead of throwing error? My guess is that this utility might just be used in a loop and an error would break the flow of training / logging.

[NicolasHug]

Sounds good, let's throw a warning then

[oxabz]

I can handle the pull request if everything is settled

[oke-aditya]

Feel free to take this @oxabz

[oxabz]

Okay so I've added the warning and made sure that the function wouldn't throw, but I'm having trouble with the automated test. The build for the unit test is giving me really a cryptic error on part of the code I havent changed. I might need a bit of help on that.