基于深度学习的交通信号灯检测系统（UI界面+YOLOv8/v7/v6/v5代码+训练数据集）
引言
交通信号灯检测是智能交通系统中不可或缺的一部分。它不仅能够提高自动驾驶汽车的安全性，还可以优化交通管理系统的效率。本文将详细介绍如何构建一个基于深度学习的交通信号灯检测系统，涵盖从环境搭建、数据准备、模型训练、系统实现到用户界面设计的完整流程，并提供实际的代码示例。
系统概述
本系统实现的主要步骤包括：
- 环境搭建
- 数据收集与处理
- 模型训练
- 系统实现
- 用户界面设计
环境搭建
首先，搭建一个适合深度学习开发的环境。本系统使用Python 3.8或以上版本，并依赖于多个深度学习和图像处理库。
安装必要的库
使用以下命令安装所需库：
pip install numpy pandas matplotlib opencv-python torch torchvision ultralytics pyqt5
数据收集与处理
数据收集
收集包含交通信号灯的图像数据集。可以使用公开的数据集，如LISA Traffic Light Dataset，或者通过摄像头自行采集。确保数据集包含各种不同状态（红灯、黄灯、绿灯）、不同角度、不同光照条件下的信号灯图像。
数据处理
将图像数据整理到指定的文件夹结构，并标注信号灯的位置。以下是示例的文件夹结构：
datasets/
├── images/
│ ├── train/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
│ ├── val/
│ │ ├── image1.jpg
│ │ ├── image2.jpg
├── labels/
├── train/
│ ├── image1.txt
│ ├── image2.txt
├── val/
├── image1.txt
├── image2.txt
每个标签文件的内容如下：
class x_center y_center width height
其中，class表示类别编号（红灯、黄灯、绿灯），x_center、y_center为归一化后的中心坐标，width和height为归一化后的宽度和高度。
数据增强
为了提高模型的泛化能力，可以对数据进行增强处理。常见的数据增强方法包括旋转、缩放、平移、镜像等。
import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
import os
# 定义数据增强变换
transform = A.Compose([
A.RandomRotate90(),
A.Flip(),
A.Transpose(),
A.OneOf([
A.IAAAdditiveGaussianNoise(),
A.GaussNoise(),
], p=0.2),
A.OneOf([
A.MotionBlur(p=0.2),
A.MedianBlur(blur_limit=3, p=0.1),
A.Blur(blur_limit=3, p=0.1),
], p=0.2),
A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
A.OneOf([
A.OpticalDistortion(p=0.3),
A.GridDistortion(p=0.1),
A.IAAPiecewiseAffine(p=0.3),
], p=0.2),
A.OneOf([
A.CLAHE(clip_limit=2),
A.IAASharpen(),
A.IAAEmboss(),
A.RandomBrightnessContrast(),
], p=0.3),
A.HueSaturationValue(p=0.3),
ToTensorV2()
])
# 处理图像
def augment_images(image_folder, output_folder):
os.makedirs(output_folder, exist_ok=True)
for filename in os.listdir(image_folder):
image_path = os.path.join(image_folder, filename)
image = cv2.imread(image_path)
augmented = transform(image=image)
augmented_image = augmented["image"].numpy().transpose(1, 2, 0)
output_path = os.path.join(output_folder, filename)
cv2.imwrite(output_path, augmented_image)
# 应用数据增强
augment_images('datasets/images/train', 'datasets/images/train_augmented')
模型训练
使用YOLO模型进行训练。
配置文件
创建一个配置文件config.yaml：
path: datasets
train: images/train
val: images/val
test: images/test
nc: 3 # 类别数，红灯、黄灯、绿灯
names: ['red', 'yellow', 'green'] # 类别名称
训练代码
使用以下代码训练模型：
from ultralytics import YOLO
# 加载模型
model = YOLO('yolov8n.pt')
# 训练模型
model.train(data='config.yaml', epochs=50, imgsz=640, batch=16, lr0=0.01)
训练过程会生成模型权重文件best.pt，该文件包含了训练好的模型参数。
系统实现
训练好的模型可以用于实时交通信号灯检测。使用OpenCV读取视频流，并调用YOLO模型进行检测。
检测代码
import cv2
from ultralytics import YOLO
# 加载训练好的模型
model = YOLO('best.pt')
# 打开视频流
cap = cv2.VideoCapture('traffic_light_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('Traffic Light Detection', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
用户界面设计
为了提高系统的易用性，我们设计了一个用户友好的界面。使用PyQt5实现用户界面，提供图像或视频播放和检测结果显示。
界面代码
以下是一个简单的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 TrafficLightDetectionUI(QWidget):
def __init__(self):
super().__init__()
self.initUI()
self.model = YOLO('best.pt')
def initUI(self):
self.setWindowTitle('Traffic Light Detection System')
self.layout = QVBoxLayout()
self.label = QLabel(self)
self.layout.addWidget(self.label)
self.button = QPushButton('Open Image or Video', self)
self.button.clicked.connect(self.open_file)
self.layout.addWidget(self.button)
self.setLayout(self.layout)
def open_file(self):
options = QFileDialog.Options()
file_path, _ = QFileDialog.getOpenFileName(self, "Open File", "", "All Files (*);;MP4 Files (*.mp4);;JPEG Files (*.jpg);;PNG Files (*.png)", options=options)
if file_path:
if file_path.endswith('.mp4'):
self.detect_traffic_light_video(file_path)
else:
self.detect_traffic_light_image(file_path)
def detect_traffic_light_image(self, file_path):
frame = cv2.imread(file_path)
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)
height, width, channel = frame.shape
bytesPerLine = 3 * width
qImg = QImage(frame.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
self.label.setPixmap(QPixmap.from
Image(qImg))
def detect_traffic_light_video(self, file_path):
cap = cv2.VideoCapture(file_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)
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))
QApplication.processEvents()
cap.release()
if __name__ == '__main__':
app = QApplication(sys.argv)
ex = TrafficLightDetectionUI()
ex.show()
sys.exit(app.exec_())
总结
本文介绍了如何构建一个基于深度学习的交通信号灯检测系统，从环境搭建、数据收集与处理、模型训练、系统实现到用户界面设计。通过步骤和代码示例，可以轻松地实现自己的交通信号灯检测系统。这不仅可以应用于自动驾驶，还可以为交通管理和安全监控提供有力的支持。
声明：本文只是简单的项目思路，如有部署的想法，想要（UI界面+YOLOv8/v7/v6/v5代码+训练数据集）的可以联系作者.