单目标检测
单目标检测(Single Object Detection)是人工智能领域中的一个重要研究方向,旨在通过计算机视觉技术,识别和定位图像中的特定目标物体。单目标检测可以应用于各种场景,如智能监控、自动驾驶、医疗影像分析等。
简单来说,单目标检测就是在确定一个目标在图片中的位置:
本文将以信号灯检测为例,介绍单目标检测的方法
环境准备
这个案例需要安装以下两个库:
pip install paddlepaddle-gpu
pip install lxml
数据集准备
本文采用如下数据集:红绿灯检测_练习_训练集(非比赛数据)_数据集-飞桨AI Studio星河社区 (baidu.com)
这个数据集共有2000张信号灯的照片,其中1000张绿灯,1000张红灯。每张照片都对应着一个xml文件,标注着信号灯在图片中的位置:
<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<annotation>
<folder>Images</folder>
<filename>green_0.jpg</filename>
<source>
<database>Unknown</database>
</source>
<size>
<width>424</width>
<height>240</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>green</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<occluded>0</occluded>
<bndbox>
<xmin>247</xmin>
<ymin>147</ymin>
<xmax>301</xmax>
<ymax>190</ymax>
</bndbox>
</object>
</annotation>
这里面,<width>和<height>标签分别定义了宽和高,<name>定义了样本的类别(red或者green),<bndbox>里的标签则是定义了信号灯的位置(矩形框)
接下来我们编写dataset.py,用于定义数据集类:
import paddle
import glob
from lxml import etree
from PIL import Image
import numpy as np
# 定义一个字典,将颜色名称映射到ID
name_to_id = {'red': 0, 'green': 1}
# 将绝对坐标转换为相对坐标
def to_labels(path):
# 读取XML文件内容
text = open(f'{path}').read().encode('utf8')
# 解析XML内容
xml = etree.HTML(text)
# 提取图像的宽度和高度
width = int(xml.xpath('//size/width/text()')[0])
height = int(xml.xpath('//size/height/text()')[0])
# 提取边界框的坐标
xmin = int(xml.xpath('//bndbox/xmin/text()')[0])
xmax = int(xml.xpath('//bndbox/xmax/text()')[0])
ymin = int(xml.xpath('//bndbox/ymin/text()')[0])
ymax = int(xml.xpath('//bndbox/ymax/text()')[0])
# 将绝对坐标转换为相对坐标
return xmin / width, ymin / height, xmax / width, ymax / height
# 定义一个PaddlePaddle数据集类
class Dataset(paddle.io.Dataset):
def __init__(self, pos='training_data'):
super().__init__() # 调用父类构造函数
# 查找指定目录下的所有.jpg图片和.xml标签文件
self.imgs = glob.glob(f'{pos}/*.jpg')
self.labels = glob.glob(f'{pos}/*.xml')
def __getitem__(self, idx):
# 根据索引获取图片和标签
img = self.imgs[idx]
label = to_labels(self.labels[idx])
# 打开图片并转换为RGB模式
pil_img = Image.open(img).convert('RGB')
# 将PIL图片转换为numpy数组,并转换为float32类型
# 同时将通道顺序从HWC转换为CHW(PaddlePaddle默认输入格式)
t = paddle.to_tensor(np.array(pil_img, dtype=np.float32).transpose((2, 0, 1)))
# 返回图片张量和标签张量
return t, paddle.to_tensor(label[:4])
def __len__(self):
# 返回数据集中图片的数量
return len(self.imgs)
训练脚本
单目标检测可以看作一个回归问题,输出4个值,用于确定目标的坐标,因此我们可以使用resnet,并指定其类别数量为4(即输出4个值),并采用MSE损失函数(因为这是回归问题),据此,可以写出训练脚本的代码:
import paddle
from dataset import Dataset
# 初始化Dataset实例,设置数据位置为'training_data'
dataset = Dataset(pos='training_data')
# 使用ResNet18网络结构,并设置输出类别数为4
net = paddle.vision.resnet18(num_classes=4)
# 将网络封装为PaddlePaddle的Model对象
model = paddle.Model(net)
# 准备模型训练,包括优化器(Adam)和损失函数(均方误差损失)
model.prepare(
paddle.optimizer.Adam(parameters=model.parameters()),
paddle.nn.MSELoss(),
)
# 训练模型,设置训练轮数为160,批处理大小为16
model.fit(dataset, epochs=160, batch_size=16, verbose=1)
# 保存模型到'output/model'路径
model.save('output/model')
可以看到,训练脚本还是非常简单的。
简单使用
使用脚本也很简单:
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from PIL import Image
import paddle
# 图片路径
img_path = 'testing_data/red_1003.jpg'
# 打开图片并转换为RGB格式
pil_img = Image.open(img_path).convert('RGB')
# 将PIL图片转换为Paddle Tensor,并调整通道顺序
t = paddle.to_tensor([np.array(pil_img, dtype=np.float32).transpose((2, 0, 1))])
# 加载ResNet18模型,并设置为4个类别
net = paddle.vision.resnet18(num_classes=4)
model = paddle.Model(net)
# 加载训练好的模型权重
model.load('output/model')
# 预测图片
pred = model.predict_batch(t)[0][0]
print(f'预测结果:{pred}')
# 根据预测结果计算边界框坐标
xmin = float(pred[0]) * 424
ymin = float(pred[1]) * 240
xmax = float(pred[2]) * 424
ymax = float(pred[3]) * 240
# 显示原始图片
plt.imshow(np.array(t[0], dtype=np.int32).transpose((1, 2, 0)))
# 定义多边形的顶点坐标(这里是预测的边界框)
vertices = np.array([[xmin, ymin], [xmin, ymax], [xmax, ymax], [xmax, ymin]])
# 创建一个多边形对象,用于绘制边界框
polygon = patches.Polygon(vertices, closed=True, edgecolor='black', facecolor='none')
# 将多边形添加到当前坐标轴上
plt.gca().add_patch(polygon)
# 显示图片和边界框
plt.show()
输出: