基于YOLOv8的安全帽检测

目标检测是一项基本的计算机视觉任务。另一方面，YOLO（You Only Look Once）是一种流行的目标检测模型，以其速度和准确性而闻名。涉及对象检测的用例非常多样化。其中之一是建筑工地安全。建筑工地经理、安全官员或监管机构可以使用它来监控并确保工人遵守安全协议，例如穿戴适当的个人防护装备 (PPE：Personal Protective Equipment)。

在这篇博文中，我们将探讨如何应用YOLO算法实现工人是否使用口罩、安全背心和安全帽的检测。该模型可以成为监控和确保工人遵守安全协议的便捷方法。本文相关的源代码可以从这个Github仓库下载。

如果你缺少训练数据，可以使用UnrealSynth合成数据生成器，只需要将安全帽的3D模型导入UnrealSynth，就可以自动生成数据集，非常方便：

https://tools.nsdt.cloud/UnrealSynth

1、下载数据集

首先，我们可以从此链接下载数据集：

点击下载按钮后，选择YOLOv8格式：

在弹出的对话框中选中下载zip到本机，然后继续：

数据集下载完成后，现在我们用新数据集训练YOLO。就这么简单。

2、训练YOLOv8模型

打开Google Colab，可以将数据集上传到 Google Drive，然后在Colab中挂载。接下来我们按如下方式在 Google Colab 中训练模型。

在运行任何脚本之前，请确保下载正确的包。可以运行以下命令来执行此操作：

!pip install ultralytics
from ultralytics import YOLO
cd /content

该命令开始训练 YOLO：

!yolo task=detect mode=train model=yolov8l.pt data='/content/drive/MyDrive/YoloDataset/Construction Site Safety.v30-raw-images_latestversion.yolov8/data.yaml' epochs=10

注意：数据集位于我的 Google Drive 中的以下路径：drive/MyDrive/YoloDataset/Construction Site Safety.v30-raw-images_latestversion.yolov8

模型训练完成后，你可以从目录下载 best.pt，如下所示：

3、设置环境

在深入实施之前，需要先设置开发环境：

安装必要的库：首先安装所需的库，包括 ultralytics、YOLO、OpenCV 和 math ：

from ultralytics import YOLO
import cv2
import math
from helper import create_video_writer

注意：helper.py 仅用于以视频格式保存输出，可以从这个 github仓库下载。

4、Pycharm实现

查看选择视频 (ppe-3.mp4) 和 best.pt 。此实现的输出保存在 ConstructionSiteSafetyOutput.mp4 中。

cap = cv2.VideoCapture('videos/ppe-3.mp4')
writer = create_video_writer(cap, "ConstructionSiteSafetyOutput.mp4")
model = YOLO("best.pt")

这里我们有分类的名称，当下载数据集时，我们可以在 data.yaml 中找到它们。

classNames = ['Excavator', 'Gloves', 'Hardhat', 'Ladder', 'Mask', 'NO-Hardhat', 'NO-Mask', 'NO-Safety Vest',
              'Person', 'SUV', 'Safety Cone', 'Safety Vest', 'bus', 'dump truck', 'fire hydrant', 'machinery',
              'mini-van', 'sedan', 'semi', 'trailer', 'truck and trailer', 'truck', 'van', 'vehicle', 'wheel loader']

在继续执行过程中，如果检测到“NO-Hardhat”、“NO-Safety Vest”和“NO-Mask”三个类别中的任何一个，它将以红色绘制边界框，但如果工人使用口罩、安全帽或安全背心，绘制绿色边界框。

myColor = (0, 0, 255)
while True:
    success, img = cap.read()
    results = model(img, stream=True)
    for r in results:
        boxes = r.boxes
        for box in boxes:
            # Bounding Box
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            w, h = x2 - x1, y2 - y1

            # Confidence
            conf = math.ceil((box.conf[0] * 100)) / 100
            # Class Name
            cls = int(box.cls[0])
            currentClass = classNames[cls]
            print(currentClass)
            if conf > 0.5:
                if currentClass =='NO-Hardhat' or currentClass =='NO-Safety Vest' or currentClass == "NO-Mask":
                    myColor = (0, 0,255)
                elif currentClass =='Hardhat' or currentClass =='Safety Vest' or currentClass == "Mask":
                    myColor =(0,255,0)
                else:
                    myColor = (255, 0, 0)


                image = cv2.putText(img, f'{classNames[cls]}', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX,
                                    1, (255, 0, 0), 2, cv2.LINE_AA)
                cv2.rectangle(img, (x1, y1), (x2, y2), myColor, 3)

    cv2.imshow("Image", img)
    writer.write(img)
    if cv2.waitKey(1) == ord("q"):
        break
cap.release()
writer.release()
cv2.destroyAllWindows()

图像格式的输出结果如下图所示：

视频格式的输出结果可以查看这个视频：

原文链接：How to check construction safety with YOLO model and Python?

BimAnt翻译整理，转载请标明出处