Paddle实现单目标检测

单目标检测

单目标检测（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()

输出：