Support single color in utils.draw_bounding_boxes (#4075)

Author: Separius

test/test_utils.py

@@ -113,6 +113,18 @@ def test_draw_boxes():
     assert_equal(img, img_cp)
 
 
+@pytest.mark.parametrize('colors', [
+    None,
+    ['red', 'blue', '#FF00FF', (1, 34, 122)],
+    'red',
+    '#FF00FF',
+    (1, 34, 122)
+])
+def test_draw_boxes_colors(colors):
+    img = torch.full((3, 100, 100), 0, dtype=torch.uint8)
+    utils.draw_bounding_boxes(img, boxes, fill=False, width=7, colors=colors)
+
+
 def test_draw_boxes_vanilla():
     img = torch.full((3, 100, 100), 0, dtype=torch.uint8)
     img_cp = img.clone()

torchvision/utils.py

@@ -141,7 +141,7 @@ def draw_bounding_boxes(
     image: torch.Tensor,
     boxes: torch.Tensor,
     labels: Optional[List[str]] = None,
-    colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None,
+    colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None,
     fill: Optional[bool] = False,
     width: int = 1,
     font: Optional[str] = None,
@@ -159,8 +159,9 @@ def draw_bounding_boxes(
             the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and
             `0 <= ymin < ymax < H`.
         labels (List[str]): List containing the labels of bounding boxes.
-        colors (List[Union[str, Tuple[int, int, int]]]): List containing the colors of bounding boxes. The colors can
-            be represented as `str` or `Tuple[int, int, int]`.
+        colors (Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]): List containing the colors
+            or a single color for all of the bounding boxes. The colors can be represented as `str` or
+            `Tuple[int, int, int]`.
         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
@@ -200,8 +201,10 @@ def draw_bounding_boxes(
     for i, bbox in enumerate(img_boxes):
         if colors is None:
             color = None
-        else:
+        elif isinstance(colors, list):
             color = colors[i]
+        else:
+            color = colors
 
         if fill:
             if color is None: