实时吸烟检测系统：基于深度学习与YOLO模型的完整实现

原创

A等天晴 2024-07-25 10:38:18 博主文章分类：计算机视觉 ©著作权

文章标签 深度学习 YOLO 人工智能 ide v8 文章分类 JavaScript 前端开发

©著作权归作者所有：来自51CTO博客作者A等天晴的原创作品，请联系作者获取转载授权，否则将追究法律责任

引言

在公共场所吸烟不仅有害吸烟者的健康，也会影响到周围的人。为了更好地管理和监控公共场所的吸烟行为，本文介绍了一种基于深度学习的吸烟检测系统。该系统包含一个用户界面，利用YOLO（You Only Look Once）v8/v7/v6/v5模型进行吸烟检测，并提供了完整的实现步骤和代码。

系统架构

数据收集和预处理
模型训练
系统实现
用户界面设计

数据收集和预处理

数据集

为了训练一个高精度的吸烟检测模型，我们需要一个包含大量吸烟和非吸烟图片的数据集。可以使用以下途径收集数据：

公开数据集：如Kaggle上的相关数据集。
自定义数据集：通过互联网、监控摄像头等途径收集图片。

数据标注

使用工具如LabelImg对数据进行标注。标注内容包括吸烟者的位置（bounding box）和标签（吸烟/非吸烟）。

# 训练数据集文件结构示例
dataset/
  ├── images/
  │   ├── train/
  │   └── val/
  └── labels/
      ├── train/
      └── val/

模型训练

YOLO模型有多个版本，本文选取YOLOv8作为示范，其他版本可以通过相似方法实现。

环境配置

安装YOLOv8所需的依赖库：

pip install ultralytics

模型训练代码

使用YOLOv8进行模型训练，假设数据已经按照YOLO的格式进行预处理和标注。

from ultralytics import YOLO

# 加载预训练的YOLOv8模型
model = YOLO('yolov8.yaml')

# 配置训练参数
model.train(data='path/to/dataset.yaml', epochs=50, imgsz=640, batch=16)

# 保存训练后的模型
model.save('best.pt')

系统实现

吸烟检测

利用训练好的模型进行吸烟检测，并实现视频流的实时检测。

import cv2
from ultralytics import YOLO

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

# 打开视频流
cap = cv2.VideoCapture('path/to/video.mp4')

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

    # 检测吸烟行为
    results = model(frame)
    for result in results:
        bbox = result['bbox']
        label = result['label']
        confidence = result['confidence']
        
        # 画框和标签
        cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
        cv2.putText(frame, f'{label} {confidence:.2f}', (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
    # 显示视频
    cv2.imshow('Smoke Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

用户界面设计

用户界面采用PyQt5实现，提供视频播放和吸烟检测结果显示。

安装PyQt5

pip install PyQt5

界面代码

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

class SmokeDetectionUI(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        
        self.model = YOLO('best.pt')
        
    def initUI(self):
        self.setWindowTitle('Smoke Detection System')
        
        self.layout = QVBoxLayout()
        
        self.label = QLabel(self)
        self.layout.addWidget(self.label)
        
        self.button = QPushButton('Open Video', self)
        self.button.clicked.connect(self.open_video)
        self.layout.addWidget(self.button)
        
        self.setLayout(self.layout)
    
    def open_video(self):
        options = QFileDialog.Options()
        video_path, _ = QFileDialog.getOpenFileName(self, "Open Video", "", "All Files (*);;MP4 Files (*.mp4)", options=options)
        
        if video_path:
            self.detect_smoke(video_path)
    
    def detect_smoke(self, video_path):
        cap = cv2.VideoCapture(video_path)
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            
            results = self.model(frame)
            for result in results:
                bbox = result['bbox']
                label = result['label']
                confidence = result['confidence']
                
                cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
                cv2.putText(frame, f'{label} {confidence:.2f}', (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            
            # 将frame转换为QImage
            height, width, channel = frame.shape
            bytesPerLine = 3 * width
            qImg = QImage(frame.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
            
            self.label.setPixmap(QPixmap.fromImage(qImg))
            cv2.waitKey(1)
        
        cap.release()

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