From 1e6682a58b4fc2054735f5cde83cd3ca58cca580 Mon Sep 17 00:00:00 2001 From: AndreiMoraru123 Date: Wed, 3 Sep 2025 20:45:17 +0300 Subject: [PATCH 1/5] add label background color parameter --- ...oxes_different_label_background_colors.png | Bin 0 -> 766 bytes test/test_utils.py | 24 ++++++++++++++++++ torchvision/utils.py | 16 +++++++++--- 3 files changed, 36 insertions(+), 4 deletions(-) create mode 100644 test/assets/fakedata/draw_boxes_different_label_background_colors.png diff --git a/test/assets/fakedata/draw_boxes_different_label_background_colors.png b/test/assets/fakedata/draw_boxes_different_label_background_colors.png new file mode 100644 index 0000000000000000000000000000000000000000..bf641e9a122d24359037ea4610f152591e706576 GIT binary patch literal 766 zcmeAS@N?(olHy`uVBq!ia0vp^DImG&K^cmER|+ojS4Tx8UVd*EX%3Z(W?Euy=;v_a`aJ zD@zO(M&CHUWtO?&ggNsp9549?`|S#gx?wA2V%`zCG2oA~*%$vc>>uQ|J#Yy=%z0bm zn&S_S(;pJ>-ESciTjtQh&{KTQ&RL@)XL?_-S+e?S}YW0P=(iZ<_s}|@qCUYizDsQ*8vAie$+ha@Mw-gJTGu2icpdoPUA3CCVd^fP>xrUr2Zeg-Gw1LxO`6+scGhu2;XN9nBIno06>pz6 z>xi-Ug9xVMLKf2wY-9n=1xJHsoYd240@`OWzKY|{5ItM*HQ^j#UK{D>(Lay`D zJzB|g+r=67-2AL~$ZVz2wMX+NZ!4M_mAZz%fV1IrU+4?5J1x2vS1Y4!L-ZVP z|NF3tdkGbDRafMsZGARRtvwy&wmwL3$sjqG20~o>x|e;)y!B)L>R06#eQNK8eM;2; cN7FxsMb{GVT(l~Z2c~NVPgg&ebxsLQ04E`1g8%>k literal 0 HcmV?d00001 diff --git a/test/test_utils.py b/test/test_utils.py index 8b6f357ce6e..223ccd524a7 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -164,6 +164,30 @@ def test_draw_boxes_with_coloured_label_backgrounds(): ) expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1) assert_equal(result, expected) + + +@pytest.mark.skipif(PILLOW_VERSION < (10, 1), reason="The reference image is only valid for PIL >= 10.1") +def test_draw_boxes_with_coloured_label_text_boxes(): + img = torch.full((3, 100, 100), 255, dtype=torch.uint8) + labels = ["a", "b", "c", "d"] + colors = ["green", "#FF00FF", (0, 255, 0), "red"] + label_colors = ["green", "red", (0, 255, 0), "#FF00FF"] + label_background_colors = ["white", "black", "yellow", "blue"] + result = utils.draw_bounding_boxes( + img, + boxes, + labels=labels, + colors=colors, + fill=True, + label_colors=label_colors, + label_background_colors=label_background_colors, + fill_labels=True + ) + path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_different_label_background_colors.png" + ) + expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1) + assert_equal(result, expected) @pytest.mark.skipif(PILLOW_VERSION < (10, 1), reason="The reference image is only valid for PIL >= 10.1") diff --git a/torchvision/utils.py b/torchvision/utils.py index 0d819ef8330..ccf446fe5b7 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -8,8 +8,8 @@ import numpy as np import torch -from PIL import __version__ as PILLOW_VERSION_STRING, Image, ImageColor, ImageDraw, ImageFont - +from PIL import Image, ImageColor, ImageDraw, ImageFont +from PIL import __version__ as PILLOW_VERSION_STRING __all__ = [ "_Image_fromarray", @@ -293,6 +293,7 @@ def draw_bounding_boxes( font: Optional[str] = None, font_size: Optional[int] = None, label_colors: Optional[Union[list[Union[str, tuple[int, int, int]]], str, tuple[int, int, int]]] = None, + label_background_colors: Optional[Union[list[Union[str, tuple[int, int, int]]], str, tuple[int, int, int]]] = None, fill_labels: bool = False, ) -> torch.Tensor: """ @@ -320,6 +321,8 @@ def draw_bounding_boxes( font_size (int): The requested font size in points. label_colors (color or list of colors, optional): Colors for the label text. See the description of the `colors` argument for details. Defaults to the same colors used for the boxes, or to black if ``fill_labels`` is True. + label_background_colors (color or list of colors, optional): Colors for the label text box fill. Defaults to the + same colors used for the boxes. Ignored when ``fill_labels`` is False. fill_labels (bool): If `True` fills the label background with specified box color (from the ``colors`` parameter). Default: False. Returns: @@ -362,6 +365,11 @@ def draw_bounding_boxes( else: label_colors = colors.copy() # type: ignore[assignment] + if fill_labels: + label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes) if label_background_colors else colors.copy() # type: ignore[assignment] + else: + label_background_colors = colors.copy() # type: ignore[assignment] + if font is None: if font_size is not None: warnings.warn("Argument 'font_size' will be ignored since 'font' is not set.") @@ -385,7 +393,7 @@ def draw_bounding_boxes( else: draw = _ImageDrawTV(img_to_draw) - for bbox, color, label, label_color in zip(img_boxes, colors, labels, label_colors): # type: ignore[arg-type] + for bbox, color, label, label_color, label_bg_color in zip(img_boxes, colors, labels, label_colors, label_background_colors): # type: ignore[arg-type] draw_method = draw.oriented_rectangle if len(bbox) > 4 else draw.rectangle fill_color = color + (100,) if fill else None draw_method(bbox, width=width, outline=color, fill=fill_color) @@ -396,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=color + (left - box_margin, top - box_margin, right + box_margin, bottom + box_margin), fill=label_bg_color ) draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=label_color, font=txt_font) # type: ignore[arg-type] From 501f757e1b0cc78936e26ad035c2ca0a95f163fa Mon Sep 17 00:00:00 2001 From: AndreiMoraru123 Date: Thu, 4 Sep 2025 19:10:45 +0300 Subject: [PATCH 2/5] format code with ufmt --- test/test_utils.py | 21 ++++++++++++--------- torchvision/utils.py | 5 ++--- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index 223ccd524a7..afdf6738d2c 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -164,7 +164,7 @@ def test_draw_boxes_with_coloured_label_backgrounds(): ) expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1) assert_equal(result, expected) - + @pytest.mark.skipif(PILLOW_VERSION < (10, 1), reason="The reference image is only valid for PIL >= 10.1") def test_draw_boxes_with_coloured_label_text_boxes(): @@ -172,19 +172,22 @@ def test_draw_boxes_with_coloured_label_text_boxes(): labels = ["a", "b", "c", "d"] colors = ["green", "#FF00FF", (0, 255, 0), "red"] label_colors = ["green", "red", (0, 255, 0), "#FF00FF"] - label_background_colors = ["white", "black", "yellow", "blue"] + label_background_colors = ["white", "black", "yellow", "blue"] result = utils.draw_bounding_boxes( - img, - boxes, - labels=labels, - colors=colors, - fill=True, + img, + boxes, + labels=labels, + colors=colors, + fill=True, label_colors=label_colors, label_background_colors=label_background_colors, - fill_labels=True + fill_labels=True, ) path = os.path.join( - os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_different_label_background_colors.png" + os.path.dirname(os.path.abspath(__file__)), + "assets", + "fakedata", + "draw_boxes_different_label_background_colors.png", ) expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1) assert_equal(result, expected) diff --git a/torchvision/utils.py b/torchvision/utils.py index ccf446fe5b7..e8c7c273b6e 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -8,8 +8,7 @@ import numpy as np import torch -from PIL import Image, ImageColor, ImageDraw, ImageFont -from PIL import __version__ as PILLOW_VERSION_STRING +from PIL import __version__ as PILLOW_VERSION_STRING, Image, ImageColor, ImageDraw, ImageFont __all__ = [ "_Image_fromarray", @@ -366,7 +365,7 @@ def draw_bounding_boxes( label_colors = colors.copy() # type: ignore[assignment] if fill_labels: - label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes) if label_background_colors else colors.copy() # type: ignore[assignment] + label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes) if label_background_colors else colors.copy() # type: ignore[assignment] else: label_background_colors = colors.copy() # type: ignore[assignment] From 2c4e798cc239c9968c94cdcd83c7b0d465f6c4e2 Mon Sep 17 00:00:00 2001 From: AndreiMoraru123 Date: Thu, 4 Sep 2025 21:00:15 +0300 Subject: [PATCH 3/5] lint with flake8 --- torchvision/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index e8c7c273b6e..443c7963e7a 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -552,7 +552,7 @@ def draw_keypoints( if visibility.shape != keypoints.shape[:-1]: raise ValueError( "keypoints and visibility must have the same dimensionality for num_instances and K. " - f"Got {visibility.shape = } and {keypoints.shape = }" + f"Got {visibility.shape=} and {keypoints.shape=}" ) original_dtype = image.dtype @@ -753,7 +753,7 @@ def _parse_colors( f"Number of colors must be equal or larger than the number of objects, but got {len(colors)} < {num_objects}." ) elif not isinstance(colors, (tuple, str)): - raise ValueError(f"`colors` must be a tuple or a string, or a list thereof, but got {colors}.") + raise ValueError(f"colors must be a tuple or a string, or a list thereof, but got {colors}.") elif isinstance(colors, tuple) and len(colors) != 3: raise ValueError(f"If passed as tuple, colors should be an RGB triplet, but got {colors}.") else: # colors specifies a single color for all objects From 3503a636f0dfe84780859b08e66485a7a2c42209 Mon Sep 17 00:00:00 2001 From: AndreiMoraru123 Date: Fri, 5 Sep 2025 11:51:34 +0300 Subject: [PATCH 4/5] update docstring for 'fill_labels' --- torchvision/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 443c7963e7a..59d2374ec8f 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -322,7 +322,8 @@ def draw_bounding_boxes( `colors` argument for details. Defaults to the same colors used for the boxes, or to black if ``fill_labels`` is True. label_background_colors (color or list of colors, optional): Colors for the label text box fill. Defaults to the same colors used for the boxes. Ignored when ``fill_labels`` is False. - fill_labels (bool): If `True` fills the label background with specified box color (from the ``colors`` parameter). Default: False. + fill_labels (bool): If `True` fills the label background with specified color (from the ``label_background_colors`` parameter, + or from the ``colors`` parameter if not specified). Default: False. Returns: img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted. From 19ae698d79bc6527fb82b8ab0d5e12a10e4fb16f Mon Sep 17 00:00:00 2001 From: AndreiMoraru123 Date: Fri, 5 Sep 2025 17:07:11 +0300 Subject: [PATCH 5/5] short circuit if stmt --- torchvision/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 59d2374ec8f..20534dec2f6 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -365,8 +365,8 @@ def draw_bounding_boxes( else: label_colors = colors.copy() # type: ignore[assignment] - if fill_labels: - label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes) if label_background_colors else colors.copy() # type: ignore[assignment] + if fill_labels and label_background_colors: + label_background_colors = _parse_colors(label_background_colors, num_objects=num_boxes) else: label_background_colors = colors.copy() # type: ignore[assignment]