fix color in draw_segmentation_masks

[rizavelioglu]

Reproduce bug:

import torch
import torchvision.utils as utils

num_masks, h, w = 4, 100, 100
img = torch.randint(0, 256, size=(3, h, w), dtype=torch.uint8)
masks = torch.randint(0, 2, (num_masks, h, w), dtype=torch.bool)

utils.draw_segmentation_masks(img, masks, colors="red")

which throws:

ValueError: There are more masks (4) than colors (3)

where the expected outcome is the drawing of all the segmentation masks in a single color (as stated in the docs).

Notice that, in the following settings the function works as intended;

• num_masks=3 and colors="red",
• num_masks=4 and colors="blue",
  meaning, in the case of num_masks=len(colors).

Currently, the respective tests pass even though the function does not work properly. This is because the test is checked in the presence of only 2 masks:

vision/test/test_utils.py

Line 205 in b78d98b

num_masks, h, w = 2, 100, 100

Therefore, for any input whose length is larger than 2, i.e. "red", "#FF00FF", (240, 10, 157), the tests will pass.

The Fix:
I have adapted the code in draw_segmentation_masks such that it also accepts a string of colors. The code was adapted similarly to how it is handled in draw_bounding_boxes, where the issue does not appear.
Specific tests passed with:

pytest test/test_utils.py -vvv -k test_draw_segmentation_masks

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/7520

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm jobs are failing with No space left on device

❌ 25 Failures

As of commit f7fd5d7:

NEW FAILURES - The following jobs have failed:

• build / linux-job (gh)
• c-source / linux-job (gh)
• onnx / linux-job (gh)
• python-source-and-configs / linux-job (gh)
• python-types / linux-job (gh)
• unittests (3.10, linux.12xlarge, cpu) / linux-job (gh)
• unittests (3.10, macos-12) / macos-job (gh)
• unittests (3.10, windows.4xlarge, cpu) / windows-job (gh)
• unittests (3.11, linux.12xlarge, cpu) / linux-job (gh)
• unittests (3.11, macos-12) / macos-job (gh)
• unittests (3.11, windows.4xlarge, cpu) / windows-job (gh)
• unittests (3.8, linux.12xlarge, cpu) / linux-job (gh)
• unittests (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.7) / linux-job (gh)
• unittests (3.8, macos-12) / macos-job (gh)
• unittests (3.8, macos-m1-12) / macos-job (gh)
• unittests (3.8, windows.4xlarge, cpu) / windows-job (gh)
• unittests (3.9, linux.12xlarge, cpu) / linux-job (gh)
• unittests (3.9, macos-12) / macos-job (gh)
• unittests (3.9, windows.4xlarge, cpu) / windows-job (gh)
• unittests-extended / linux-job (gh)
• unittests-prototype (3.10, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.11, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.8, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.7) / linux-job (gh)
• unittests-prototype (3.9, linux.12xlarge, cpu) / linux-job (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pmeier]

Thanks @rizavelioglu for the PR. The idea to align with the bounding box functionality SGTM. IMO, we should have a def _parse_colors(colors, num_objects: int) function that handles this. With that, we don't need to duplicate the parsing in both functions. Since it would be a "generalization" of _generate_color_palette, we could inline that into the new function. Thoughts?

[rizavelioglu]

Thank you @pmeier for your message!
I applied the changes and added the _parse_colors function. but with an additional argument: output_format.
I did so because draw_bounding_boxes expects a list of tuple, whereas draw_segmentation_masks expects a list of torch.tensor. What do you think?

[pmeier]

Answered your question inline. Looking good so far!

[pmeier]

Only minor things left. Thanks a lot @rizavelioglu! LGTM if relevant CI is green.

[pmeier]

Sorry for the delay @rizavelioglu. Two more minor things that I found in a final pass, but they are non-blocking. Let me know if you want to address them. Otherwise we can also merge as is.

[rizavelioglu]

Thanks @pmeier, been waiting for your reply :)
Sure, I will apply the changes.