Change default colors in draw_segmentation_masks()

🚀 The feature

The current utils.draw_segmentation_masks() uses _generate_color_palette(num_masks) by default if the colors of the masks are not specified by the user. _generate_color_palette(num_masks) has a behavior that it always give black as the first color. As an example:

# https://github.com/pytorch/vision/blob/b403bfc771e0caf31efd06d43860b09004f4ac61/torchvision/utils.py#L533
def _generate_color_palette(num_objects: int):
    palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
    return [tuple((i * palette) % 255) for i in range(num_objects)]

# generate colors for 1-3 objects
for i in range(1, 4):
    print(_generate_color_palette(i))

yields

[(tensor(0), tensor(0), tensor(0))]
[(tensor(0), tensor(0), tensor(0)), (tensor(1), tensor(127), tensor(31))]
[(tensor(0), tensor(0), tensor(0)), (tensor(1), tensor(127), tensor(31)), (tensor(2), tensor(254), tensor(62))]

The proposal is to change this behavior such that black is never used in the color palette. The reason is sometimes users need to draw the segmentation masks on a pure black background, and the current behavior will hide the first object since it is also black, and it may be hard for the users to notice that an object is missing.

Motivation, pitch

Here is an example of when I'm trying to draw segmentation masks on a black background.

An image (from COCO) for reference:

The default behavior:

viz = utils.draw_segmentation_masks(
        torch.zeros(img.shape, dtype=torch.uint8), 
        masks=masks, 
        alpha=1)

The desired behavior would be something like:

viz = utils.draw_segmentation_masks(
        torch.zeros(img.shape, dtype=torch.uint8), 
        masks=masks, 
        alpha=1, 
        colors=[tuple(_) for _ in np.random.randint(100, 256, size=(len(masks), 3)).tolist()])

Alternatives

No response

Additional context

The easiest way to modify this is to iterate from the second element of the current _generate_color_palette() function, i.e.,

def _generate_color_palette_v2(num_objects: int):
    palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
    return [tuple((i * palette) % 255) for i in range(1,num_objects+1)]

which will give

However, personally, I think the colors are still not visually distinct enough (most colors are green/blue-ish), so I wonder what's the intuition behind the current implementation, and if we can somehow improve it.

I'm happy to make a diff, please comment if you have any thoughts.

The relevant entry point is torchvision/utils.py, where draw_segmentation_masks() calls _generate_color_palette(num_masks); start there and compare the current palette with the proposed nonzero sequence. Done means the default palette never assigns black while still producing colors for every mask; the issue also flags visual distinctness as an open design question.