Skip to content

Commit 0153fca

Browse files
author
pytorchbot
committed
2025-04-17 nightly release (f799a53)
1 parent b82cf99 commit 0153fca

File tree

6 files changed

+46
-14
lines changed

6 files changed

+46
-14
lines changed
680 Bytes
Loading

test/test_ops.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1384,6 +1384,10 @@ def test_bbox_cxcywhr_to_xyxyxyxy(self):
13841384
box_xyxyxyxy = ops.box_convert(box_tensor, in_fmt="cxcywhr", out_fmt="xyxyxyxy")
13851385
torch.testing.assert_close(box_xyxyxyxy, exp_xyxyxyxy)
13861386

1387+
# Reverse conversion
1388+
box_cxcywhr = ops.box_convert(box_xyxyxyxy, in_fmt="xyxyxyxy", out_fmt="cxcywhr")
1389+
torch.testing.assert_close(box_cxcywhr, box_tensor)
1390+
13871391
def test_bbox_xywhr_to_xyxyxyxy(self):
13881392
box_tensor = torch.tensor([[4, 5, 4, 2, 90]], dtype=torch.float)
13891393
exp_xyxyxyxy = torch.tensor([[4, 5, 4, 1, 6, 1, 6, 5]], dtype=torch.float)
@@ -1392,6 +1396,10 @@ def test_bbox_xywhr_to_xyxyxyxy(self):
13921396
box_xyxyxyxy = ops.box_convert(box_tensor, in_fmt="xywhr", out_fmt="xyxyxyxy")
13931397
torch.testing.assert_close(box_xyxyxyxy, exp_xyxyxyxy)
13941398

1399+
# Reverse conversion
1400+
box_xywhr = ops.box_convert(box_xyxyxyxy, in_fmt="xyxyxyxy", out_fmt="xywhr")
1401+
torch.testing.assert_close(box_xywhr, box_tensor)
1402+
13951403
@pytest.mark.parametrize("inv_infmt", ["xwyh", "cxwyh", "xwyhr", "cxwyhr", "xxxxyyyy"])
13961404
@pytest.mark.parametrize("inv_outfmt", ["xwcx", "xhwcy", "xwcxr", "xhwcyr", "xyxyxxyy"])
13971405
def test_bbox_invalid(self, inv_infmt, inv_outfmt):

test/test_utils.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,23 @@ def test_draw_boxes_with_coloured_labels():
131131
assert_equal(result, expected)
132132

133133

134+
@pytest.mark.skipif(PILLOW_VERSION < (10, 1), reason="The reference image is only valid for PIL >= 10.1")
135+
def test_draw_boxes_with_coloured_label_backgrounds():
136+
img = torch.full((3, 100, 100), 255, dtype=torch.uint8)
137+
labels = ["a", "b", "c", "d"]
138+
colors = ["green", "#FF00FF", (0, 255, 0), "red"]
139+
label_colors = ["green", "red", (0, 255, 0), "#FF00FF"]
140+
result = utils.draw_bounding_boxes(
141+
img, boxes, labels=labels, colors=colors, fill=True, label_colors=label_colors, fill_labels=True
142+
)
143+
144+
path = os.path.join(
145+
os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_different_label_fill_colors.png"
146+
)
147+
expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
148+
assert_equal(result, expected)
149+
150+
134151
@pytest.mark.parametrize("fill", [True, False])
135152
def test_draw_boxes_dtypes(fill):
136153
img_uint8 = torch.full((3, 100, 100), 255, dtype=torch.uint8)

torchvision/ops/_box_convert.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -178,10 +178,9 @@ def _box_xyxyxyxy_to_xywhr(boxes: Tensor) -> Tensor:
178178
x1, y1, x3, y3, x2, y2, x4, y4 = boxes.unbind(-1)
179179
r_rad = torch.atan2(y1 - y3, x3 - x1)
180180
r = r_rad * 180 / torch.pi
181-
cos, sin = torch.cos(r_rad), torch.sin(r_rad)
182181

183-
w = (x2 - x1) * cos + (y1 - y2) * sin
184-
h = (x2 - x1) * sin + (y2 - y1) * cos
182+
w = ((x3 - x1) ** 2 + (y1 - y3) ** 2).sqrt()
183+
h = ((x3 - x2) ** 2 + (y3 - y2) ** 2).sqrt()
185184

186185
boxes = torch.stack((x1, y1, w, h, r), dim=-1)
187186

torchvision/transforms/v2/functional/_meta.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -252,13 +252,13 @@ def _xyxyxyxy_to_xywhr(xyxyxyxy: torch.Tensor, inplace: bool) -> torch.Tensor:
252252
xyxyxyxy = xyxyxyxy.float()
253253

254254
r_rad = torch.atan2(xyxyxyxy[..., 1].sub(xyxyxyxy[..., 3]), xyxyxyxy[..., 2].sub(xyxyxyxy[..., 0]))
255-
cos, sin = r_rad.cos(), r_rad.sin()
256-
# x1, y1, x3, y3, (x2 - x1), (y2 - y1) x4, y4
257-
xyxyxyxy[..., 4:6].sub_(xyxyxyxy[..., :2])
258-
# (x2 - x1) * cos + (y1 - y2) * sin = w
259-
xyxyxyxy[..., 2] = xyxyxyxy[..., 4].mul(cos).sub(xyxyxyxy[..., 5].mul(sin))
260-
# (x2 - x1) * sin + (y2 - y1) * cos = h
261-
xyxyxyxy[..., 3] = xyxyxyxy[..., 5].mul(cos).add(xyxyxyxy[..., 4].mul(sin))
255+
# x1, y1, (x3 - x1), (y3 - y1), (x2 - x3), (y2 - y3) x4, y4
256+
xyxyxyxy[..., 4:6].sub_(xyxyxyxy[..., 2:4])
257+
xyxyxyxy[..., 2:4].sub_(xyxyxyxy[..., :2])
258+
# sqrt((x3 - x1) ** 2 + (y1 - y3) ** 2) = w
259+
xyxyxyxy[..., 2] = xyxyxyxy[..., 2].pow(2).add(xyxyxyxy[..., 3].pow(2)).sqrt()
260+
# sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2) = h
261+
xyxyxyxy[..., 3] = xyxyxyxy[..., 4].pow(2).add(xyxyxyxy[..., 5].pow(2)).sqrt()
262262
xyxyxyxy[..., 4] = r_rad.div_(torch.pi).mul_(180.0)
263263
return xyxyxyxy[..., :5].to(dtype)
264264

torchvision/utils.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ def draw_bounding_boxes(
162162
font: Optional[str] = None,
163163
font_size: Optional[int] = None,
164164
label_colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None,
165+
fill_labels: bool = False,
165166
) -> torch.Tensor:
166167

167168
"""
@@ -186,7 +187,8 @@ def draw_bounding_boxes(
186187
`/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
187188
font_size (int): The requested font size in points.
188189
label_colors (color or list of colors, optional): Colors for the label text. See the description of the
189-
`colors` argument for details. Defaults to the same colors used for the boxes.
190+
`colors` argument for details. Defaults to the same colors used for the boxes, or to black if ``fill_labels`` is True.
191+
fill_labels (bool): If `True` fills the label background with specified box color (from the ``colors`` parameter). Default: False.
190192
191193
Returns:
192194
img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted.
@@ -223,8 +225,8 @@ def draw_bounding_boxes(
223225
)
224226

225227
colors = _parse_colors(colors, num_objects=num_boxes)
226-
if label_colors:
227-
label_colors = _parse_colors(label_colors, num_objects=num_boxes) # type: ignore[assignment]
228+
if label_colors or fill_labels:
229+
label_colors = _parse_colors(label_colors if label_colors else "black", num_objects=num_boxes) # type: ignore[assignment]
228230
else:
229231
label_colors = colors.copy() # type: ignore[assignment]
230232

@@ -259,7 +261,13 @@ def draw_bounding_boxes(
259261
draw.rectangle(bbox, width=width, outline=color)
260262

261263
if label is not None:
262-
margin = width + 1
264+
box_margin = 1
265+
margin = width + box_margin
266+
if fill_labels:
267+
left, top, right, bottom = draw.textbbox((bbox[0] + margin, bbox[1] + margin), label, font=txt_font)
268+
draw.rectangle(
269+
(left - box_margin, top - box_margin, right + box_margin, bottom + box_margin), fill=color
270+
)
263271
draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=label_color, font=txt_font) # type: ignore[arg-type]
264272

265273
out = F.pil_to_tensor(img_to_draw)

0 commit comments

Comments
 (0)