基于YOLOv5深度学习的安全帽检测系统：YOLOv5 + 数据集 + UI界面

深度学习实战项目

于 2024-10-11 12:37:11 发布

阅读量332

点赞数 12

分类专栏： YOLOv10/v8/v7/v6/v5实战项目文章标签： YOLO 深度学习 ui 人工智能目标检测 yolov5 安全帽检测

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

版权

YOLOv10/v8/v7/v6/v5实战项目专栏收录该内容

180 篇文章 0 订阅

订阅专栏

引言

安全帽是建筑工地和工业现场中必不可少的个人防护装备，能够有效降低头部受伤的风险。为了确保工人的安全，及时检测佩戴安全帽的情况显得尤为重要。本文将详细介绍如何构建一个基于YOLOv5的安全帽检测系统，包括数据集准备、模型训练、UI界面设计及系统的实现。通过该系统，用户可以快速、准确地判断工人是否佩戴安全帽，从而提高工作场所的安全性。

1. 项目概述

本项目的主要目标是构建一个安全帽检测系统，能够自动识别是否佩戴安全帽。系统的主要组成部分包括：

数据集准备
YOLOv5模型训练
UI界面设计
系统实现与测试

2. 数据集准备

2.1 数据收集

构建安全帽检测系统的第一步是收集包含安全帽和未佩戴安全帽的图像数据集。数据集应包括不同环境下的工人图像，以增强模型的鲁棒性。

数据收集的方式可以包括：

手动拍摄：在工地或工业现场拍摄工人佩戴安全帽和未佩戴安全帽的图像。
公开数据集：使用Kaggle、Google Dataset Search等平台上的公开安全帽数据集。

示例数据集来源

Helmet Detection Dataset

2.2 数据标注

数据收集完成后，需要对图像进行标注。可以使用一些开源的标注工具（如LabelImg、Roboflow）来生成YOLO格式的标签文件。每个标签文件应与相应的图像文件同名，内容包括每个安全帽的类别和位置信息（边界框坐标）。

YOLO格式标签示例

0 0.5 0.5 0.4 0.4
1 0.7 0.3 0.2 0.3

每一行表示一个安全帽，格式为：

<类索引> <中心点x> <中心点y> <宽度> <高度>

2.3 数据集划分

将数据集划分为训练集和验证集，通常的划分比例为80%训练集和20%验证集。创建一个data.yaml文件，用于YOLOv5模型训练时使用，文件内容如下：

# data.yaml
train: ./data/train  # 训练集路径
val: ./data/val      # 验证集路径

nc: 1                # 类别数量
names: ['安全帽']    # 类别名称

3. YOLOv5模型训练

3.1 环境准备

确保你的计算环境中安装了所需的依赖项。可以使用以下命令创建一个Python虚拟环境并安装依赖项：

# 创建虚拟环境
python -m venv helmet_env
# 激活虚拟环境
# Windows
helmet_env\Scripts\activate
# macOS/Linux
source helmet_env/bin/activate

# 安装依赖
pip install torch torchvision torchaudio
pip install -U -r requirements.txt

3.2 下载YOLOv5模型

可以从YOLOv5的GitHub仓库克隆代码并进入该目录：

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

3.3 模型训练

使用以下命令启动模型训练：

python train.py --img 640 --batch 16 --epochs 100 --data data.yaml --weights yolov5s.pt --device 0

--img：输入图像大小
--batch：批处理大小
--epochs：训练轮数
--data：数据集配置文件路径
--weights：预训练权重文件路径
--device：使用的设备（GPU或CPU）

训练完成后，模型权重将保存在runs/train/exp/weights/目录下。

3.4 验证模型效果

使用以下命令验证模型的效果，查看在验证集上的表现：

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

3.5 模型推理

在训练完成后，可以对新的图像进行推理，以检测工人是否佩戴安全帽。使用以下命令进行推理：

python detect.py --weights runs/train/exp/weights/best.pt --source path/to/image.jpg --img 640

4. UI界面设计

为了便于用户使用系统，我们将构建一个简单的Web界面。使用Flask作为后端框架，HTML/CSS进行前端设计。

4.1 安装Flask

在虚拟环境中安装Flask：

pip install Flask

4.2 创建Flask应用

在项目根目录下创建一个名为app.py的Python文件，内容如下：

from flask import Flask, render_template, request
import cv2
import numpy as np
import torch

app = Flask(__name__)
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt')

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    file = request.files['file']
    img = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)

    # 模型推理
    results = model(img)
    result_img = np.array(results.render()[0])  # 获取绘制结果
    cv2.imwrite('static/output.jpg', result_img)  # 保存结果图像

    return render_template('result.html', filename='static/output.jpg')

if __name__ == '__main__':
    app.run(debug=True)

4.3 HTML页面设计

创建一个名为templates的文件夹，里面创建index.html和result.html两个文件。

`index.html`

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>安全帽检测</title>
</head>
<body>
    <h1>上传工人图像进行安全帽检测</h1>
    <form action="/upload" method="POST" enctype="multipart/form-data">
        <input type="file" name="file" accept="image/*" required>
        <button type="submit">检测</button>
    </form>
</body>
</html>

`result.html`

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>检测结果</title>
</head>
<body>
    <h1>检测结果</h1>
    <img src="{{ filename }}" alt="Detection Result">
    <br>
    <a href="/">返回</a>
</body>
</html>