基于YOLOv8的摄像头吸烟检测系统,提供数据集源码+pyqt系统

在安防和公共环境监控领域,吸烟行为的检测至关重要。今天就来跟大家分享基于YOLOv8搭建摄像头吸烟检测系统,还会附上数据集源码以及PyQt系统相关内容。

YOLOv8 简介

YOLO(You Only Look Once)系列一直是目标检测领域的佼佼者,YOLOv8 更是在速度和精度上有了进一步提升。它基于之前版本的架构进行优化,能够快速准确地识别图像和视频流中的目标物体。

数据集准备

要训练一个准确的吸烟检测模型,优质的数据集是关键。下面展示一些简单的数据处理代码示例:

import os
import shutil

# 假设原始数据集路径
original_dataset_path = 'original_smoking_dataset'
# 划分后的训练集、验证集、测试集路径
train_path = 'train_smoking_dataset'
val_path = 'val_smoking_dataset'
test_path = 'test_smoking_dataset'

# 创建目录
os.makedirs(train_path, exist_ok=True)
os.makedirs(val_path, exist_ok=True)
os.makedirs(test_path, exist_ok=True)

# 简单划分数据集,假设按 8:1:1 划分
file_list = os.listdir(original_dataset_path)
total_files = len(file_list)
train_files = file_list[:int(total_files * 0.8)]
val_files = file_list[int(total_files * 0.8):int(total_files * 0.9)]
test_files = file_list[int(total_files * 0.9):]

# 移动文件到相应目录
for file in train_files:
    shutil.copy(os.path.join(original_dataset_path, file), os.path.join(train_path, file))
for file in val_files:
    shutil.copy(os.path.join(original_dataset_path, file), os.path.join(val_path, file))
for file in test_files:
    shutil.copy(os.path.join(original_dataset_path, file), os.path.join(test_path, file))

这段代码主要完成了将原始数据集按比例划分为训练集、验证集和测试集,并分别存储到不同目录的操作。os.makedirs 用于创建目录,如果目录已存在则不会报错。通过计算文件总数并按比例切片,我们确定了每个子集包含的文件,最后用 shutil.copy 将文件复制到相应位置。

YOLOv8 训练

使用准备好的数据集,就可以开始 YOLOv8 的训练了。在终端中执行如下命令:

yolo detect train data=smoking_dataset.yaml model=yolov8n.pt epochs=50 imgsz=640

这里 data 参数指定数据集配置文件路径,model 选择预训练模型(这里以 yolov8n.pt 为例,是YOLOv8的轻量级模型),epochs 表示训练轮数,imgsz 是输入图像大小。训练过程中,模型会不断学习数据集中吸烟和非吸烟场景的特征。

摄像头实时检测

训练好模型后,就可以利用摄像头进行实时检测了。下面是Python代码示例:

import cv2
from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('best.pt')

# 打开摄像头
cap = cv2.VideoCapture(0)

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

    results = model(frame)
    annotated_frame = results[0].plot()

    cv2.imshow('Smoking Detection', annotated_frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

代码中,首先通过 YOLO('best.pt') 加载训练好的模型(best.pt 是训练后保存的最佳权重文件)。接着使用 cv2.VideoCapture(0) 打开默认摄像头。在循环中,不断读取摄像头的帧,将帧输入模型进行检测,model(frame) 返回检测结果。results[0].plot() 会在原帧上绘制检测到的目标框等信息。最后通过 cv2.imshow 显示处理后的帧,按下 q 键退出循环并释放摄像头资源、关闭窗口。

PyQt 系统集成

为了让吸烟检测系统有更友好的用户界面,我们可以使用 PyQt 进行集成。以下是一个简单的 PyQt 窗口示例代码:

import sys
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout
from PyQt5.QtGui import QPixmap, QImage
import cv2
from ultralytics import YOLO


class SmokingDetectionApp(QWidget):
    def __init__(self):
        super().__init__()
        self.model = YOLO('best.pt')
        self.cap = cv2.VideoCapture(0)
        self.initUI()

    def initUI(self):
        self.setWindowTitle('吸烟检测系统')
        self.setGeometry(100, 100, 800, 600)

        self.label = QLabel(self)
        self.label.setGeometry(0, 0, 800, 600)

        layout = QVBoxLayout()
        layout.addWidget(self.label)
        self.setLayout(layout)

        self.update_frame()

    def update_frame(self):
        ret, frame = self.cap.read()
        if ret:
            results = self.model(frame)
            annotated_frame = results[0].plot()

            height, width, channel = annotated_frame.shape
            bytes_per_line = 3 * width
            q_image = QImage(annotated_frame.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            pixmap = QPixmap.fromImage(q_image)
            self.label.setPixmap(pixmap)

        self.timer = self.startTimer(30)

    def timerEvent(self, event):
        self.update_frame()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = SmokingDetectionApp()
    ex.show()
    sys.exit(app.exec_())

这个PyQt代码创建了一个简单的窗口应用。SmokingDetectionApp 类继承自 QWidget,在初始化时加载模型并打开摄像头。initUI 方法设置窗口标题、大小,并添加一个 QLabel 用于显示视频帧。updateframe 方法从摄像头读取帧,进行检测并转换为 QPixmap 显示在 QLabel 上,同时通过 startTimer 定时更新画面。timerEvent 会在定时器触发时调用 updateframe 实现连续的视频流显示。

基于YOLOv8的摄像头吸烟检测系统,提供数据集源码+pyqt系统

通过上述步骤,我们成功搭建了基于YOLOv8的摄像头吸烟检测系统,从数据集处理到模型训练,再到实时检测以及PyQt系统集成,希望能给大家在相关项目开发上带来一些启发。

这样一套完整的系统源码和数据集,对于安防监控、公共场合禁烟管理等场景有着很高的实用价值,感兴趣的小伙伴可以根据实际需求进一步优化和扩展。

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐