diff --git a/test/test_utils.py b/test/test_utils.py index 1439a974368..b08810895bf 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -122,7 +122,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): diff --git a/torchvision/utils.py b/torchvision/utils.py index 399dc3fcc5a..d45f8c2a7f5 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -199,10 +199,16 @@ def draw_bounding_boxes( txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size) + if colors is None: + if labels is None: + colors = _generate_color_palette(len(img_boxes)) + else: + assert len(labels) == len(img_boxes) + label_color_map = dict(zip(labels, _generate_color_palette(len(labels)))) + 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