基于YOLOv8的施工安全帽及安全背心检测

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

推荐：用 NSDT编辑器 快速搭建可编程3D场景

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

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()

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

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

---

原文链接：基于YOLOv8的安全帽检测 — BimAnt