Add label background (#9018)

Co-authored-by: Nicolas Hug <[email protected]>

Author: AntoineSimoulin

test/assets/fakedata/draw_boxes_different_label_fill_colors.png

@@ -131,6 +131,23 @@ def test_draw_boxes_with_coloured_labels():
     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_backgrounds():
+    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"]
+    result = utils.draw_bounding_boxes(
+        img, boxes, labels=labels, colors=colors, fill=True, label_colors=label_colors, fill_labels=True
+    )
+
+    path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_different_label_fill_colors.png"
+    )
+    expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
+    assert_equal(result, expected)
+
+
 @pytest.mark.parametrize("fill", [True, False])
 def test_draw_boxes_dtypes(fill):
     img_uint8 = torch.full((3, 100, 100), 255, dtype=torch.uint8)

test/test_utils.py

@@ -162,6 +162,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,
+    fill_labels: bool = False,
 ) -> torch.Tensor:
 
     """
@@ -186,7 +187,8 @@ def draw_bounding_boxes(
             `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
         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.
+            `colors` argument for details.  Defaults to the same colors used for the boxes, or to black if ``fill_labels`` is True.
+        fill_labels (bool): If `True` fills the label background with specified box color (from the ``colors`` parameter). Default: False.
 
     Returns:
         img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted.
@@ -223,8 +225,8 @@ def draw_bounding_boxes(
         )
 
     colors = _parse_colors(colors, num_objects=num_boxes)
-    if label_colors:
-        label_colors = _parse_colors(label_colors, num_objects=num_boxes)  # type: ignore[assignment]
+    if label_colors or fill_labels:
+        label_colors = _parse_colors(label_colors if label_colors else "black", num_objects=num_boxes)  # type: ignore[assignment]
     else:
         label_colors = colors.copy()  # type: ignore[assignment]
 
@@ -259,7 +261,13 @@ def draw_bounding_boxes(
             draw.rectangle(bbox, width=width, outline=color)
 
         if label is not None:
-            margin = width + 1
+            box_margin = 1
+            margin = width + box_margin
+            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
+                )
             draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=label_color, font=txt_font)  # type: ignore[arg-type]
 
     out = F.pil_to_tensor(img_to_draw)