基于YOLOv8的安全帽检测系统

作者

本文将介绍如何利用YOLOv8（You Only Look Once）实现安全帽检测系统。整个流程包括环境配置、数据集标定、模型训练、测试与验证，并通过算法优化提升检测精度。

ps: 本文为自用学习资料，仅在此网站分享，仅供学习参考，未经允许，不得转载。本文所参考的资料文献将在文末、文中列出。

环境配置与数据集标定

在进行YOLO模型训练前，首先需要配置环境。若环境未配置，您可以参考下方链接进行配置。接下来，我们使用 labelimg 工具进行数据集的标注和处理。

零基础教程：使用yolov8训练自己的目标检测数据集_yolov8训练数据需要多少-CSDN博客

标定数据集

标定是将目标框选并标注类别。以安全帽为例，我们需要在每张图片中框出安全帽的位置。多个目标需要分别标注，并确保每个目标都被框出，避免遗漏。（eg: 提供100张图片就要保证每张图片都被标注到）

一旦完成标定后，需要将数据集划分为训练集、验证集和测试集。以下是数据划分的Python代码：

import os
import random
import shutil

# 设置文件路径和划分比例
root_path = "./voc_yolo/"
image_dir = "./JPEGImages/"
label_dir = "./Annotations/"
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

# 创建训练集、验证集和测试集目录
os.makedirs("images/train", exist_ok=True)
os.makedirs("images/val", exist_ok=True)
os.makedirs("images/test", exist_ok=True)
os.makedirs("labels/train", exist_ok=True)
os.makedirs("labels/val", exist_ok=True)
os.makedirs("labels/test", exist_ok=True)

# 获取所有图像文件名
image_files = os.listdir(image_dir)
total_images = len(image_files)
random.shuffle(image_files)

# 计算划分数量
train_count = int(total_images * train_ratio)
val_count = int(total_images * val_ratio)
test_count = total_images - train_count - val_count

# 划分训练集
train_images = image_files[:train_count]
for image_file in train_images:
    label_file = image_file[:image_file.rfind(".")] + ".txt"
    shutil.copy(os.path.join(image_dir, image_file), "images/train/")
    shutil.copy(os.path.join(label_dir, label_file), "labels/train/")

# 划分验证集
val_images = image_files[train_count:train_count+val_count]
for image_file in val_images:
    label_file = image_file[:image_file.rfind(".")] + ".txt"
    shutil.copy(os.path.join(image_dir, image_file), "images/val/")
    shutil.copy(os.path.join(label_dir, label_file), "labels/val/")

# 划分测试集
test_images = image_files[train_count+val_count:]
for image_file in test_images:
    label_file = image_file[:image_file.rfind(".")] + ".txt"
    shutil.copy(os.path.join(image_dir, image_file), "images/test/")
    shutil.copy(os.path.join(label_dir, label_file), "labels/test/")

# 生成训练集图片路径txt文件
with open("train.txt", "w") as file:
    file.write("\n".join([root_path + "images/train/" + image_file for image_file in train_images]))

# 生成验证集图片路径txt文件
with open("val.txt", "w") as file:
    file.write("\n".join([root_path + "images/val/" + image_file for image_file in val_images]))

# 生成测试集图片路径txt文件
with open("test.txt", "w") as file:
    file.write("\n".join([root_path + "images/test/" + image_file for image_file in test_images]))

print("数据划分完成！")

关于数据集的划分：按7：2：1划分

训练集（Training Set）： 训练集是用于训练模型的数据集合。目标检测模型通过在训练集中的图像中进行学习，并优化其参数，从而识别目标的特征和背景信息。训练集通常包含大量标注数据，其中目标物体的位置信息和类别都已明确标记。模型从这些样本中学习目标的外观和形态，进而能够对新图像进行目标检测。

验证集（Validation Set）： 验证集的主要作用是在训练过程中评估模型的性能并调整超参数。它是从与训练集不同的图像中随机选取的一部分数据。通过在验证集上测试，可以实时监控模型训练的进展，并进行超参数的调整，如学习率和正则化参数等。验证集上的结果能帮助确定最佳模型，并防止出现过拟合现象。验证集的标注数据也包括目标物体的位置信息和类别。

测试集（Test Set）： 测试集用于在训练完成后对模型进行最终评估。它与训练集和验证集完全不同，是用于衡量训练好的目标检测模型在新数据上的表现的独立数据集。测试集中的图像对于模型而言是全新的，模型会在这些图像上生成目标检测的预测结果。测试集也包含目标物体的位置和类别标签。通过比较模型在测试集上的表现，可以评估其准确率、召回率、精确率等指标，从而了解模型在实际应用中的泛化能力和有效性。

这三个数据集在目标检测训练中各自有着至关重要的作用。训练集负责模型的学习和参数调整，验证集用于超参数的优化和模型选择，测试集则用于最终的性能评估和泛化能力测试。它们的目的是确保训练出的模型能够准确地在新的图像中检测到目标物体。

将数据集放到下述位置：

配置训练用的 YAML 文件

在YOLO中，数据集的配置文件采用YAML格式。例如，创建 helmet.yaml 文件以供训练使用。在该文件中定义了目标类别及其名称，如何创建？

如图找到coco128文件，复制并在同路径中粘贴一个新的配置文件命名helmet.yaml，这里存的就是一会儿训练用的配置文件：

训练模型

训练模型的python代码如下 （需要新建train.py并放到正确位置，见上图共用）：

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('path_to_your_model/yolov8n.pt')

# 使用自定义数据集进行训练
results = model.train(
    data='path_to_your_dataset/helmet.yaml', 
    epochs=30, 
    imgsz=640, 
    device=[0], 
    workers=0, 
    lr0=0.01, 
    batch=8, 
    amp=False
)

参数说明：

epochs (epochs=30)

  • 训练的轮数，即模型将在整个训练集上训练多少次。一般10到50次，取决于数据集的多少，一般数据集少训练次数少；数据集多训练次数多。每次训练都会对模型参数进行更新，一般要识别的目标对象是多个目标时通常选择较大的值以确保模型有足够的时间学习（30到50次）。

imgsz (imgsz=640)

  • 输入图像的大小（以像素为单位）。这是训练期间图像会被缩放的大小。对于YOLO模型，常用的尺寸包括416、640和1280等。较大的图像尺寸通常会导致更高的计算需求，但也可能带来更精细的目标检测性能。

device (device=[0])

  • 指定训练所使用的硬件设备。在这里，[0] 表示使用第一个GPU设备进行训练。如果你有多个GPU，可以指定多个设备，如[0, 1]。如果没有GPU，也可以设置为 cpu，即仅使用CPU进行训练，将device=[0]改为device=’cpu’即可

workers (workers=0)

  • 用于数据加载的子进程数量。设置为0时表示不使用任何额外的进程加载数据。如果设置为更大的数字，将会加速数据加载，尤其是当训练数据集非常大时，通常推荐设置为一个合理的值（如4或8）。

lr0 (lr0=0.01) 学习率

  • 初始学习率，控制优化器在每次参数更新时步长的大小。较大的学习率可以使模型更快收敛，但可能导致不稳定的训练过程。较小的学习率可能会带来更慢的收敛速度，但通常更稳定。同时如果学习率变小，批次轮次也应该同样增加。学习率区间为0到1，一般设置为0.01到0.05.
  • 有关学习率，详情参考以下两个链接：
  • 深度学习笔记(五)：学习率过大过小对于网络训练有何影响以及如何解决-CSDN博客
  • https://www.bilibili.com/video/BV1d8411E77J/?spm_id_from=333.337.search-card.all.click&vd_source=2d33bb0cc56df660aef4d5b952259a01

batch (batch=8)

  • 每次训练时输入到网络的样本数量（批次大小）。批次大小越大，训练过程中每次更新模型参数时计算的数据量就越多，这通常有助于稳定训练过程，但也会增加内存消耗。较小的批次大小会加快训练，但可能导致不稳定的梯度更新。

amp (amp=False)

  • 这个参数决定是否使用自动混合精度训练（Automatic Mixed Precision）。启用AMP能够在不牺牲训练精度的情况下减少显存的占用，从而加速训练过程。通常在使用GPU时，开启AMP会带来更高的性能。如果你的显卡支持FP16精度计算，可以开启该选项。设置为False则表示禁用。

训练完成后，会生成runs，而新训练的模型将在 weights 目录下的 best.pt，即为我们在后续使用的模型。

测试与验证

一旦模型训练完成，就可以对新图像进行检测。这里使用Claude 生成了制作展示测试界面的窗口，附上代码如下：

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk
import cv2
from ultralytics import YOLO
from datetime import datetime

class HelmetDetectorGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("安全帽检测系统")
        self.root.geometry("1200x800")

        # 初始化变量
        self.model = None
        self.current_image = None
        self.photo_image = None
        self.is_model_loaded = False

        self.create_gui()

    def create_gui(self):
        # 创建控制面板和图像显示区
        pass

    def load_model(self):
        pass

    def select_image(self):
        pass

    def display_image(self, img):
        pass

    def detect_objects(self):
        pass

    def save_result(self):
        pass

def main():
    root = tk.Tk()
    app = HelmetDetectorGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()

此GUI程序通过 tkinter 创建一个简单的界面，让用户加载训练好的模型，选择图片并进行安全帽检测。

总结与展望

通过以上步骤，我们可以构建一个基于YOLOv8的安全帽检测系统。尽管模型可能无法达到100%的识别精度，进一步的优化和高级算法可以帮助提高识别效果与置信度。

参考资料网站：

零基础教程：使用yolov8训练自己的目标检测数据集_yolov8训练数据需要多少-CSDN博客

深度学习笔记(五)：学习率过大过小对于网络训练有何影响以及如何解决-CSDN博客

实战 | YOLOv8自定义数据集训练实现手势识别 (标注+训练+预测 保姆级教程)

推荐公众号：OpenCV 与 AI深度学习

发表评论