Repository: pytorch/vision
Pull request: Random colors for drawing boxes
Closes #4528 Supersedes #4658

~~Work in progress will detail it soon.~~
Done! awaiting review.

:) Feels so good to be healthy and code again.

cc @pmeier (I guess Nicolas will return next week)


Problem statement:
Generate colors from random color palette in draw_bounding_boxes
### 🚀 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 :smiley:  to join.

Hints from issue discussion:
I can take this if that's okay.
I'm glad that you want to take it. 
Just need a green signal from the maintainers @NicolasHug @datumbox 



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.
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.


> 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.
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!
Hey @gessha  are you working on this?
Or
@ABD-01 would you like to take this up?
Yes, I can if no one's working on it.

I guess the main change involved will be here
https://github.com/pytorch/vision/blob/a75dc89abbbb233469d1dbe2ee00ee012dd70522/torchvision/utils.py#L203-L204
Just to make the variable `colors` be a list returned from [`_generate_color_palette`](https://github.com/pytorch/vision/blob/a75dc89abbbb233469d1dbe2ee00ee012dd70522/torchvision/utils.py#L303), right?

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. 


@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 :smiley: 

```

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.

```

Commit message: Fix bugs

Patch to review:
diff --git a/test/test_utils.py b/test/test_utils.py
index 30f144d8206..c671d2ee8bc 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -124,7 +124,7 @@ def test_draw_boxes_vanilla():
     img = torch.full((3, 100, 100), 0, dtype=torch.uint8)
     img_cp = img.clone()
     boxes_cp = boxes.clone()
-    result = utils.draw_bounding_boxes(img, boxes, fill=False, width=7)
+    result = utils.draw_bounding_boxes(img, boxes, fill=False, width=7, colors="white")
 
     path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_vanilla.png")
     if not os.path.exists(path):
@@ -149,7 +149,11 @@ def test_draw_invalid_boxes():
     img_tp = ((1, 1, 1), (1, 2, 3))
     img_wrong1 = torch.full((3, 5, 5), 255, dtype=torch.float)
     img_wrong2 = torch.full((1, 3, 5, 5), 255, dtype=torch.uint8)
+    img_correct = torch.zeros((3, 10, 10), dtype=torch.uint8)
     boxes = torch.tensor([[0, 0, 20, 20], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float)
+    labels_wrong = ["one", "two"]
+    colors_wrong = ["pink", "blue"]
+
     with pytest.raises(TypeError, match="Tensor expected"):
         utils.draw_bounding_boxes(img_tp, boxes)
     with pytest.raises(ValueError, match="Tensor uint8 expected"):
@@ -158,6 +162,10 @@ def test_draw_invalid_boxes():
         utils.draw_bounding_boxes(img_wrong2, boxes)
     with pytest.raises(ValueError, match="Only grayscale and RGB images are supported"):
         utils.draw_bounding_boxes(img_wrong2[0][:2], boxes)
+    with pytest.raises(ValueError, match="Number of boxes and labels mismatch."):
+        utils.draw_bounding_boxes(img_correct, boxes, labels_wrong)
+    with pytest.raises(ValueError, match="Number of colors should be greater or equal to the number of boxes."):
+        utils.draw_bounding_boxes(img_correct, boxes, colors=colors_wrong)
 
 
 @pytest.mark.parametrize(
diff --git a/torchvision/utils.py b/torchvision/utils.py
index cbbccae0b95..526a4b5c45a 100644
--- a/torchvision/utils.py
+++ b/torchvision/utils.py
@@ -169,6 +169,8 @@ def draw_bounding_boxes(
         colors (color or list of colors, optional): List containing the colors
             of the boxes or single color for all boxes. The color can be represented as
             PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+            By default, random colors are generated for boxes.
+            If labels are provided, boxes with same labels have same color.
         fill (bool): If `True` fills the bounding box with specified color.
         width (int): Width of bounding box.
         font (str): A filename containing a TrueType font. If the file is not found in this filename, the loader may
@@ -191,34 +193,45 @@ def draw_bounding_boxes(
     elif image.size(0) not in {1, 3}:
         raise ValueError("Only grayscale and RGB images are supported")
 
+    if labels:
+        if len(labels) != boxes.size(0):
+            raise ValueError("Number of boxes and labels mismatch. Please specify labels for each box.")
+
+    # List of colors should have minimum colors.
+    if colors:
+        if isinstance(colors, list):
+            if len(colors) < boxes.size(0):
+                raise ValueError("Number of colors should be greater or equal to the number of boxes.")
+
+    # Handle Grayscale images
     if image.size(0) == 1:
         image = torch.tile(image, (3, 1, 1))
 
     ndarr = image.permute(1, 2, 0).cpu().numpy()
     img_to_draw = Image.fromarray(ndarr)
-
     img_boxes = boxes.to(torch.int64).tolist()
 
     if fill:
         draw = ImageDraw.Draw(img_to_draw, "RGBA")
-
     else:
         draw = ImageDraw.Draw(img_to_draw)
 
     txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size)
 
+    if colors is None:
+        colors = _generate_random_color_palette(len(img_boxes))
+        if labels is not None:
+            label_color_map = dict(zip(labels, colors))
+            colors = [label_color_map[label] for label in labels]
+
     for i, bbox in enumerate(img_boxes):
-        if colors is None:
-            color = None
-        elif isinstance(colors, list):
+        if isinstance(colors, list):
             color = colors[i]
         else:
             color = colors
 
         if fill:
-            if color is None:
-                fill_color = (255, 255, 255, 100)
-            elif isinstance(color, str):
+            if isinstance(color, str):
                 # This will automatically raise Error if rgb cannot be parsed.
                 fill_color = ImageColor.getrgb(color) + (100,)
             elif isinstance(color, tuple):
@@ -387,6 +400,30 @@ def _generate_color_palette(num_masks: int):
     return [tuple((i * palette) % 255) for i in range(num_masks)]
 
 
+def _generate_random_color() -> Tuple[int, int, int]:
+    """
+    Generates a random RGB Color.
+
+    Returns:
+        color Tuple(int, int, int): A Tuple containing Random RGB Color.
+    """
+    return tuple(torch.randperm(256)[:3].tolist())
+
+
+def _generate_random_color_palette(num_colors: int) -> List[Tuple[int, int, int]]:
+    """
+    Generates a random RGB Color palette.
+
+    Args:
+        num_colors (int): Integer denoting number of random RGB colors to generate
+
+    Returns:
+        colors(List[Tuple[int, int, int]): A List of Tuples each containing random RGB color.
+            Each color is Tuple containing RGB values.
+    """
+    return [_generate_random_color() for i in range(num_colors)]
+
+
 def _log_api_usage_once(obj: Any) -> None:
     if not obj.__module__.startswith("torchvision"):
         return


Code review comments:
diff_hunk: @@ -387,6 +400,30 @@ def _generate_color_palette(num_masks: int):
     return [tuple((i * palette) % 255) for i in range(num_masks)]
 
 
+def _generate_random_color() -> Tuple[int, int, int]:
+    """
+    Generates a random RGB Color.
+
+    Returns:
+        color Tuple(int, int, int): A Tuple containing Random RGB Color.
+    """
+    return tuple(torch.randperm(256)[:3].tolist())
path: torchvision/utils.py
text: @author:
Alternative is to use random and generate a tuple. We need to use tuple and not list as PIL accepts only Tuples

```
return (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
```

Let me know if this is better.


@author:
Also do you need tests for generating random color and color palette?
I can add those as well.

Is it worth considering making these two utilities (`generate_random_color` and `generate_random_color_palette`) as part of public API?



@author:
A nice improvement can be to extend this utility and allow to generate random RGBA color as well? 
But I don't see any great use case.


@user1:
What about

```suggestion
    return tuple(torch.randint(256, (3,)).tolist())
```

Using `torch.randperm` will never give you the same value for two channels, but I don't think this is required here?

@author:
Yes you are right Philip. I didn't know we could do this.

diff_hunk: @@ -191,34 +193,45 @@ def draw_bounding_boxes(
     elif image.size(0) not in {1, 3}:
         raise ValueError("Only grayscale and RGB images are supported")
 
+    if labels:
+        if len(labels) != boxes.size(0):
+            raise ValueError("Number of boxes and labels mismatch. Please specify labels for each box.")
+
+    # List of colors should have minimum colors.
+    if colors:
+        if isinstance(colors, list):
+            if len(colors) < boxes.size(0):
+                raise ValueError("Number of colors should be greater or equal to the number of boxes.")
path: torchvision/utils.py
text: @author:
This is used to silently throw list index out of range error, if number of colors were less than boxes.
I guess it's better to throw an error.

diff_hunk: @@ -169,6 +169,8 @@ def draw_bounding_boxes(
         colors (color or list of colors, optional): List containing the colors
             of the boxes or single color for all boxes. The color can be represented as
             PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+            By default, random colors are generated for boxes.
+            If labels are provided, boxes with same labels have same color.
path: torchvision/utils.py
text: @user1:
I'm concerned with this API change because the behaviour changes depending on whether the labels are defined. Effectively the new API tries to do things for the user instead of letting them decide what to do. I think it would be best to stick to generating a random palette for as many masks as we detected. This will be similar to what we do on the segmentation drawing tool and will enable us to reuse methods such as `_generate_color_palette()`.



@user3:
@user1 would you mind detailing what you mean by "enable us to reuse methods such as _generate_color_palette()" ?

Having one box color per label is a default behaviour that makes sense to me, and @user2-aditya followed my earlier suggestion of https://github.com/pytorch/vision/issues/4528#issuecomment-933232255. It avoids users some extra code.

`draw_segmentation_masks()` doesn't accept a `label` parameter, so it's not unexpected that its behaviour differs from `draw_bounding_boxes()` in that respect. 

@user1:
@user3 I was thinking that `utils._generate_color_palette(num_masks)` can be used to generate a colour per mask (not per label) as currently happens on the segmentation case. 

Concerning the lack of support for labels on segmentation, this is indeed the situation with our utility but doesn't have to be the case going forwards. Other libraries such as [detectron2](https://github.com/facebookresearch/detectron2) offer visualizations for segmentation and masks that display the labels.

TBH I haven't followed the original discussions (I try to catch up coming back from PTO). My thoughts on this is that if we want to keep the utility simple, we could let the users define one colour per mask. If they need to apply a more complex logic (such as keeping the same colour per label), this should be done outside of the utility using their own custom logic and provide the list of colours to our util. This would align the approach across segmentation and detection. Let me know your thoughts.

@user3:
> Concerning the lack of support for labels on segmentation, this is indeed the situation with our utility but doesn't have to be the case going forwards. Other libraries such as detectron2 offer visualizations for segmentation and masks that display the labels.

I agree. And if we were to support labels in `draw_segmentation_masks()`, I think it would also make sense to have one color per label by default :)

I also agree that complexity is a concern, both regarding the code and regarding what the user would expect from that function. As far as I can tell the only extra code that handling labels adds are these few lines below:


```py
        if labels is not None:
            label_color_map = dict(zip(labels, colors))
            colors = [label_color_map[label] for label in labels]
```

which I feel is minimal.

@author:
Users still can draw colors for their labels. If both labels and colors are passed. The `ith` box would get the `ith` color, irrespective of the labels.
I guess the current improvement just improves how the the default colors are handled. The need to pass a color palette reduces.

diff_hunk: @@ -191,34 +193,45 @@ def draw_bounding_boxes(
     elif image.size(0) not in {1, 3}:
         raise ValueError("Only grayscale and RGB images are supported")
 
+    if labels:
path: torchvision/utils.py
text: @author:
We need to check these. As agreed. If labels are specified, all the boxes should have labels.