YOLOv9 在自定义数据集上进行目标检测 | 计算机视觉项目-51CTO.COM

在快速发展的计算机视觉领域，目标检测是一个基石应用程序。在众多的目标检测算法中，YOLOv9已经成为一个强大且多功能的解决方案，提供了实时检测能力，并具有令人印象深刻的准确性。在这篇全面的指南中，我们将探索在自定义数据集上训练YOLOv9的过程，并在测试数据上进行推理。

通过本教程的学习，你将深入了解YOLOv9的工作原理以及如何使用自定义数据集在你的项目中实现它。那么，让我们开始一起探索YOLOv9的强大功能吧！

激活函数系列

Softmax激活函数
Sigmoid/Logistic激活函数
双曲正切（Tanh）激活函数
修正线性单元（ReLU）激活函数

数据集：https://www.kaggle.com/datasets/snehilsanyal/construction-site-safety-image-dataset-roboflow

克隆YOLOv9仓库

!git clone https://github.com/SkalskiP/yolov9.git

让我们使用git clone命令从GitHub克隆YOLOv9仓库，并开始探索使用YOLOv9进行实时目标检测。Git是一个版本控制系统，允许开发人员跟踪他们的代码库的更改，与他人协作和管理项目历史。

克隆过程完成后，我们将在机器上拥有YOLOv9仓库的本地副本。这将使我们能够探索代码库，进行修改，并使用YOLOv9算法进行我们自己的目标检测任务。

dataDir = '/content/css-data/' # css-data is the unzip path of the dataset
workingDir = '/content/' # Working Dir in google colab

在这里，我们为项目设置目录路径。变量dataDir保存我们CSS（建筑工地安全）数据存储的目录路径。我们将其设置为'/content/css-data/'，表明我们的CSS数据集文件位于'/content/'目录内的名为'css-data'的目录中。

接下来，我们有变量workingDir，它定义了我们工作目录的路径。

num_classes = 10
classes = ['Hardhat', 'Mask', 'NO-Hardhat', 'NO-Mask', 'NO-Safety Vest', 'Person', 'Safety Cone', 'Safety Vest', 'machinery', 'vehicle']

定义了两个变量：

num_classes设置为整数值10，表示分类任务中的类别数。
classes是一个包含10个字符串元素的列表，代表不同的类别：'Hardhat'（安全帽）、'Mask'（口罩）、'NO-Hardhat'（无安全帽）、'NO-Mask'（无口罩）、'NO-Safety Vest'（无安全背心）、'Person'（人员）、'Safety Cone'（安全锥）、'Safety Vest'（安全背心）、'machinery'（机械）和'vehicle'（车辆）。


import yaml
import os

file_dict = {
    'train': os.path.join(dataDir, 'train'),
    'val': os.path.join(dataDir, 'valid'),
    'test': os.path.join(dataDir, 'test'),
    'nc': num_classes,
    'names': classes
}
with open(os.path.join(workingDir,'yolov9', 'data.yaml'), 'w+') as f:
  yaml.dump(file_dict, f)

定义了一个名为file_dict的字典，包含以下键值对：

‘train’：训练数据目录的路径，通过连接‘dataDir’变量和‘train’形成。
‘val’：验证数据目录的路径，通过连接‘dataDir’变量和‘valid’形成。
‘test’：测试数据目录的路径，通过连接‘dataDir’变量和‘test’形成。
‘nc’：类别数，来源于‘num_classes’变量。
‘names’：类别名称列表，来源于‘classes’变量。

在‘workingDir/yolov9’目录内以写入模式打开名为‘data.yaml’的文件。使用YAML格式将file_dict字典的内容写入'data.yaml'文件。如果文件不存在，则创建一个新文件。

下载YOLOv9权重

!wget  https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt

我们现在正在获取YOLOv9的预训练权重文件。使用!wget，我们可以直接从网络环境中检索文件。

在这个命令中，我们正在下载名为yolov9-c.pt的文件，该文件来自托管在GitHub上的YOLOv9仓库的v0.1版本。这些预训练权重对于使用学习到的参数初始化我们的YOLOv9模型至关重要，使其能够准确检测对象。

下载完成后，我们将能够使用这些权重来微调我们的自定义数据集进行目标检测任务，节省了从头开始训练模型所需的时间和计算资源。

cd yolov9

现在，让我们导航到‘yolov9’目录。‘yolov9’是我们刚刚从GitHub克隆的仓库。这样做，我们将获得‘yolov9’模型的代码，我们将用于在自定义数据集上训练yolov9模型。

!pip install -r requirements.txt -q

接下来，我们使用pip包管理器安装项目所需的Python包和依赖项。!pip install命令允许我们从Python包索引（PyPI）或指定的要求文件安装Python包。

在这里，-r requirements.txt告诉pip安装'requirements.txt'文件中列出的所有包。该文件包含训练yolov9模型所必需的Python库列表。

-q标志代表'quiet'模式，这意味着pip将在没有错误的情况下静默执行安装过程，不会显示详细输出。这有助于保持安装过程的简洁明了，特别是在脚本或自动化任务中。

!python train_dual.py --workers 8 --batch 4  --img 640 --epochs 50 --data /content/yolov9/data.yaml --weights /content/yolov9-e.pt --device 0 --cfg /content/yolov9/models/detect/yolov9.yaml --hyp /content/yolov9/data/hyps/hyp.scratch-high.yaml

这个命令启动了YOLOv9模型的训练过程。但在执行单元格之前，我们需要修改两个参数，第一个参数是—data，我们将在—data参数中传递上面创建的data.yaml文件路径，第二个参数是—cfg，我们需要对yolov9.yaml文件进行一些更改，因为这个原始的yolov9.yaml文件是为COCO数据集配置的，COCO数据集包含80个类别，而我们的自定义数据集有10个类别，所以我们需要在yolov9.yaml中调整nc参数并传递nc=10，因为我们有10个类别，而不是80，我希望你明白我的意思—hyp包含模型的超参数，我们不需要修改它，所以保持原样。

现在让我们使用!python train_dual.py微调我们的自定义数据集上的yolov9模型：

--workers 8：这个参数指定了训练期间数据预处理的数据加载工作器的数量。更高的数字可以加快数据加载和预处理的速度。
--batch 4：这个参数设置了训练的批次大小。批次大小决定了每次训练迭代中使用的样本数量。
--img 640：这个参数指定了训练的输入图像大小。在这种情况下，图像将在训练期间调整为640像素的高度和宽度。
--epochs 5：这个参数设置了训练的周期数。一个周期是完整地遍历整个训练数据集一次。
--data /content/yolov9/data.yaml：这个参数指定了包含数据集配置信息的YAML文件的路径。
--weights /content/yolov9-c.pt：这个参数指定了用于在训练前初始化模型的预训练权重文件（yolov9-c.pt）的路径。
--device 0：这个参数指定了用于训练的设备。这里，0可能表示第一个可用于训练的GPU设备。
--cfg /content/yolov9/models/detect/yolov9.yaml：这个参数指定了模型配置文件（yolov9.yaml）的路径，该文件定义了YOLOv9模型的架构和参数。
--hyp /content/yolov9/data/hyps/hyp.scratch-high.yaml：这个参数指定了超参数文件（hyp.scratch-high.yaml）的路径，该文件包含训练期间模型的调整参数。

推理

!python detect.py --img 640 --conf 0.1 --device 0 --weights /content/yolov9/runs/train/exp2/weights/best.pt --source /content/css-data/test/images/004763_jpg.rf.46484e6ca73caeaa9de45822cf1085a9.jpg

from IPython.display import Image
Image(filename="/content/yolov9/runs/detect/exp2/004763_jpg.rf.46484e6ca73caeaa9de45822cf1085a9.jpg", width=600)