在这里插入图片描述
在这里插入图片描述
如何使用YOLOv8进行安全帽和安全背心穿戴识别,并提供详细的训练代码和数据集准备步骤。在这里插入图片描述

项目结构
深色版本
safety_gear_detection/
├── dataset/
│ ├── images/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
│ ├── labels/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
├── weights/
│ └── best.pt
├── src/
│ ├── train.py
│ ├── predict.py
│ ├── utils.py
├── requirements.txt
└── README.md

  1. 安装依赖
    首先,确保你已经安装了必要的库。创建一个requirements.txt文件,内容如下:

深色版本
opencv-python
numpy
pillow
ultralytics
然后,使用以下命令安装依赖:

bash
深色版本
pip install -r requirements.txt
2. 数据集准备
确保你的数据集已经按照以下结构组织:

深色版本
dataset/
├── images/
│ ├── train/
│ ├── val/
│ └── test/
└── labels/
├── train/
├── val/
└── test/
每个文件夹中包含对应的图像文件和标签文件。确保所有图像文件都是.jpg格式,而标签文件是.txt格式(YOLO格式),并且它们的名字与对应的图像文件相同。

  1. 数据集配置文件
    创建一个数据集配置文件(如safety_gear_dataset.yaml),该文件定义了数据集的基本信息,包括路径、类别等。示例配置如下:

yaml
深色版本

训练和验证的数据集路径

train: dataset/images/train
val: dataset/images/val
test: dataset/images/test

标签路径

labels_train: dataset/labels/train
labels_val: dataset/labels/val
labels_test: dataset/labels/test

类别名称

names:
0: vest
1: no-helmet
2: helmet
3: no-vest

类别数量

nc: 4
4. 训练代码
4.1 train.py
这是训练模型的主程序文件。

python
深色版本
import subprocess

定义路径

dataset_config = “safety_gear_dataset.yaml”
model_config = “yolov8.yaml”
pretrained_weights = “yolov8x.pt”
output_dir = “runs/detect/train/exp”

训练模型

train_command = [
“python”, “ultralytics/yolo/v8/detect/train.py”,
“–data”, dataset_config,
“–cfg”, model_config,
“–weights”, pretrained_weights,
“–batch-size”, “16”,
“–epochs”, “200”
]

subprocess.run(train_command)
5. 模型评估
训练完成后,可以通过验证集来评估模型的性能。YOLOv8提供了方便的命令来进行模型评估:

bash
深色版本
python ultralytics/yolo/v8/detect/val.py --data safety_gear_dataset.yaml --weights runs/detect/train/exp/weights/best.pt
6. 模型推理
你可以使用训练好的模型进行推理,检测新的图像中的目标。示例命令如下:

bash
深色版本
python ultralytics/yolo/v8/detect/predict.py --source path/to/your/image.jpg --weights runs/detect/train/exp/weights/best.pt --conf 0.25
7. 详细代码示例(Python脚本)
7.1 predict.py
这是进行推理的脚本。

python
深色版本
import cv2
import numpy as np
from ultralytics import YOLO

加载模型

model = YOLO(‘runs/detect/train/exp/weights/best.pt’)

定义类别名称

class_names = [‘vest’, ‘no-helmet’, ‘helmet’, ‘no-vest’]

预测图像

def predict_image(image_path):
image = cv2.imread(image_path)
results = model(image)
for result in results:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = box.conf[0].item()
cls = int(box.cls[0].item())
label = f’{class_names[cls]} {conf:.2f}’
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return image

主函数

if name == “main”:
image_path = “path/to/your/image.jpg”
result_image = predict_image(image_path)
cv2.imshow(‘Safety Gear Detection’, result_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
8. 运行项目
确保你的数据集和训练好的权重已经放在相应的文件夹中。
在项目根目录下运行以下命令启动训练:
bash
深色版本
python src/train.py
训练完成后,运行以下命令进行评估:
bash
深色版本
python ultralytics/yolo/v8/detect/val.py --data safety_gear_dataset.yaml --weights runs/detect/train/exp/weights/best.pt
运行以下命令进行推理:
bash
深色版本
python src/predict.py
9. 详细注释
train.py
导入库:导入subprocess库用于调用命令行命令。
定义路径:定义数据集配置文件、模型配置文件、预训练权重和输出目录。
训练模型:使用subprocess.run调用YOLOv8的训练命令。
predict.py
导入库:导入OpenCV、NumPy和YOLOv8。
加载模型:使用YOLO类加载训练好的权重。
定义类别名称:定义类别名称。
预测图像:定义一个函数predict_image,用于对单张图像进行安全帽和安全背心穿戴识别,并在图像上绘制检测框和标签。
主函数:定义主函数,读取图像路径,调用predict_image函数进行检测,并显示结果。

Logo

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

更多推荐