Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion test/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ def test_draw_segmentation_masks_errors(device):
utils.draw_segmentation_masks(image=img, masks=masks_bad_shape)
with pytest.raises(ValueError, match="Number of colors must be equal or larger than the number of objects"):
utils.draw_segmentation_masks(image=img, masks=masks, colors=[])
with pytest.raises(ValueError, match="`colors` must be a tuple or a string, or a list thereof"):
with pytest.raises(ValueError, match="colors must be a tuple or a string, or a list thereof"):
bad_colors = np.array(["red", "blue"]) # should be a list
utils.draw_segmentation_masks(image=img, masks=masks, colors=bad_colors)
with pytest.raises(ValueError, match="If passed as tuple, colors should be an RGB triplet"):
Expand Down
6 changes: 3 additions & 3 deletions torchvision/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -359,14 +359,14 @@ def draw_bounding_boxes(
f"Number of boxes ({num_boxes}) and labels ({len(labels)}) mismatch. Please specify labels for each box."
)

colors = _parse_colors(colors, num_objects=num_boxes)
colors = _parse_colors(colors, num_objects=num_boxes) # type: ignore[assignment]
if label_colors or fill_labels:
label_colors = _parse_colors(label_colors if label_colors else "black", num_objects=num_boxes) # type: ignore[assignment]
else:
label_colors = colors.copy() # type: ignore[assignment]

if fill_labels and label_background_colors:
label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes)
label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes) # type: ignore[assignment]
else:
label_background_colors = colors.copy() # type: ignore[assignment]

Expand Down Expand Up @@ -404,7 +404,7 @@ def draw_bounding_boxes(
if fill_labels:
left, top, right, bottom = draw.textbbox((bbox[0] + margin, bbox[1] + margin), label, font=txt_font)
draw.rectangle(
(left - box_margin, top - box_margin, right + box_margin, bottom + box_margin), fill=label_bg_color
(left - box_margin, top - box_margin, right + box_margin, bottom + box_margin), fill=label_bg_color # type: ignore[arg-type]
)
draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=label_color, font=txt_font) # type: ignore[arg-type]

Expand Down
Loading