From 6cb46e4639ef9407abbd8849a7bbd44d2393a103 Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Sat, 2 Oct 2021 22:54:30 +0530 Subject: [PATCH 01/10] Add str param --- torchvision/utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 494661e6ad8..b98e5f8daf8 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -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,8 +242,9 @@ 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)``. + colors: Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]: List containing the colors + of the masks or single color for all 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. From 663fa4c5cd27d15b2651fe3702679a6acbe90d2d Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Sat, 2 Oct 2021 22:54:47 +0530 Subject: [PATCH 02/10] Update test to include str --- test/test_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/test_utils.py b/test/test_utils.py index 37829b906f1..b9ce69d2b94 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -168,6 +168,7 @@ def test_draw_invalid_boxes(): @pytest.mark.parametrize('colors', [ None, + 'blue', ['red', 'blue'], ['#FF00FF', (1, 34, 122)], ]) @@ -197,6 +198,9 @@ def test_draw_segmentation_masks(colors, alpha): if colors is None: colors = utils._generate_color_palette(num_masks) + if isinstance(colors, str): + colors = [colors] + # Make sure each mask draws with its own color for mask, color in zip(masks, colors): if isinstance(color, str): From 2c34a52ad3984c3422d202ae5e3bea73142b20cb Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Sun, 3 Oct 2021 01:07:23 +0530 Subject: [PATCH 03/10] Fix mypy --- torchvision/utils.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index b98e5f8daf8..13f195e07d4 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -289,18 +289,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)] + return [tuple(((i * palette) % 255)) for i in range(num_masks)] From b9da8f3f6108ae39aad5cab9293b6e3ec413f18b Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Sun, 3 Oct 2021 01:14:04 +0530 Subject: [PATCH 04/10] Remove a small bracket --- torchvision/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 13f195e07d4..aa7cc14c3ee 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -302,4 +302,4 @@ def draw_segmentation_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)] + return [tuple((i * palette) % 255) for i in range(num_masks)] From 0482b7aa42c60dad6ba5817ec579a58c084665a1 Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Sun, 3 Oct 2021 17:15:03 +0530 Subject: [PATCH 05/10] Test more robustly --- test/test_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/test_utils.py b/test/test_utils.py index b9ce69d2b94..b15231cc9e5 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -169,6 +169,8 @@ def test_draw_invalid_boxes(): @pytest.mark.parametrize('colors', [ None, 'blue', + '#FF00FF', + (1, 34, 122), ['red', 'blue'], ['#FF00FF', (1, 34, 122)], ]) @@ -198,7 +200,7 @@ def test_draw_segmentation_masks(colors, alpha): if colors is None: colors = utils._generate_color_palette(num_masks) - if isinstance(colors, str): + if isinstance(colors, str) or isinstance(colors, tuple): colors = [colors] # Make sure each mask draws with its own color From dec72707e8553aa31ddd8e73e1bdfde68763a4c5 Mon Sep 17 00:00:00 2001 From: Aditya Oke Date: Mon, 4 Oct 2021 17:43:26 +0530 Subject: [PATCH 06/10] Update docstring and test: --- test/test_utils.py | 3 +-- torchvision/utils.py | 13 ++++++------- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index b15231cc9e5..02cbc91f74a 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -199,8 +199,7 @@ def test_draw_segmentation_masks(colors, alpha): if colors is None: colors = utils._generate_color_palette(num_masks) - - if isinstance(colors, str) or isinstance(colors, tuple): + elif isinstance(colors, str) or isinstance(colors, tuple): colors = [colors] # Make sure each mask draws with its own color diff --git a/torchvision/utils.py b/torchvision/utils.py index aa7cc14c3ee..db880b0aea3 100644 --- a/torchvision/utils.py +++ b/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 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 @@ -242,11 +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: Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]: List containing the colors - of the masks or single color for all masks. The colors can be represented as + colors: (color 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)``. - 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. + By default, random colors are generated for each mask. Returns: img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top. From de943a992783c42bd7cc688d0386664722f02bef Mon Sep 17 00:00:00 2001 From: Aditya Oke <47158509+oke-aditya@users.noreply.github.com> Date: Mon, 4 Oct 2021 17:48:45 +0530 Subject: [PATCH 07/10] Apply suggestions from code review Co-authored-by: Nicolas Hug --- torchvision/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index db880b0aea3..76926669b93 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -159,7 +159,7 @@ 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: (color list of colors, optional): List containing the colors + 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. @@ -242,7 +242,7 @@ 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: (color list of colors, optional): List containing the colors + 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. From 456bc9272f8cc82cde3fc93897240dfc02c7aeab Mon Sep 17 00:00:00 2001 From: Aditya Oke <47158509+oke-aditya@users.noreply.github.com> Date: Tue, 5 Oct 2021 18:09:04 +0530 Subject: [PATCH 08/10] Update torchvision/utils.py Small docstring fix --- torchvision/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 46045a111a6..4e6f570e7cd 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -243,7 +243,7 @@ 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: (color or list of colors, optional): List containing the colors + 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. From 2489feaceb47729d7e761799c317ea20612edf01 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 5 Oct 2021 13:43:28 +0100 Subject: [PATCH 09/10] Update torchvision/utils.py --- torchvision/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 4e6f570e7cd..f6c16f8ec3f 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -160,7 +160,7 @@ 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: (color or list of colors, optional): List containing the colors + 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. From 36683cb6032690f2794ea13571b5aa61b39d2fea Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 5 Oct 2021 14:49:00 +0100 Subject: [PATCH 10/10] remove unnecessary renaming --- torchvision/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 4e6f570e7cd..e54b7463bb6 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -293,8 +293,8 @@ def draw_segmentation_masks( img_to_draw = image.detach().clone() # TODO: There might be a way to vectorize this - for mask, color_tensor in zip(masks, colors_): - img_to_draw[:, mask] = color_tensor[:, None] + for mask, color in zip(masks, colors_): + img_to_draw[:, mask] = color[:, None] out = image * (1 - alpha) + img_to_draw * alpha return out.to(out_dtype)