Generate colors from random color palette in draw_bounding_boxes

[oke-aditya]

🚀 The feature

The current default white color looks bad. Maybe we can generate colors from palette as we do in drawing masks?

Motivation, pitch

It should be simple tweak to _generate_color_palette to pass num_boxes and get a List of RGB values. Instead of List[Tensor]. The fix is simple and that can be passed to get RGB colors.

Alternatives

Users can anyway use color="white" to get white boxes.

Additional context

It's simple issue and good-first. We can probably wait for some new contributor 😃 to join.

[gessha]

I can take this if that's okay.

[oke-aditya]

I'm glad that you want to take it.
Just need a green signal from the maintainers @NicolasHug @datumbox

[NicolasHug]

Thanks for the proposal @oke-aditya and @gessha for volunteering.

That sounds good to me, we should just make sure that the same color is generated for all boxes with the same label.

[oke-aditya]

Well I'm not sure how we would ensure that same color is generated for all boxes with same label.
My thoughts were to generate color palette equal to number of boxes. And then we simply pick ith color from the palette.

We don't keep consistency of randomly generating same color for same label in drawing segmentation masks too.

[NicolasHug]

Well I'm not sure how we would ensure that same color is generated for all boxes with same label.

We can generate num_labels colors and then assign one color per label. Obviously we don't need to do that if labels isn't specified

We don't keep consistency of randomly generating same color for same label in drawing segmentation masks too.

True but draw_segmentation_masks doesn't accept a label anyway, as it probably doesn't need one.

[oke-aditya]

Great, I think it should be feasible with some refactor. As of now we plot each bbox and then the labels separately.

@gessha feel free to send a PR!

[oke-aditya]

Hey @gessha are you working on this?
Or
@ABD-01 would you like to take this up?

[ABD-01]

Yes, I can if no one's working on it.

I guess the main change involved will be here

vision/torchvision/utils.py

Lines 203 to 204 in a75dc89

	if colors is None:
	color = None

Just to make the variable colors be a list returned from _generate_color_palette, right?

[oke-aditya]

Yes and use the colors the generated from the palette in draw_bounding_boxes.
Create colors upto num_labels.

Map each label to a color.

Since labels is a list and generated color palette is also a list. Both have the same length.
(As we generated colors equal to num_labels)

Use these colors to draw boxes.

[oke-aditya]

@ABD-01 I spent some time and thought a bit more.

Here is a brief pseudocode (highly untested, written on paper), this can be simplified / optimized as per needs 😃

if colors is None:
   if labels is None:
       # We have no choice to create colors equal to number of boxes
       # gen_colors is  List[Tuple[int, int, int]]

        gen_colors = genereate_color_palette(len(img_boxes))
   else:
      # One way is to generate number of colors equal to number of unique labels.
      # Or alternatively, just generate excess colors and map each label just once.
      gen_colors = generate_color_palette(len(labels)) (or probably even len(img_boxes))
      
     # Since dictionary will have only unique keys.  
      label_color_map = {}
      for label, color in zip(label, gen_colors):
          label_color_map[label] = color

# Now we have color palette we can enumerate over boxes.
for i, bbox in enumerate(img_boxes):
    if colors is None:
        if labels is None:
            # No labels so just use the generated box colors
           color = gen_colors[i]
            # We have labels, so simply make use of label_color_map
        else:
           color = label_color_map[labels[i]]     
   else:
          # Main code continues.