基于深度学习的安全帽检测系统:YOLOv5/v6/v7/v8/v10模型实现与UI界面集成、数据集-CSDN博客

本文链接：https://blog.csdn.net/m0_52343631/article/details/142665368

1. 引言

在各类建筑和工业现场，安全帽是保障工人安全的基本防护装备。尽管安全帽的佩戴在一定程度上提高了工人的安全性，但由于人为因素，仍然可能出现工人未佩戴安全帽的情况。因此，开发一个自动检测安全帽佩戴状态的系统显得尤为重要。本博客将详细介绍如何构建一个基于深度学习的安全帽检测系统，使用YOLO系列模型（YOLOv5、YOLOv6等）来实现对安全帽佩戴的检测，并提供完整的代码实现、训练数据集配置及用户界面的设计。

4.3 配置文件 hardhat_dataset.yaml

6.2 hardhat_detection_ui.py 文件

2. 项目概述

本项目的目标是实现一个安全帽检测系统，具体功能包括：

图像上传和检测：用户可以上传图像，系统会检测安全帽的佩戴情况。
实时视频检测：通过摄像头实时监控工人是否佩戴安全帽。
友好的用户界面：使用Tkinter构建图形用户界面，方便用户操作。

2.1 系统架构

系统架构如下：

+---------------------+
|    用户界面         |
| (Tkinter GUI)      |
+---------------------+
           |
           v
+---------------------+
|     YOLO模型        |
| (安全帽检测模型)   |
+---------------------+
           |
           v
+---------------------+
|   图像处理模块      |
| (图像上传与显示)    |
+---------------------+
           |
           v
+---------------------+
|    数据集           |
| (安全帽图像数据)   |
+---------------------+

3. 项目结构

项目的文件结构如下：

hardhat_detection/
├── dataset/
│   ├── images/
│   │   ├── train/
│   │   └── val/
│   ├── labels/
│   │   ├── train/
│   │   └── val/
│   └── hardhat_dataset.yaml
├── yolov5/
│   ├── ... # YOLOv5相关文件
├── hardhat_detection_ui.py
└── requirements.txt

4. 数据集准备

4.1 数据集收集

您可以通过公开数据集（如Kaggle的安全帽检测数据集）或自行拍摄来收集数据。数据集应包括佩戴安全帽和未佩戴安全帽的图像。

4.2 数据集标注

使用工具（如LabelImg）对图像进行标注。标注完成后，生成YOLO格式的标签文件，文件名与图像名相同，扩展名为.txt，内容格式如下：

<class_id> <x_center> <y_center> <width> <height>

这里的<class_id>是类别的索引，YOLO格式的标签文件中坐标的单位是相对图像尺寸的比例。

4.3 配置文件 `hardhat_dataset.yaml`

该文件用于配置数据集的路径和类别信息，内容如下：

# hardhat_dataset.yaml
train: ./images/train
val: ./images/val

nc: 2
names: ['no_hardhat', 'hardhat']

4.4 数据集示例

训练集：包含各种环境下工人佩戴和未佩戴安全帽的图像。
验证集：用于模型验证，确保模型在未见数据上的表现。

5. YOLO模型训练

5.1 安装YOLOv5

首先，克隆YOLOv5代码库：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

5.2 训练模型

在YOLOv5目录下，使用以下命令进行模型训练：

python train.py --img 640 --batch 16 --epochs 50 --data ../dataset/hardhat_dataset.yaml --weights yolov5s.pt

参数说明：

--img 640：输入图像大小为640x640。
--batch 16：批量大小为16。
--epochs 50：训练轮数为50。
--data：数据集配置文件路径。
--weights：初始权重文件（可以选择预训练模型）。

5.3 训练过程监控

训练过程中，YOLOv5会输出训练日志，包括损失值、精度等信息。您可以使用TensorBoard进行可视化：

tensorboard --logdir runs/train

5.4 模型评估

训练完成后，您可以使用以下命令评估模型的性能：

python val.py --weights runs/train/exp/weights/best.pt --data ../dataset/hardhat_dataset.yaml --img 640

6. 用户界面设计

6.1 设计思路

用户界面采用Tkinter进行设计，提供上传图片和摄像头实时检测的功能，界面简洁直观，易于操作。

6.2 `hardhat_detection_ui.py` 文件

以下是用户界面的完整代码。此代码使用Tkinter创建一个简单的GUI，用户可以通过上传图片或使用摄像头进行检测。

import tkinter as tk
from tkinter import filedialog
import torch
import cv2
import numpy as np
from PIL import Image, ImageTk

class HardHatDetectionApp:
    def __init__(self, root, model_path):
        self.root = root
        self.root.title("安全帽检测系统")
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path, force_reload=True)

        # 创建画布
        self.canvas = tk.Canvas(root, width=640, height=480)
        self.canvas.pack()

        # 创建按钮
        self.upload_button = tk.Button(root, text="上传图片", command=self.upload_image)
        self.upload_button.pack(side=tk.LEFT)

        self.camera_button = tk.Button(root, text="摄像头检测", command=self.detect_from_camera)
        self.camera_button.pack(side=tk.LEFT)

        # 创建结果标签
        self.result_label = tk.Label(root, text="检测结果:")
        self.result_label.pack()

    def upload_image(self):
        file_path = filedialog.askopenfilename()
        if file_path:
            img = Image.open(file_path)
            img = img.resize((640, 480))
            self.img_tk = ImageTk.PhotoImage(img)
            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.img_tk)

            # YOLO模型检测
            self.perform_detection(img)

    def detect_from_camera(self):
        cap = cv2.VideoCapture(0)

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            # YOLO模型检测
            results = self.model(frame)

            # 将检测结果绘制在帧上
            self.draw_boxes(frame, results)

            # 将BGR转为RGB以适应PIL图像
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img_pil = Image.fromarray(frame)
            img_tk = ImageTk.PhotoImage(img_pil)
            self.canvas.create_image(0, 0, anchor=tk.NW, image=img_tk)

            # 更新界面
            self.root.update()

        cap.release()

    def perform_detection(self, img):
        results = self.model(np.array(img))
        detection_results = results.pandas().xyxy[0]
        self.result_label.config(text=f"检测结果:\n{detection_results}")

    def draw_boxes(self, frame, results):
        for i in range(len(results.xyxy[0])):
            x1, y1, x2, y2, conf, cls = results.xyxy[0][i]
            label = f"{self.model.names[int(cls)]} {conf:.2f}"
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            cv2.putText(frame, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 启动应用
if __name__ == "__main__":
    root = tk.Tk()
    app = HardHatDetectionApp(root, model_path='runs/train/exp/weights/best.pt')  # 请根据您的模型路径修改
    root.mainloop()

6.3 代码说明

Tkinter Canvas: 用于显示图像或视频流。
上传图片: 用户通过文件对话框选择图片，程序读取并处理该图片。
摄像头检测: 实时捕捉摄像头画面，并使用YOLO模型检测安全帽状态。
检测结果: 将检测框和标签绘制在视频流或图片上。

7. 部署与测试

7.1 部署

环境准备：确保安装了必要的Python库，可以使用requirements.txt安装：
```
pip install -r requirements.txt
```
运行程序：在项目根目录下运行用户界面脚本：
```
python hardhat_detection_ui.py
```