使用YOLOv5构建安全帽检测和识别系统

摘要：在本博客中，我们将探讨如何使用YOLOv5构建一个实时安全帽检测和识别系统。这个系统可以在工地等场所有效地检测工人是否佩戴了安全帽，从而提高安全意识和预防安全事故的发生。

目录

1.介绍

2.YOLOv5简介

3.数据集准备

4.YOLOv5模型训练

4.1首先，我们需要将YOLOv5的代码克隆到本地：

4.2 创建数据集配置文件

4.3 创建或修改YOLOv5配置文件

4.4训练YOLOv5模型

5.模型评估与调优

6.系统部署与实时检测

7.总结

---

正文：

1.介绍

安全帽检测和识别系统对于确保工地安全至关重要。这种系统可以实时监控场景中的人员，检测他们是否佩戴了安全帽。在本文中，我们将介绍如何使用YOLOv5构建一个这样的系统。

2.YOLOv5简介

YOLOv5（You Only Look Once version 5）是一种快速且准确的实时目标检测算法。相较于其前几代，YOLOv5具有更低的计算复杂性、更高的精度和更快的速度。此外，它支持多种尺寸的模型，可以根据硬件资源和性能需求进行选择。

3.数据集准备

为了训练我们的安全帽检测和识别系统，我们需要一个包含带安全帽和不带安全帽的人员图像的数据集。数据集应该包含两个类别：“带安全帽”和“不带安全帽”。在收集数据集后，我们需要将其分为训练集、验证集和测试集，并为每张图像提供标签数据。

4.YOLOv5模型训练

4.1首先，我们需要将YOLOv5的代码克隆到本地：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5

然后，安装YOLOv5所需的依赖库： 

pip install -r requirements.txt

4.2 创建数据集配置文件

我们需要为YOLOv5创建一个数据集配置文件，以指定训练集、验证集的路径和类别信息。在yolov5/data目录下，创建一个名为helmet_detection.yaml的文件，内容如下：

train: ../helmet_dataset/train/images
val: ../helmet_dataset/val/images

nc: 2
names: ["with_helmet", "without_helmet"]

这里，train和val字段分别指定训练集和验证集图像的路径。nc表示类别数量，本例中为2（带安全帽和不带安全帽）。names字段包含每个类别的名称。

4.3 创建或修改YOLOv5配置文件

为了训练YOLOv5模型，我们需要一个合适的配置文件。YOLOv5提供了多种预定义的配置文件，例如yolov5s.yaml（小型）、yolov5m.yaml（中型）和yolov5l.yaml（大型）。根据硬件性能和目标检测精度的需求，您可以选择合适的配置文件。在本例中，我们使用yolov5s.yaml。

确保配置文件中的类别数量与数据集配置文件中的一致。在yolov5/models目录下，找到yolov5s.yaml文件，并将其中的nc字段修改为2（与helmet_detection.yaml中的类别数量一致）：

nc: 2  # number of classes

 

4.4训练YOLOv5模型

准备好数据集配置文件和模型配置文件后，我们可以开始训练YOLOv5模型。在命令行中运行以下命令：

python train.py --img 640 --batch 16 --epochs 100 --data helmet_detection.yaml --cfg yolov5s.yaml --weights yolov5s.pt --name helmet_detection

其中：

• --img 640：指定输入图像的大小为640x640像素。
• --batch 16：指定每个批次的图像数量为16。
• --epochs 100：指定训练轮数为100轮。

5.模型评估与调优

训练完成后，我们可以评估模型的性能。我们可以使用验证集上的准确率、召回率、F1分数和mAP（平均精度）等指标来评估模型性能。若性能不佳，可以尝试更改模型结构、使用更大的预训练模型或增加训练数据。

6.系统部署与实时检测

训练好的YOLOv5模型可以部署到服务器、工业设备或移动设备上。在部署过程中，我们需要注意处理实时视频流、图像预处理和后处理等方面的问题。以下是一个简单的实时检测示例代码：

import cv2
import torch
from PIL import Image

# Load the trained model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', source='local')

# Initialize the camera
cap = cv2.VideoCapture(0)

while True:
    # Capture frame-by-frame
    ret, frame = cap.read()

    # Convert the frame to PIL image
    pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

    # Perform the inference
    results = model(pil_img)

    # Draw the results on the frame
    frame = results.render()

    # Display the resulting frame
    cv2.imshow('frame', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# When everything is done, release the capture
cap.release()
cv2.destroyAllWindows()

 

这段代码首先加载训练好的模型，然后使用OpenCV捕获摄像头视频流。对于每一帧，我们将其转换为PIL图像，并使用YOLOv5模型进行推理。最后，我们将检测结果绘制在帧上并显示。

7.总结

本文介绍了如何使用YOLOv5构建一个实时安全帽检测和识别系统。我们讨论了数据集准备、模型训练、评估、调优、部署和实时检测等关键步骤。通过这个系统，我们可以有效地检测工地上工人是否佩戴了安全帽，从而提高安全意识和预防安全事故的发生。