使用PyTorch实现目标检测边界框转换与可视化

一、引言

在目标检测任务中，边界框（Bounding Box）的坐标表示与转换是核心基础操作。本文将演示如何：

实现边界框的两种表示形式（角点坐标 vs 中心坐标）之间的转换
使用Matplotlib在图像上可视化边界框
验证坐标转换的正确性

二、环境准备

import torch
from d2l import torch as d2l

PyTorch 1.12+
d2l 0.17.0+
Matplotlib 3.5+

三、代码实现与解析

1. 图像加载与显示

d2l.set_figsize()  # 设置画布大小
img = d2l.plt.imread('./data/catdog.jpg')  # 读取图像
d2l.plt.imshow(img)  # 显示图像

2. 坐标转换函数

角点坐标 → 中心坐标

def box_corner_to_center(boxes):x1, y1, x2, y2 = boxes[:,0], boxes[:,1], boxes[:,2], boxes[:,3]cx = (x1 + x2) / 2  # 中心x坐标cy = (y1 + y2) / 2  # 中心y坐标w = x2 - x1         # 宽度h = y2 - y1         # 高度return torch.stack((cx, cy, w, h), axis=-1)

中心坐标 → 角点坐标

def box_center_to_corner(boxes):cx, cy, w, h = boxes[:,0], boxes[:,1], boxes[:,2], boxes[:,3]x1 = cx - w/2  # 左上角xy1 = cy - h/2  # 左上角yx2 = cx + w/2  # 右下角xy2 = cy + h/2  # 右下角yreturn torch.stack((x1, y1, x2, y2), axis=-1)

3. 转换验证

dog_bbox = [10.0, 50.0, 80.0, 110.0]  # 狗的边界框(左上x,左上y,右下x,右下y)
cat_bbox = [90.0, 40.0, 160.0, 110.0] # 猫的边界框
boxes = torch.tensor([dog_bbox, cat_bbox])# 验证转换的正确性
print(box_center_to_corner(box_corner_to_center(boxes)) == boxes)

输出结果：

tensor([[True, True, True, True],[True, True, True, True]])

4. 边界框可视化

def bbox_to_rect(bbox, color):return d2l.plt.Rectangle(xy=(bbox[0], bbox[1]),  # 左上角坐标width=bbox[2]-bbox[0],  # 宽度height=bbox[3]-bbox[1], # 高度fill=False, edgecolor=color, linewidth=2)fig = d2l.plt.imshow(img)
fig.axes.add_patch(bbox_to_rect(dog_bbox, 'blue'))  # 添加蓝色狗框
fig.axes.add_patch(bbox_to_rect(cat_bbox, 'red'))   # 添加红色猫框

可视化效果：

X轴范围：0-150
Y轴范围：0-125
蓝色矩形框标注狗的位置
红色矩形框标注猫的位置

四、完整代码

import torch
from d2l import torch as d2l# 设置画布并加载图像
d2l.set_figsize()
img = d2l.plt.imread('./data/catdog.jpg')
d2l.plt.imshow(img)# 坐标转换函数
def box_corner_to_center(boxes):x1, y1, x2, y2 = boxes[:,0], boxes[:,1], boxes[:,2], boxes[:,3]cx, cy = (x1+x2)/2, (y1+y2)/2w, h = x2-x1, y2-y1return torch.stack((cx, cy, w, h), axis=-1)def box_center_to_corner(boxes):cx, cy, w, h = boxes[:,0], boxes[:,1], boxes[:,2], boxes[:,3]x1, y1 = cx-w/2, cy-h/2x2, y2 = cx+w/2, cy+h/2return torch.stack((x1, y1, x2, y2), axis=-1)# 定义边界框并验证转换
dog_bbox, cat_bbox = [10.0,50.0,80.0,110.0], [90.0,40.0,160.0,110.0]
boxes = torch.tensor([dog_bbox, cat_bbox])
print(box_center_to_corner(box_corner_to_center(boxes)) == boxes)# 可视化边界框
def bbox_to_rect(bbox, color):return d2l.plt.Rectangle((bbox[0], bbox[1]), bbox[2]-bbox[0], bbox[3]-bbox[1], fill=False, edgecolor=color, linewidth=2)fig = d2l.plt.imshow(img)
fig.axes.add_patch(bbox_to_rect(dog_bbox, 'blue'))
fig.axes.add_patch(bbox_to_rect(cat_bbox, 'red'))