Fix annotation of draw_segmentation_masks

test/test_utils.py

@@ -168,6 +168,9 @@ def test_draw_invalid_boxes():
 
 @pytest.mark.parametrize('colors', [
     None,
+    'blue',
+    '#FF00FF',
+    (1, 34, 122),
     ['red', 'blue'],
     ['#FF00FF', (1, 34, 122)],
 ])
@@ -196,6 +199,8 @@ def test_draw_segmentation_masks(colors, alpha):
 
     if colors is None:
         colors = utils._generate_color_palette(num_masks)
+    elif isinstance(colors, str) or isinstance(colors, tuple):
+        colors = [colors]
 
     # Make sure each mask draws with its own color
     for mask, color in zip(masks, colors):

torchvision/utils.py

@@ -159,9 +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 (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]`.
+        colors: (color or list of colors, optional): List containing the colors
+            of the boxes or single color for all boxes. The color can be represented as
+            PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
         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
@@ -230,7 +230,7 @@ def draw_segmentation_masks(
     image: torch.Tensor,
     masks: torch.Tensor,
     alpha: float = 0.8,
-    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,
 ) -> torch.Tensor:
 
     """
@@ -242,10 +242,10 @@ def draw_segmentation_masks(
         masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool.
         alpha (float): Float number between 0 and 1 denoting the transparency of the masks.
             0 means full transparency, 1 means no transparency.
-        colors (list or None): List containing the colors of the masks. The colors can
-            be represented as PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
-            When ``masks`` has a single entry of shape (H, W), you can pass a single color instead of a list
-            with one element. By default, random colors are generated for each mask.
+        colors: (color or list of colors, optional): List containing the colors
+            of the masks or single color for all masks. The color can be represented as
+            PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+            By default, random colors are generated for each mask.
 
     Returns:
         img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top.
@@ -288,18 +288,17 @@ def draw_segmentation_masks(
     for color in colors:
         if isinstance(color, str):
             color = ImageColor.getrgb(color)
-        color = torch.tensor(color, dtype=out_dtype)
-        colors_.append(color)
+        colors_.append(torch.tensor(color, dtype=out_dtype))
 
     img_to_draw = image.detach().clone()
     # TODO: There might be a way to vectorize this
-    for mask, color in zip(masks, colors_):
-        img_to_draw[:, mask] = color[:, None]
+    for mask, color_tensor in zip(masks, colors_):
+        img_to_draw[:, mask] = color_tensor[:, None]
 
     out = image * (1 - alpha) + img_to_draw * alpha
     return out.to(out_dtype)
 
 
-def _generate_color_palette(num_masks):
+def _generate_color_palette(num_masks: int):
     palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
     return [tuple((i * palette) % 255) for i in range(num_masks)]