数据集包括3个类:在这里插入图片描述

carry objects(携带工具)、helmet(戴头盔)、no-helmet(没带头盔)
按8:1:1比例划分为train:5110张,val:639张,test:639张
yolo系列可以直接拿来进行训练
本数据集面对的是复杂场景的安全帽图片,
数据集中包含有煤矿场景,因图片目标较为单一,为扩充数据集,对其他安全帽数据集的图片进行了复杂场景的处理,同时重新使用新的类标注图片。在这里插入图片描述
在这里插入图片描述
如何使用 YOLOv8 训练矿井安全帽检测数据集的详细步骤。YOLOv8 是 YOLO 系列的最新版本,具有更高的精度和更快的推理速度。

一、数据准备

  1. 数据集结构
    假设你的数据集已经按照 8:1:1 的比例划分好,并且结构如下:

深色版本
mine_safety_helmet_dataset/
├── images/
│ ├── train/
│ │ ├── 000001.jpg
│ │ ├── 000002.jpg
│ │ └── …
│ ├── val/
│ │ ├── 000001.jpg
│ │ ├── 000002.jpg
│ │ └── …
│ └── test/
│ ├── 000001.jpg
│ ├── 000002.jpg
│ └── …
├── labels/
│ ├── train/
│ │ ├── 000001.txt
│ │ ├── 000002.txt
│ │ └── …
│ ├── val/
│ │ ├── 000001.txt
│ │ ├── 000002.txt
│ │ └── …
│ └── test/
│ ├── 000001.txt
│ ├── 000002.txt
│ └── …
└── mine_safety_helmet.yaml
images/ 目录包含输入图像。
labels/ 目录包含对应的标注文件,每个标注文件中的每一行表示一个目标框,格式为 class_id x_center y_center width height。
mine_safety_helmet.yaml 是数据集的配置文件。
2. 配置文件
创建一个 mine_safety_helmet.yaml 文件,内容如下:

yaml
深色版本

数据集路径

path: ./mine_safety_helmet_dataset

类别名称

names:
0: carry_objects
1: helmet
2: no_helmet

数据集划分

train: images/train
val: images/val
test: images/test
二、安装依赖
确保你已经安装了 YOLOv8 所需的依赖。可以使用以下命令安装:

bash
深色版本
pip install ultralytics
三、训练模型

  1. 导入库
    python
    深色版本
    import torch
    from ultralytics import YOLO
  2. 加载模型
    python
    深色版本

加载预训练的YOLOv8模型

model = YOLO(‘yolov8n.yaml’) # 你可以选择其他预训练模型,如’yolov8s.yaml’、'yolov8m.yaml’等
3. 配置训练参数
python
深色版本

配置训练参数

model.train(
data=‘mine_safety_helmet.yaml’, # 数据集配置文件
epochs=100, # 训练轮数
imgsz=640, # 输入图像大小
batch=16, # 批量大小
name=‘mine_safety_helmet’, # 训练结果保存的文件夹名称
device=‘cuda’ # 使用GPU进行训练
)
四、评估模型
训练完成后,可以使用验证集来评估模型的性能。

python
深色版本

评估模型

results = model.val(data=‘mine_safety_helmet.yaml’, imgsz=640, batch=16, device=‘cuda’)

打印评估结果

print(results)
五、保存模型
训练完成后,模型会自动保存到指定的文件夹中。你也可以手动保存模型:

python
深色版本

保存模型

model.save(‘mine_safety_helmet_model.pt’)
六、完整代码
将上述所有步骤整合成一个完整的脚本:

python
深色版本
import torch
from ultralytics import YOLO

加载预训练的YOLOv8模型

model = YOLO(‘yolov8n.yaml’) # 你可以选择其他预训练模型,如’yolov8s.yaml’、'yolov8m.yaml’等

配置训练参数

model.train(
data=‘mine_safety_helmet.yaml’, # 数据集配置文件
epochs=100, # 训练轮数
imgsz=640, # 输入图像大小
batch=16, # 批量大小
name=‘mine_safety_helmet’, # 训练结果保存的文件夹名称
device=‘cuda’ # 使用GPU进行训练
)

评估模型

results = model.val(data=‘mine_safety_helmet.yaml’, imgsz=640, batch=16, device=‘cuda’)

打印评估结果

print(results)

保存模型

model.save(‘mine_safety_helmet_model.pt’)
七、运行脚本
将上述代码保存为一个 Python 脚本文件(例如 train_mine_safety_helmet.py),然后在命令行中运行:

bash
深色版本
python train_mine_safety_helmet.py
八、注意事项
数据集路径:确保数据集路径正确,特别是图片和标注文件的路径。
硬件资源:训练大型模型可能需要较多的计算资源,建议使用 GPU。
超参数调整:根据实际情况调整学习率、批量大小等超参数,以获得更好的训练效果。
数据增强:为了提高模型的泛化能力,可以在训练过程中使用数据增强技术,如随机裁剪、翻转等。
希望这些步骤能帮助你成功训练一个矿井安全帽检测模型。

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐