Skip to content

Fix IndexError on draw_segmentation_masks and draw_bounding_boxes #5857

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
May 23, 2022
+26 −0
Merged

Fix IndexError on draw_segmentation_masks and draw_bounding_boxes #5857

Changes from 3 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b00b522
Fixing the IndexError in draw_segmentation_masks
oxabz Apr 21, 2022
6e695ad
fixing the bug on draw_bounding_boxes
Apr 21, 2022
b51cbb2
Changing fstring to normal string
Apr 21, 2022
a1b91cb
Removing unecessary conversion
Apr 22, 2022
d9e4ee1
Adding test for the change
Apr 22, 2022
7adff47
Adding a test for draw seqmentation mask
Apr 22, 2022
462ba93
Fixing small mistake
Apr 22, 2022
32a12b5
Fixing an error in the tests
Apr 22, 2022
f527d4e
removing useless imports
Apr 28, 2022
82c11c4
Merge branch 'main' into fix-draw-seg-mask-index-error
NicolasHug May 23, 2022
0467907
ufmt
NicolasHug May 23, 2022
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
12 changes: 10 additions & 2 deletions torchvision/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,10 @@ def draw_bounding_boxes(

num_boxes = boxes.shape[0]

if num_boxes == 0:
warnings.warn("boxes doesn't contain any box. No box was drawn")
return image.to(torch.uint8)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you don't ned to do this as the image dtype is unit8 in the input.


if labels is None:
labels: Union[List[str], List[None]] = [None] * num_boxes # type: ignore[no-redef]
elif len(labels) != num_boxes:
Expand Down Expand Up @@ -306,6 +310,12 @@ def draw_segmentation_masks(
if colors is not None and num_masks > len(colors):
raise ValueError(f"There are more masks ({num_masks}) than colors ({len(colors)})")

out_dtype = torch.uint8
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same. No need of this, only uint8 images are allowed


if num_masks == 0:
warnings.warn("masks doesn't contain any mask. No mask was drawn")
return image.to(out_dtype)

if colors is None:
colors = _generate_color_palette(num_masks)

Expand All @@ -316,8 +326,6 @@ def draw_segmentation_masks(
if isinstance(colors[0], tuple) and len(colors[0]) != 3:
raise ValueError("It seems that you passed a tuple of colors instead of a list of colors")

out_dtype = torch.uint8

colors_ = []
for color in colors:
if isinstance(color, str):
Expand Down