基于paddlepaddle的yolo基本实现

基于paddlepaddle的yolo基本实现

引言

在这篇博客中,我们将深入探讨如何使用PaddlePaddle来实现YOLO(You Only Look Once)模型。YOLO是一种流行的实时目标检测算法,它以其速度和准确性而闻名。我们将使用ResNet18作为骨干网络,并一步步构建整个模型。

数据集:https://aistudio.baidu.com/datasetdetail/94809

构建骨干网络:ResNet18

首先,我们从构建骨干网络ResNet18开始。ResNet(残差网络)通过引入残差学习来解决深层网络中的退化问题。在这个模型中,我们使用了多个卷积层、批归一化(Batch Normalization)、ReLU激活函数和下采样来构建网络。每一层的细节如下所示:

  • 初始卷积层和池化:这一层使用了一个大的卷积核(7x7)和步长为2,以及一个最大池化层,以减小特征图的尺寸并提取初始特征。
  • 残差块:ResNet的核心是残差块,它允许信息直接从早期层传递到后期层。在这个模型中,我们有多个残差块,每个块包含两个3x3卷积层,后跟批归一化和ReLU激活。
  • 下采样:在某些残差块之后,我们使用步长为2的卷积进行下采样,以减少特征图的尺寸并增加深度。
import paddle
import paddle.nn as nn


# 定义一个名为ResNet18的自定义神经网络类,继承自nn.Layer
class ResNet18(nn.Layer):
    def __init__(self, in_channels=3):
        super().__init__()

        # 第一层卷积层,输入通道数为in_channels,输出通道数为64,卷积核大小为7x7,步长为2,填充为3
        self.conv1 = nn.Conv2D(in_channels=in_channels, out_channels=64, kernel_size=7, stride=2, padding=3)
        # 最大池化层,池化核大小为3x3,步长为2,填充为1
        self.maxpool = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)

        # 定义第2层的第1个卷积层,输入通道数为64,输出通道数为64,卷积核大小为3x3,步长为1,填充为1
        self.conv2_1 = nn.Conv2D(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.norm2_1 = nn.BatchNorm2D(num_features=64)  # 批量归一化层
        self.relu2_1 = nn.ReLU()  # ReLU激活函数

        # 定义第2层的第2个卷积层,输入通道数为64,输出通道数为64,卷积核大小为3x3,步长为1,填充为1
        self.conv2_2 = nn.Conv2D(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.norm2_2 = nn.BatchNorm2D(num_features=64)
        self.relu2_2 = nn.ReLU()

        # 定义第3层的第1个卷积层,输入通道数为64,输出通道数为64,卷积核大小为3x3,步长为1,填充为1
        self.conv3_1 = nn.Conv2D(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.norm3_1 = nn.BatchNorm2D(num_features=64)
        self.relu3_1 = nn.ReLU()

        # 定义第3层的第2个卷积层,输入通道数为64,输出通道数为64,卷积核大小为3x3,步长为1,填充为1
        self.conv3_2 = nn.Conv2D(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.norm3_2 = nn.BatchNorm2D(num_features=64)
        self.relu3_2 = nn.ReLU()

        # 定义第4层的第1个卷积层,输入通道数为64,输出通道数为128,卷积核大小为3x3,步长为2,填充为1
        self.conv4_1 = nn.Conv2D(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1)
        self.norm4_1 = nn.BatchNorm2D(num_features=128)
        self.relu4_1 = nn.ReLU()

        # 定义第4层的第2个卷积层,输入通道数为128,输出通道数为128,卷积核大小为3x3,步长为1,填充为1
        self.conv4_2 = nn.Conv2D(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)
        self.norm4_2 = nn.BatchNorm2D(num_features=128)
        self.relu4_2 = nn.ReLU()

        # 下采样操作,将第3层的特征图尺寸减半,用于与第4层的特征图相加
        self.downsample3_4 = nn.Conv2D(in_channels=64, out_channels=128, kernel_size=1, stride=2, padding=0)

        # 定义第5层的第1个卷积层,输入通道数为128,输出通道数为128,卷积核大小为3x3,步长为1,填充为1
        self.conv5_1 = nn.Conv2D(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)
        self.norm5_1 = nn.BatchNorm2D(num_features=128)
        self.relu5_1 = nn.ReLU()

        # 定义第5层的第2个卷积层,输入通道数为128,输出通道数为128,卷积核大小为3x3,步长为1,填充为1
        self.conv5_2 = nn.Conv2D(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)
        self.norm5_2 = nn.BatchNorm2D(num_features=128)
        self.relu5_2 = nn.ReLU()

        # 定义第6层的第1个卷积层,输入通道数为128,输出通道数为256,卷积核大小为3x3,步长为2,填充为1
        self.conv6_1 = nn.Conv2D(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
        self.norm6_1 = nn.BatchNorm2D(num_features=256)
        self.relu6_1 = nn.ReLU()

        # 定义第6层的第2个卷积层,输入通道数为256,输出通道数为256,卷积核大小为3x3,步长为1,填充为1
        self.conv6_2 = nn.Conv2D(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
        self.norm6_2 = nn.BatchNorm2D(num_features=256)
        self.relu6_2 = nn.ReLU()

        # 下采样操作,将第5层的特征图尺寸减半,用于与第6层的特征图相加
        self.downsample5_6 = nn.Conv2D(in_channels=128, out_channels=256, kernel_size=1, stride=2, padding=0)

        # 定义第7层的第1个卷积层,输入通道数为256,输出通道数为256,卷积核大小为3x3,步长为1,填充为1
        self.conv7_1 = nn.Conv2D(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
        self.norm7_1 = nn.BatchNorm2D(num_features=256)
        self.relu7_1 = nn.ReLU()

        # 定义第7层的第2个卷积层,输入通道数为256,输出通道数为256,卷积核大小为3x3,步长为1,填充为1
        self.conv7_2 = nn.Conv2D(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
        self.norm7_2 = nn.BatchNorm2D(num_features=256)
        self.relu7_2 = nn.ReLU()

        # 定义第8层的第1个卷积层,输入通道数为256,输出通道数为512,卷积核大小为3x3,步长为2,填充为1
        self.conv8_1 = nn.Conv2D(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1)
        self.norm8_1 = nn.BatchNorm2D(num_features=512)
        self.relu8_1 = nn.ReLU()

        # 定义第8层的第2个卷积层,输入通道数为512,输出通道数为512,卷积核大小为3x3,步长为1,填充为1
        self.conv8_2 = nn.Conv2D(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
        self.norm8_2 = nn.BatchNorm2D(num_features=512)
        self.relu8_2 = nn.ReLU()

        # 下采样操作,将第7层的特征图尺寸减半,用于与第8层的特征图相加
        self.downsample7_8 = nn.Conv2D(in_channels=256, out_channels=512, kernel_size=1, stride=2, padding=0)

        # 定义第9层的第1个卷积层,输入通道数为512,输出通道数为512,卷积核大小为3x3,步长为1,填充为1
        self.conv9_1 = nn.Conv2D(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
        self.norm9_1 = nn.BatchNorm2D(num_features=512)
        self.relu9_1 = nn.ReLU()

        # 定义第9层的第2个卷积层,输入通道数为512,输出通道数为512,卷积核大小为3x3,步长为1,填充为1
        self.conv9_2 = nn.Conv2D(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
        self.norm9_2 = nn.BatchNorm2D(num_features=512)
        self.relu9_2 = nn.ReLU()

    # 定义前向传播方法,接受输入x
    def forward(self, x):
        x = self.conv1(x)  # 第1层卷积
        x = self.maxpool(x)  # 最大池化

        h = x  # 将当前特征图保存在h中,用于后续的跳跃连接

        x = self.conv2_1(x)  # 第2层的第1个卷积
        x = self.norm2_1(x)  # 批量归一化
        x = self.relu2_1(x)  # ReLU激活

        x = self.conv2_2(x)  # 第2层的第2个卷积
        x = self.norm2_2(x)  # 批量归一化
        x = self.relu2_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv3_1(x)  # 第3层的第1个卷积
        x = self.norm3_1(x)  # 批量归一化
        x = self.relu3_1(x)  # ReLU激活

        x = self.conv3_2(x)  # 第3层的第2个卷积
        x = self.norm3_2(x)  # 批量归一化
        x = self.relu3_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv4_1(x)  # 第4层的第1个卷积
        x = self.norm4_1(x)  # 批量归一化
        x = self.relu4_1(x)  # ReLU激活

        x = self.conv4_2(x)  # 第4层的第2个卷积
        x = self.norm4_2(x)  # 批量归一化
        h = self.downsample3_4(h)  # 第3层到第4层的下采样
        x = self.relu4_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv5_1(x)  # 第5层的第1个卷积
        x = self.norm5_1(x)  # 批量归一化
        x = self.relu5_1(x)  # ReLU激活

        x = self.conv5_2(x)  # 第5层的第2个卷积
        x = self.norm5_2(x)  # 批量归一化
        x = self.relu5_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv6_1(x)  # 第6层的第1个卷积
        x = self.norm6_1(x)  # 批量归一化
        x = self.relu6_1(x)  # ReLU激活

        x = self.conv6_2(x)  # 第6层的第2个卷积
        x = self.norm6_2(x)  # 批量归一化
        h = self.downsample5_6(h)  # 第5层到第6层的下采样
        x = self.relu6_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv7_1(x)  # 第7层的第1个卷积
        x = self.norm7_1(x)  # 批量归一化
        x = self.relu7_1(x)  # ReLU激活

        x = self.conv7_2(x)  # 第7层的第2个卷积
        x = self.norm7_2(x)  # 批量归一化
        x = self.relu7_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv8_1(x)  # 第8层的第1个卷积
        x = self.norm8_1(x)  # 批量归一化
        x = self.relu8_1(x)  # ReLU激活

        x = self.conv8_2(x)  # 第8层的第2个卷积
        x = self.norm8_2(x)  # 批量归一化
        h = self.downsample7_8(h)  # 第7层到第8层的下采样
        x = self.relu8_2(x + h)  # 加上跳跃连接并经过ReLU激活

        h = x  # 将当前特征图保存在h中

        x = self.conv9_1(x)  # 第9层的第1个卷积
        x = self.norm9_1(x)  # 批量归一化
        x = self.relu9_1(x)  # ReLU激活

        x = self.conv9_2(x)  # 第9层的第2个卷积
        x = self.norm9_2(x)  # 批量归一化
        x = self.relu9_2(x + h)  # 加上跳跃连接并经过ReLU激活

        return x  # 返回最终的特征图作为网络的输出

YOLO模型的实现

YOLO模型的核心思想是将目标检测问题转换为单个回归问题。这意味着模型直接在图片上预测边界框和类别概率。

  • YOLO层:我们在ResNet18的基础上添加了一个YOLO层。这个层包含一个1x1的卷积,用于将深层特征图转换为预测向量。
  • 预测向量:预测向量包含每个网格单元的偏移量、尺寸、置信度和类别概率。
import paddle
import paddle.nn as nn


# 定义一个名为YOLO的自定义神经网络类,继承自nn.Layer
class YOLO(nn.Layer):
    def __init__(self, backbone, channels=512, num_classes=1):
        super().__init__()

        # YOLO模型的主干网络,通常是一个预训练的卷积神经网络,用于特征提取
        self.backbone = backbone

        # 用于预测目标框的卷积层,输入通道数为channels,输出通道数为4(目标框的位置信息) + 1(目标存在的置信度) + num_classes(目标的类别数量)
        self.conv = nn.Conv2D(in_channels=channels, out_channels=4 + 1 + num_classes, kernel_size=1, stride=1,
                              padding=0)

        # 用于将预测的目标框的位置信息中的xy坐标映射到[0, 1]的范围,以表示相对于图像的位置
        self.sigmoid = nn.Sigmoid()

        # 用于确保目标框的宽度和高度始终为正数
        self.relu = nn.ReLU()

    # 定义前向传播方法,接受输入x
    def forward(self, x):
        x = self.backbone(x)  # 通过主干网络提取特征图

        x = self.conv(x)  # 使用卷积层进行目标框的预测

        # 提取目标框的位置信息中的xy坐标,并将其映射到[0, 1]的范围
        offset_xy = self.sigmoid(x[:, :2, :, :])

        # 提取目标框的宽度和高度信息,并确保始终为正数
        wh = self.relu(x[:, 2:4, :, :])

        # 提取目标存在的置信度信息,映射到[0, 1]的范围
        confidence = self.sigmoid(x[:, 4:5, :, :])

        # 提取目标的类别信息,映射到[0, 1]的范围
        classes = self.sigmoid(x[:, 5:, :, :])

        # 返回预测的目标框信息:位置偏移、宽高、置信度和类别概率
        return offset_xy, wh, confidence, classes

数据集处理

import os

data = os.listdir("data/images")
# print(data)

# 划分训练集和测试集
train_data = data[:int(len(data) * 0.8)]
test_data = data[int(len(data) * 0.8):]

# 如果已经存在train.txt和test.txt,先删除
if os.path.exists("train.txt"):
    os.remove("train.txt")
if os.path.exists("test.txt"):
    os.remove("test.txt")

# 写入train.txt和test.txt

with open("train.txt", "w") as f:
    for i in train_data:
        img_path = os.path.join("data/images", i)
        xml_path = os.path.join("data/Annotations", i.replace("jpg", "xml"))
        f.write(img_path + " " + xml_path + "\n")

with open("test.txt", "w") as f:
    for i in test_data:
        img_path = os.path.join("data/images", i)
        xml_path = os.path.join("data/Annotations", i.replace("jpg", "xml"))
        f.write(img_path + " " + xml_path + "\n")



数据集处理

为了训练我们的模型,我们需要准备并处理数据集。我们首先将数据集分为训练集和测试集,然后创建了对应的文本文件来存储图像和标注文件的路径。

  • 数据集类:我们定义了一个MyDataset类,它从给定的文本文件中读取图像和标注,并在需要时应用变换。
import cv2  # 导入OpenCV库用于图像处理
import xml.etree.ElementTree as ET  # 导入ElementTree库用于解析XML
import numpy as np  # 导入NumPy库用于数值计算
import paddle  # 导入PaddlePaddle库
from paddle.io import Dataset  # 导入PaddlePaddle的Dataset类


# 自定义数据集类,继承自PaddlePaddle的Dataset类
class MyDataset(Dataset):
    def __init__(self, txt_path, transform=None):
        super().__init__()
        self.transform = transform  # 数据增强的函数,可选
        self.data = []  # 存储图像和标注文件路径的列表
        with open(txt_path) as f:
            for line in f.readlines():
                self.data.append(line.strip().split(" "))  # 读取txt文件中的每一行,分割为图像路径和XML标注文件路径

    def __getitem__(self, idx):
        im = cv2.imread(self.data[idx][0])  # 读取图像,使用OpenCV库
        gt_bbox = self._get_xml(self.data[idx][1])  # 解析XML标注文件,获取目标框信息
        sample = {"image": im, "gt_bbox": np.array(gt_bbox, dtype=np.float64)}  # 构建样本字典,包括图像和目标框
        if self.transform:
            sample = self.transform(sample)  # 如果定义了数据增强函数,对样本进行数据增强操作
        return sample  # 返回样本字典

    def _get_xml(self, xml_path):
        root = ET.ElementTree(file=xml_path).getroot()  # 解析XML文件获取根节点
        object_list = root.findall("object")  # 查找所有object标签,每个标签对应一个目标物体
        gt_bbox = []  # 存储目标框的列表
        for o in object_list:
            bndbox = o.find("bndbox")  # 查找目标框坐标信息
            xmin = bndbox.find("xmin").text  # 获取xmin标签的文本内容,即目标框的左上角x坐标
            ymin = bndbox.find("ymin").text  # 获取ymin标签的文本内容,即目标框的左上角y坐标
            xmax = bndbox.find("xmax").text  # 获取xmax标签的文本内容,即目标框的右下角x坐标
            ymax = bndbox.find("ymax").text  # 获取ymax标签的文本内容,即目标框的右下角y坐标
            gt_bbox.append([eval(xmin), eval(ymin), eval(xmax), eval(ymax)])  # 将目标框坐标转换为浮点数并添加到列表中
        return gt_bbox  # 返回目标框的列表

    def __len__(self):
        return len(self.data)  # 返回数据集的长度,即样本数量

train_dataset = MyDataset("train.txt")

sample = train_dataset[0]
print(sample["image"].shape)
print(sample["gt_bbox"])
(397, 599, 3)
[[243. 189. 414. 290.]]

数据增强

数据增强是提高模型泛化能力的关键步骤。在本项目中,我们使用了PaddlePaddle的变换库来实现简单的数据增强,例如调整大小、归一化和重新排列维度。

from paddle.vision.transforms import Compose  # 导入Compose类,用于组合多个变换操作
from ppdet.data.transform import operators as ops  # 导入ppdet库中的数据变换操作

# 训练数据的变换操作列表
train_transforms = Compose([
    ops.Resize(target_size=[512, 512], keep_ratio=False),  # 调整图像大小为512x512,不保持宽高比
    ops.NormalizeImage(),  # 对图像进行归一化,将像素值缩放到0到1之间
    ops.Permute(),  # 调整图像通道顺序,通常是从HWC(Height x Width x Channels)到CHW(Channels x Height x Width)
])

# 测试数据的变换操作列表
test_transforms = Compose([
    ops.Resize(target_size=[512, 512], keep_ratio=False),  # 调整图像大小为512x512,不保持宽高比
    ops.NormalizeImage(),  # 对图像进行归一化,将像素值缩放到0到1之间
    ops.Permute(),  # 调整图像通道顺序,通常是从HWC(Height x Width x Channels)到CHW(Channels x Height x Width)
])

train_dataset = MyDataset("train.txt", transform=train_transforms)
test_dataset = MyDataset("test.txt", transform=test_transforms)

批处理函数

为了高效地训练我们的模型,我们定义了一个批处理函数,它将一批数据转换为模型可以理解的格式。

def collate_fn(batch):
    images = []  # 存储图像数据
    gt_bboxs = []  # 存储标注框数据
    for id, item in enumerate(batch):
        for bbox in item["gt_bbox"].tolist():  # 遍历每个样本中的标注框
            gt_bboxs.append([id, 0, *bbox])  # 将标注框的信息添加到gt_bboxs列表中,格式为:[样本ID, 类别ID, xmin, ymin, xmax, ymax]
        images.append(item["image"])  # 将图像添加到images列表中

    images = paddle.to_tensor(np.array(images, dtype=np.float32))  # 将图像列表转换为PaddlePaddle张量

    return images, gt_bboxs  # 返回图像张量和标注框列表


# 创建自定义数据集对象并加载数据
train_dataset = MyDataset("train.txt", transform=train_transforms)

# 创建数据加载器,设置批量大小为4,shuffle参数为True表示在每个epoch开始前对数据进行随机重排
train_loader = paddle.io.DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)

# 遍历数据加载器的第一个批次
for batch_id, data in enumerate(train_loader()):
    images, gt_bboxs = data
    print(images.shape)  # 打印图像张量的形状
    print(gt_bboxs)  # 打印标注框列表
    break

辅助函数

我们还实现了一些辅助函数来帮助处理数据和评估模型性能:

  • gt_bbox2gt_tensor:将标注的边界框转换为训练时使用的张量格式。
  • pred_tensor2pred_bbox:将模型的输出张量转换为可解释的边界框格式。
def gt_bbox2gt_tensor(gt_bbox, out_h, out_w, in_h, in_w, batch_size, num_classes):
    """
    将边界框数据转换为训练目标检测模型时所需的张量格式。

    gt_bbox: 边界框的列表,每个边界框的格式为 [batch_id, class_id, x1, y1, x2, y2]。
    out_h: 网络输出张量的高度。
    out_w: 网络输出张量的宽度。
    in_h: 输入图像的高度。
    in_w: 输入图像的宽度。
    batch_size: 批量处理的图像数量。
    num_classes: 目标类别的总数。
    """

    # 初始化存储边界框中心位置偏移量的张量。
    offset_xy = paddle.zeros([batch_size, 2, out_h, out_w])

    # 初始化存储边界框宽度和高度的张量。
    wh = paddle.zeros([batch_size, 2, out_h, out_w])

    # 初始化存储边界框存在的置信度的张量。
    confidence = paddle.zeros([batch_size, 1, out_h, out_w])

    # 初始化存储各个类别的张量。
    classes = paddle.zeros([batch_size, num_classes, out_h, out_w])

    # 遍历每个边界框并填充上述张量。
    for box in gt_bbox:
        # 解析边界框的各个组成部分。
        batch_id, class_id, x1, y1, x2, y2 = box

        # 计算边界框中心的 x, y 坐标。
        center_x = (x1 + x2) / 2 / in_w * out_w
        center_y = (y1 + y2) / 2 / in_h * out_h

        # 计算并存储中心位置的偏移量。
        offset_xy[batch_id, 0, int(center_y), int(center_x)] = center_x - int(center_x)
        offset_xy[batch_id, 1, int(center_y), int(center_x)] = center_y - int(center_y)

        # 计算并存储边界框的宽度和高度。
        wh[batch_id, 0, int(center_y), int(center_x)] = (x2 - x1) / in_w * out_w
        wh[batch_id, 1, int(center_y), int(center_x)] = (y2 - y1) / in_h * out_h

        # 在相应位置标记置信度为 1,表示该位置有物体。
        confidence[batch_id, 0, int(center_y), int(center_x)] = 1

        # 标记该物体所属的类别。
        classes[batch_id, class_id, int(center_y), int(center_x)] = 1

    # 返回处理后的张量。
    return offset_xy, wh, confidence, classes

def pred_tensor2pred_bbox(offset_xy, wh, confidence, classes, in_h, in_w, confidence_threshold=0.001):
    """
    将模型输出的张量转换为预测的边界框、置信度和类别信息。

    offset_xy: 形状为 [N, 2, out_h, out_w] 的张量,包含每个网格中心位置的偏移量预测。
    wh: 形状为 [N, 2, out_h, out_w] 的张量,包含每个边界框的宽度和高度预测。
    confidence: 形状为 [N, 1, out_h, out_w] 的张量,表示每个网格单元包含物体的置信度。
    classes: 形状为 [N, num_classes, out_h, out_w] 的张量,表示每个网格单元中物体可能属于各个类别的概率。
    in_h, in_w: 输入图像的高度和宽度。
    confidence_threshold: 置信度阈值,用于确定是否认为网格中包含物体。
    """

    N, _, out_h, out_w = offset_xy.shape  # 提取张量的形状,获取批次大小N和输出特征图的尺寸out_h, out_w。

    object_mask = confidence > confidence_threshold  # 创建一个对象掩码,标识每个网格单元是否包含物体。

    classes = paddle.argmax(classes, axis=1, keepdim=True)  # 对类别预测进行argmax操作,找到每个网格单元最可能的类别。

    x_grid = paddle.arange(0, out_w).reshape([1, -1]) + paddle.zeros([out_h, 1])  # 创建网格的x坐标。
    y_grid = paddle.arange(0, out_h).reshape([-1, 1]) + paddle.zeros([1, out_w])  # 创建网格的y坐标。

    pred_bbox = []  # 初始化用于存储预测边界框的列表。
    pred_scores = []  # 初始化用于存储预测置信度的列表。
    pred_classes = []  # 初始化用于存储预测类别的列表。

    for i in range(N):  # 遍历每个图像样本。
        sub_object_mask = object_mask[i, 0, :, :]  # 获取当前图像的对象掩码。

        # 提取当前图像的偏移量、网格坐标、边界框尺寸、置信度和类别信息。
        o_x = offset_xy[i, 0, :, :][sub_object_mask].numpy()
        o_y = offset_xy[i, 1, :, :][sub_object_mask].numpy()
        x_g = x_grid[sub_object_mask].numpy()
        y_g = y_grid[sub_object_mask].numpy()
        c_x = ((o_x + x_g) / out_w * in_w).tolist()
        c_y = ((o_y + y_g) / out_h * in_h).tolist()
        w = (wh[i, 0, :, :][sub_object_mask].numpy() / out_w * in_w).tolist()
        h = (wh[i, 0, :, :][sub_object_mask].numpy() / out_h * in_h).tolist()
        s = confidence[i, 0, :, :][sub_object_mask].numpy().tolist()
        c = classes[i, 0, :, :][sub_object_mask].numpy().tolist()

        sub_bbox = []
        sub_scores = []
        sub_classes = []
        for j in range(len(o_x)):  # 遍历当前图像中所有检测到的对象。
            # 计算并存储每个边界框的坐标、置信度和类别。
            sub_bbox.append([
                c_x[j] - w[j] / 2,  # 边界框左上角x坐标。
                c_y[j] - h[j] / 2,  # 边界框左上角y坐标。
                c_x[j] + w[j] / 2,  # 边界框右下角x坐标。
                c_y[j] + h[j] / 2  # 边界框右下角y坐标。
            ])
            sub_scores.append(s[j])
            sub_classes.append(c[j])

        pred_bbox.append(sub_bbox)
        pred_scores.append(sub_scores)
        pred_classes.append(sub_classes)

    return pred_bbox, pred_scores, pred_classes  # 返回预测的边界框、置信度和类别信息。

# shape: [out_h, out_w]
x_grid = paddle.arange(0, 7).reshape([1, -1]) + paddle.zeros([7, 1])
y_gride = paddle.arange(0, 7).reshape([-1, 1]) + paddle.zeros([1, 7])
print(x_grid)
print(y_gride)

损失函数

YOLO模型使用了一种特殊的损失函数,它结合了坐标损失、置信度损失和分类损失。

class YOLOLoss(nn.Layer):
    def __init__(self):
        super().__init__()
        # 使用均方误差作为损失函数,不进行求和或平均,以便于后续操作
        self.mse_loss = nn.MSELoss(reduction='none')
        # 设置坐标损失的权重系数
        self.lambda_coord = 5.
        # 设置没有目标的损失的权重系数
        self.lambda_noobj = 0.5

    def forward(self, offset_xy, wh, confidence, classes, gt_offset_xy, gt_wh, gt_confidence, gt_classes):
        # 识别出有物体的网格(目标掩码)
        object_mask = gt_confidence > 0
        # 计算预测的偏移量(offset_xy)与真实值(gt_offset_xy)之间的损失,并仅对有目标的网格求和
        offset_loss = self.mse_loss(offset_xy, gt_offset_xy)[
            (object_mask.astype('float32') + paddle.zeros_like(offset_xy)) > 0].sum()
        # 计算预测的宽高(wh)与真实的宽高(gt_wh)之间的损失,并仅对有目标的网格求和
        wh_loss = self.mse_loss(paddle.sqrt(wh + 1e-6), paddle.sqrt(gt_wh + 1e-6))[
            (object_mask.astype('float32') + paddle.zeros_like(offset_xy)) > 0].sum()
        # 计算预测的置信度(confidence)与真实置信度(gt_confidence)之间的损失
        confidence_loss = self.mse_loss(confidence, gt_confidence)
        # 对有目标的网格中的置信度损失求和
        obj_c_loss = confidence_loss[object_mask].sum()
        # 对没有目标的网格中的置信度损失求和
        noobj_c_loss = confidence_loss[object_mask == False].sum()
        # 计算预测的类别(classes)与真实类别(gt_classes)之间的损失,并仅对有目标的网格求和
        classes_loss = self.mse_loss(classes, gt_classes)[
            (object_mask.astype('float32') + paddle.zeros_like(classes)) > 0].sum()

        # 计算总损失,其中包括坐标损失、有目标的置信度损失、无目标的置信度损失和类别损失
        total_loss = (
                                 offset_loss + wh_loss) * self.lambda_coord + obj_c_loss + noobj_c_loss * self.lambda_noobj + classes_loss

        return total_loss

#测试
offset_xy = paddle.rand([4, 2, 7, 7])
wh = paddle.rand([4, 2, 7, 7])
confidence = paddle.rand([4, 1, 7, 7])
classes = paddle.rand([4, 1, 7, 7])
gt_offset_xy = paddle.rand([4, 2, 7, 7])
gt_wh = paddle.rand([4, 2, 7, 7])
gt_confidence = paddle.rand([4, 1, 7, 7])
gt_classes = paddle.rand([4, 1, 7, 7])
loss = YOLOLoss()
total_loss = loss(offset_xy, wh, confidence, classes, gt_offset_xy, gt_wh, gt_confidence, gt_classes)
print(total_loss)

训练和评估

我们使用了PaddlePaddle的优化器和训练循环来训练模型,并使用了特定的度量标准来评估模型性能。

from ppdet.metrics.map_utils import DetectionMAP

class Metric:
    def __init__(self, num_classes):
        # 初始化检测评估指标类,设置类别数量,并指定类别名称(在这里只有一个类别,标记为'fall')
        self.d_map = DetectionMAP(num_classes, catid2name={0: 'fall'})

    def __call__(self, pred_bbox, pred_scores, pred_label, gt_bbox, gt_label):
        # 在每次评估前重置检测指标
        self.d_map.reset()
        for i in range(len(pred_bbox)):
            # 更新评估指标,根据预测的边界框、分数、标签和真实的边界框、标签
            self.d_map.update(pred_bbox[i], pred_scores[i], pred_label[i], gt_bbox[i], gt_label[i])

        # 累计计算评估指标
        self.d_map.accumulate()

        # 返回平均精度(mean Average Precision)
        return self.d_map.get_map()

def nms(pred_bbox, pred_scores, pred_classes):
    # 初始化新的预测结果列表
    new_pred_bbox = []
    new_pred_scores = []
    new_pred_classes = []

    for i in range(len(pred_bbox)):
        # 为每个图像添加一个新的结果列表
        new_pred_bbox.append([])
        new_pred_scores.append([])
        new_pred_classes.append([])
        # 检查是否有预测边界框
        if len(pred_bbox[i]) > 0:
            # 应用非极大值抑制(NMS),以减少重叠的边界框
            idxs = paddle.vision.ops.nms(boxes=paddle.to_tensor(pred_bbox[i]))
            for j in idxs:
                # 将NMS后的边界框、分数、类别添加到新列表中
                new_pred_bbox[-1].append(pred_bbox[i][j])
                new_pred_scores[-1].append(pred_scores[i][j])
                new_pred_classes[-1].append(pred_classes[i][j])
    # 返回经过NMS处理后的预测结果
    return new_pred_bbox, new_pred_scores, new_pred_classes

import paddle

# 基础配置
num_classes = 1  # 设置类别数量为1
batch_size = 32  # 设置批量大小为32
learning_rate = 0.01  # 设置学习率为0.01

# 模型
resnet18 = ResNet18()  # 创建一个ResNet18作为YOLO模型的骨干网络
yolo = YOLO(backbone=resnet18, channels=512, num_classes=num_classes)  # 使用ResNet18骨干网络创建YOLO模型

# 数据
# 创建训练数据集,指定数据集文件和转换函数
train_dataset = MyDataset('train.txt', train_transforms)
# 创建训练数据加载器,用于在训练过程中加载数据
train_dataloader = paddle.io.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
# 创建测试数据集,指定数据集文件和转换函数
test_dataset = MyDataset('test.txt', test_transforms)
# 创建测试数据加载器,用于在测试过程中加载数据
test_dataloader = paddle.io.DataLoader(test_dataset, batch_size=1, collate_fn=collate_fn)

# 评价函数
metric = Metric(num_classes=num_classes)  # 初始化评估指标对象

# 损失函数
loss_fn = YOLOLoss()  # 初始化YOLO损失函数

# 优化器
# 使用Adam优化器,并设置学习率和优化的参数
optimizer = paddle.optimizer.Adam(learning_rate=learning_rate, parameters=yolo.parameters())

# 设置训练的总轮数
epochs = 50

for epoch in range(epochs):
    # 开始训练模式
    yolo.train()
    train_total_loss = 0
    train_total_ap = 0
    print('----------------------- Train -----------------------')
    for batch_id, batch in enumerate(train_dataloader):
        # 从模型中获取预测的边界框、宽高、置信度和类别
        offset_xy, wh, confidence, classes = yolo(batch[0])
        # 获取输入图片的尺寸信息
        N, _, in_h, in_w = batch[0].shape
        # 获取预测结果的尺寸信息
        out_h, out_w = offset_xy.shape[2:]
        # 将真实的标注信息转换为用于训练的张量格式
        gt_offset_xy, gt_wh, gt_confidence, gt_classes = gt_bbox2gt_tensor(batch[1], out_h, out_w, in_h, in_w, N, num_classes)
        # 将预测得到的张量转换为预测框
        pred_bbox, pred_scores, pred_classes = pred_tensor2pred_bbox(offset_xy, wh, confidence, classes, in_h, in_w, 0.001)
        # 应用非极大值抑制(NMS)
        pred_bbox, pred_scores, pred_classes = nms(pred_bbox, pred_scores, pred_classes)

        # 计算损失
        step_loss = loss_fn(offset_xy, wh, confidence, classes, gt_offset_xy, gt_wh, gt_confidence, gt_classes)
        # 反向传播
        step_loss.backward()
        # 更新模型参数
        optimizer.step()
        # 清除梯度
        optimizer.clear_grad()

        # 将数据读取的标注信息转换为需要的格式
        gt_bbox = []
        gt_label = []
        for j in range(N):
            gt_bbox.append([])
            gt_label.append([])
        for item in batch[1]:
            gt_bbox[item[0]].append(item[2:])
            gt_label[item[0]].append(item[1])

        # 计算平均精度(AP)
        ap = metric(pred_bbox, pred_scores, pred_classes, gt_bbox, gt_label)
        # 记录累计损失和平均精度
        train_total_loss += step_loss.item()
        train_total_ap += ap

        # 定期打印训练状态
        if batch_id % 50 == 0:
            print(f'Train epoch/epochs:{epoch + 1}/{epochs} batch_id/total_batch:{batch_id + 1}/{len(train_dataloader)} loss:{step_loss.item()} ap: {ap}')
    # 每个epoch结束后打印总体训练状态
    print(f'Train epoch/epochs:{epoch + 1}/{epochs} loss:{train_total_loss / len(train_dataloader)} ap:{train_total_ap / len(train_dataloader)}')
    # 保存模型参数
    paddle.save(yolo.state_dict(), 'yolo.pdparams')

    # 开始测试模式
    yolo.eval()
    test_total_loss = 0
    test_total_ap = 0
    print('----------------------- Test -----------------------')
    for batch_id, batch in enumerate(test_dataloader):
        # 同样的过程应用于测试数据
        offset_xy, wh, confidence, classes = yolo(batch[0])
        N, _, in_h, in_w = batch[0].shape
        out_h, out_w = offset_xy.shape[2:]
        gt_offset_xy, gt_wh, gt_confidence, gt_classes = gt_bbox2gt_tensor(batch[1], out_h, out_w, in_h, in_w, N, num_classes)
        pred_bbox, pred_scores, pred_classes = pred_tensor2pred_bbox(offset_xy, wh, confidence, classes, in_h, in_w, 0.001)
        pred_bbox, pred_scores, pred_classes = nms(pred_bbox, pred_scores, pred_classes)

        step_loss = loss_fn(offset_xy, wh, confidence, classes, gt_offset_xy, gt_wh, gt_confidence, gt_classes)
        gt_bbox = []
        gt_label = []
        for j in range(N):
            gt_bbox.append([])
            gt_label.append([])
        for item in batch[1]:
            gt_bbox[item[0]].append(item[2:])
            gt_label[item[0]].append(item[1])

        ap = metric(pred_bbox, pred_scores, pred_classes, gt_bbox, gt_label)
        test_total_loss += step_loss.item()
        test_total_ap += ap

    # 打印测试结果
    print(f'test epoch/epochs:{epoch + 1}/{epochs} loss:{test_total_loss / len(test_dataloader)} ap:{test_total_ap / len(test_dataloader)}')

模型预测和可视化

最后,我们展示了如何使用训练好的模型进行预测,并在图像上可视化预测的边界框。

import cv2
import matplotlib.pyplot as plt
import os

num_classes = 1  # 设置类别数量为1
# 创建YOLO模型实例
resnet18 = ResNet18()
yolo = YOLO(backbone=resnet18, channels=512, num_classes=num_classes)
# 加载训练好的模型参数
yolo.set_state_dict(paddle.load('yolo.pdparams'))

# 准备测试数据
test_dataset = MyDataset('test.txt', test_transforms)
idx = 1  # 选择要可视化的样本索引

# 获取指定索引的测试样本
sample = test_dataset[idx]
# 读取图片并调整尺寸到模型输入尺寸
image = cv2.imread(test_dataset.data[idx][0])
image = cv2.resize(image, dsize=[512, 512])
# 使用matplotlib显示原始图片
plt.imshow(image)
plt.show()

# 将图片输入模型进行预测
offset_xy, wh, confidence, classes = yolo(paddle.to_tensor([sample['image']]))
# 根据预测结果生成预测的边界框
pred_bbox, pred_scores, pred_classes = pred_tensor2pred_bbox(offset_xy, wh, confidence, classes, 512, 512, 0.001)
# 应用非极大值抑制(NMS)处理重叠的边界框
pred_bbox, pred_scores, pred_classes = nms(pred_bbox, pred_scores, pred_classes)

# 在图片上标记真实边界框(绿色框)
for box in sample['gt_bbox']:
    cv2.rectangle(image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 4)

# 在图片上标记预测的边界框(红色框)
for box in pred_bbox[0]:
    cv2.rectangle(image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 255), 4)

# 使用matplotlib显示标记后的图片
plt.imshow(image)
plt.show()

你可能感兴趣的:(AI,paddlepaddle,YOLO,人工智能)