如何用YOLOv8训练一个识别安全帽的模型？

训练一个识别安全帽的模型是一个典型的 目标检测任务，可以使用 YOLOv8 来实现。以下是详细的步骤：

1. 准备数据集

训练模型的第一步是准备一个包含安全帽图像的数据集。数据集需要满足以下要求：

图像：包含安全帽的图像。
标注：每张图像需要标注安全帽的位置（边界框）和类别。

（1）数据收集

收集包含安全帽的图像，可以从公开数据集中获取，或者自己拍摄。
公开数据集推荐：
- Safety Helmet Dataset
- COCO 数据集（如果包含安全帽类别）。

（2）数据标注

使用标注工具（如 LabelImg、CVAT 或 Roboflow）对图像进行标注。
标注格式：YOLO 格式（每个图像对应一个 .txt 文件，包含边界框和类别信息）。

标注文件示例（image1.txt）：

0 0.5 0.5 0.2 0.3  # 类别编号 中心点x 中心点y 宽度 高度

2. 安装 Ultralytics

确保你已经安装了 Ultralytics：

pip install ultralytics

3. 准备配置文件

YOLOv8 需要一个配置文件来定义数据集和训练参数。

（1）数据集配置文件

创建一个 .yaml 文件（如 helmet_dataset.yaml），内容如下：

# 数据集路径
path: D:/datasets/helmet  # 数据集根目录
train: images/train       # 训练集图像路径
val: images/val           # 验证集图像路径# 类别信息
names:0: helmet  # 安全帽类别

（2）目录结构

确保数据集目录结构如下：

D:/datasets/helmet/
├── images/
│   ├── train/  # 训练集图像
│   └── val/    # 验证集图像
├── labels/
│   ├── train/  # 训练集标注
│   └── val/    # 验证集标注
└── helmet_dataset.yaml  # 数据集配置文件

4. 训练模型

使用 YOLOv8 训练模型非常简单。以下是训练代码：

from ultralytics import YOLO# 加载预训练模型
model = YOLO("yolov8n.pt")  # 使用 YOLOv8n 模型# 训练模型
results = model.train(data="D:/datasets/helmet/helmet_dataset.yaml",  # 数据集配置文件epochs=50,  # 训练轮数imgsz=640,  # 图像大小batch=16,   # 批量大小name="helmet_detection"  # 训练任务名称
)

5. 验证模型

训练完成后，可以使用验证集评估模型性能：

# 验证模型
metrics = model.val()
print(metrics.box.map)  # 打印 mAP 值

6. 使用模型进行推理

训练完成后，可以使用模型对新的图像或视频进行推理：

# 加载训练好的模型
model = YOLO("runs/detect/helmet_detection/weights/best.pt")# 对图像进行推理
results = model("D:/datasets/helmet/images/test/image1.jpg")# 显示结果
results[0].show()# 保存结果
results[0].save("output.jpg")