基于YOLOv8/YOLOv7/YOLOv6/YOLOv5的交通标志识别系统详解（深度学习模型+UI界面代码+训练数据集）

摘要

交通标志识别是智能交通系统和自动驾驶领域的关键技术之一。本文详细介绍了一个基于YOLO系列算法（YOLOv8、v7、v6、v5）的交通标志识别系统，涵盖深度学习模型原理、UI界面设计、训练数据集构建以及完整代码实现。通过对比不同版本YOLO算法的性能表现，我们展示了如何构建一个高效、准确的交通标志识别系统，并提供了完整的实现代码和训练数据集参考。

关键词：YOLO，交通标志识别，深度学习，计算机视觉，目标检测，智能交通

目录

摘要

1. 引言

1.1 研究背景与意义

1.2 研究现状

2. 系统设计与架构

2.1 系统总体架构

2.2 技术栈

3. 数据集构建与处理

3.1 数据集选择与准备

3.2 数据预处理流程

3.3 数据标注格式

4. YOLO模型详解与实现

4.1 YOLOv5模型架构

4.2 YOLOv8改进特性

4.3 损失函数设计

5. 模型训练与优化

5.1 训练配置

5.2 训练脚本

5.3 超参数优化

6. 用户界面设计

6.1 PyQt5界面实现

6.2 Gradio Web界面

7. 模型评估与对比

7.1 评估指标

7.2 YOLO版本对比实验

8. 部署与优化

8.1 模型量化

8.2 TensorRT加速

1. 引言

1.1 研究背景与意义

随着智能交通系统（ITS）和自动驾驶技术的快速发展，交通标志识别（Traffic Sign Recognition, TSR）已成为车辆环境感知系统的核心组成部分。准确、实时的交通标志识别对于保障道路安全、提高交通效率具有重要意义。

传统的交通标志识别方法主要基于颜色、形状等手工特征，但这些方法在复杂环境（如光照变化、遮挡、天气条件变化）下表现不佳。近年来，基于深度学习的交通标志识别方法展现出显著优势，尤其是YOLO（You Only Look Once）系列算法以其高速度和高精度成为目标检测领域的主流选择。

1.2 研究现状

YOLO系列算法自2016年首次提出以来，经历了多次迭代升级：

• YOLOv5：采用PyTorch框架，简化了模型部署

• YOLOv6：专注于工业应用优化

• YOLOv7：在速度和精度上达到新平衡

• YOLOv8：最新版本，提供更优的性能和灵活性

交通标志识别领域的公开数据集主要包括德国交通标志检测基准（GTSDB）、德国交通标志识别基准（GTSRB）、比利时交通标志数据集（BTSD）等。这些数据集为算法训练和评估提供了标准基准。

2. 系统设计与架构

2.1 系统总体架构

我们的交通标志识别系统采用模块化设计，主要包括以下组件：

1. 数据采集与预处理模块：负责图像采集、标注和数据增强

2. 深度学习模型模块：基于YOLO系列算法的目标检测模型

3. 模型训练与优化模块：训练策略、超参数优化

4. 推理与部署模块：模型量化、加速和部署

5. 用户界面模块：基于PyQt5/Gradio的交互界面

2.2 技术栈

• 深度学习框架：PyTorch 1.12+

• 编程语言：Python 3.8+

• 可视化库：OpenCV, Matplotlib

• UI框架：PyQt5 或 Gradio

• 开发环境：CUDA 11.3+, cuDNN 8.2+

3. 数据集构建与处理

3.1 数据集选择与准备

我们选择以下公开数据集作为基础：

1. GTSDB（German Traffic Sign Detection Benchmark）：

   • 包含900张图像，43类交通标志

   • 图像尺寸：1360×800像素

   • 训练集/测试集：600/300

2. TT100K（Tsinghua-Tencent 100K）：

   • 包含10万张图像，221类交通标志

   • 中国道路场景，更符合实际应用需求

3. 自定义数据集扩展：

   • 补充收集了不同天气、光照条件下的交通标志图像

   • 使用数据增强技术扩充数据集规模

3.2 数据预处理流程

python

import cv2 import numpy as np import albumentations as A from albumentations.pytorch import ToTensorV2 class TrafficSignDataset: def __init__(self, image_paths, labels, transform=None): self.image_paths = image_paths self.labels = labels self.transform = transform # 定义数据增强管道 self.train_transform = A.Compose([ A.Resize(640, 640), A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.RGBShift(p=0.2), A.Blur(blur_limit=3, p=0.1), A.MedianBlur(blur_limit=3, p=0.1), A.ToGray(p=0.05), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2() ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'])) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image = cv2.imread(self.image_paths[idx]) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if self.transform: augmented = self.transform( image=image, bboxes=self.labels[idx]['bboxes'], class_labels=self.labels[idx]['labels'] ) image = augmented['image'] bboxes = augmented['bboxes'] labels = augmented['class_labels'] return image, bboxes, labels

3.3 数据标注格式

我们采用YOLO格式进行标注，每个标注文件包含：

• 类别ID

• 边界框中心点坐标（归一化）

• 边界框宽度和高度（归一化）

text

# 标注文件示例 (traffic_sign_001.txt) 0 0.5 0.5 0.2 0.3 # 类别0，中心点(0.5,0.5)，宽0.2，高0.3 1 0.3 0.7 0.15 0.25 # 类别1，中心点(0.3,0.7)，宽0.15，高0.25

4. YOLO模型详解与实现

4.1 YOLOv5模型架构

python

import torch import torch.nn as nn from models.common import Conv, Bottleneck, C3, SPPF class YOLOv5(nn.Module): def __init__(self, num_classes=43): super(YOLOv5, self).__init__() # 主干网络 self.backbone = nn.Sequential( Conv(3, 64, 6, 2, 2), # 0 Conv(64, 128, 3, 2), # 1 C3(128, 128, 3), # 2 Conv(128, 256, 3, 2), # 3 C3(256, 256, 6), # 4 Conv(256, 512, 3, 2), # 5 C3(512, 512, 9), # 6 Conv(512, 1024, 3, 2), # 7 C3(1024, 1024, 3), # 8 SPPF(1024, 1024, 5) # 9 ) # 颈部网络 self.neck = Neck() # 检测头 self.head = Detect(num_classes) def forward(self, x): x = self.backbone(x) x = self.neck(x) return self.head(x)

4.2 YOLOv8改进特性

YOLOv8在以下方面进行了改进：

1. 新的骨干网络：CSPDarknet53 -> CSPDarknet-AA

2. 无锚框检测：简化了训练流程

3. 新的损失函数：Distribution Focal Loss + CIoU Loss

4. 任务对齐分配：改进的标签分配策略

python

# YOLOv8模型定义 class YOLOv8(nn.Module): def __init__(self, cfg='yolov8s.yaml', ch=3, nc=43): super().__init__() self.model, self.save = parse_model(deepcopy(cfg), ch=[ch]) def forward(self, x, augment=False, profile=False): return self._forward_once(x, profile)

4.3 损失函数设计

python

class YOLOLoss: def __init__(self, num_classes=43): self.num_classes = num_classes self.box_loss = CIoULoss() self.cls_loss = nn.BCEWithLogitsLoss() self.dfl_loss = DistributionFocalLoss() def __call__(self, preds, targets): # 解构预测 pred_distri, pred_scores = preds[:, :self.reg_max * 4], preds[:, self.reg_max * 4:] # 边界框损失 loss_box = self.box_loss(pred_boxes, target_boxes) # 分类损失 loss_cls = self.cls_loss(pred_scores, target_cls) # DFL损失（YOLOv8特有） loss_dfl = self.dfl_loss(pred_distri, target_boxes) return loss_box + loss_cls + loss_dfl

5. 模型训练与优化

5.1 训练配置

yaml

# 训练配置文件 (train_config.yaml) train: &train # 数据配置 data: traffic_sign.yaml epochs: 300 batch_size: 16 imgsz: 640 # 优化器配置 optimizer: AdamW lr0: 0.001 lrf: 0.01 momentum: 0.937 weight_decay: 0.0005 # 数据增强 hsv_h: 0.015 hsv_s: 0.7 hsv_v: 0.4 degrees: 0.0 translate: 0.1 scale: 0.5 shear: 0.0 perspective: 0.0 flipud: 0.0 fliplr: 0.5 mosaic: 1.0 mixup: 0.0 # 模型保存 save_period: 50 pretrained: true freeze: 0

5.2 训练脚本

python

import yaml from ultralytics import YOLO import torch def train_yolov8(): # 加载模型 model = YOLO('yolov8s.pt') # 预训练模型 # 训练配置 config = yaml.load(open('train_config.yaml'), Loader=yaml.FullLoader) # 开始训练 results = model.train( data='traffic_sign.yaml', epochs=300, imgsz=640, batch=16, patience=50, save=True, save_period=50, pretrained=True, optimizer='AdamW', lr0=0.001, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3, warmup_momentum=0.8, box=7.5, cls=0.5, dfl=1.5, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0, device='0', # GPU设备 workers=8, project='traffic_sign_detection', name='yolov8_train', exist_ok=True, verbose=True ) return results if __name__ == '__main__': train_yolov8()

5.3 超参数优化

python

import optuna def objective(trial): # 超参数搜索空间 lr = trial.suggest_float('lr', 1e-5, 1e-2, log=True) batch_size = trial.suggest_categorical('batch_size', [8, 16, 32]) weight_decay = trial.suggest_float('weight_decay', 1e-5, 1e-3, log=True) # 训练模型 model = YOLO('yolov8s.pt') results = model.train( data='traffic_sign.yaml', epochs=100, batch=batch_size, lr0=lr, weight_decay=weight_decay, verbose=False ) # 返回验证集mAP return results.box.map # 运行超参数优化 study = optuna.create_study(direction='maximize') study.optimize(objective, n_trials=50)

6. 用户界面设计

6.1 PyQt5界面实现

python

import sys from PyQt5.QtWidgets import * from PyQt5.QtCore import * from PyQt5.QtGui import * import cv2 import torch from ultralytics import YOLO class TrafficSignApp(QMainWindow): def __init__(self): super().__init__() self.model = None self.initUI() self.load_model() def initUI(self): # 主窗口设置 self.setWindowTitle('交通标志识别系统') self.setGeometry(100, 100, 1200, 800) # 中央部件 central_widget = QWidget() self.setCentralWidget(central_widget) # 主布局 main_layout = QHBoxLayout() # 左侧图像显示区域 self.image_label = QLabel() self.image_label.setAlignment(Qt.AlignCenter) self.image_label.setStyleSheet("border: 2px solid #ccc;") main_layout.addWidget(self.image_label, 3) # 右侧控制面板 control_panel = QVBoxLayout() # 文件选择按钮 file_btn = QPushButton('选择图像') file_btn.clicked.connect(self.open_image) control_panel.addWidget(file_btn) # 视频选择按钮 video_btn = QPushButton('选择视频') video_btn.clicked.connect(self.open_video) control_panel.addWidget(video_btn) # 摄像头按钮 camera_btn = QPushButton('开启摄像头') camera_btn.clicked.connect(self.open_camera) control_panel.addWidget(camera_btn) # 模型选择 model_label = QLabel('选择模型:') control_panel.addWidget(model_label) self.model_combo = QComboBox() self.model_combo.addItems(['YOLOv5s', 'YOLOv6s', 'YOLOv7', 'YOLOv8s']) control_panel.addWidget(self.model_combo) # 置信度阈值 conf_label = QLabel('置信度阈值:') control_panel.addWidget(conf_label) self.conf_slider = QSlider(Qt.Horizontal) self.conf_slider.setRange(10, 90) self.conf_slider.setValue(25) control_panel.addWidget(self.conf_slider) # 显示统计信息 self.stats_label = QLabel('统计信息') self.stats_label.setStyleSheet("border: 1px solid #ccc; padding: 10px;") control_panel.addWidget(self.stats_label) # 退出按钮 exit_btn = QPushButton('退出') exit_btn.clicked.connect(self.close) control_panel.addWidget(exit_btn) # 添加到主布局 control_widget = QWidget() control_widget.setLayout(control_panel) main_layout.addWidget(control_widget, 1) central_widget.setLayout(main_layout) def load_model(self): """加载预训练模型""" model_name = self.model_combo.currentText() model_path = f'models/{model_name.lower()}.pt' try: self.model = YOLO(model_path) QMessageBox.information(self, '成功', f'{model_name}模型加载成功！') except Exception as e: QMessageBox.critical(self, '错误', f'模型加载失败: {str(e)}') def open_image(self): """打开图像文件""" file_path, _ = QFileDialog.getOpenFileName( self, '选择图像', '', '图像文件 (*.jpg *.jpeg *.png *.bmp)' ) if file_path: self.process_image(file_path) def process_image(self, image_path): """处理单张图像""" # 读取图像 image = cv2.imread(image_path) image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 模型推理 conf_threshold = self.conf_slider.value() / 100 results = self.model(image_rgb, conf=conf_threshold) # 绘制检测结果 annotated_image = results[0].plot() # 显示图像 height, width, channel = annotated_image.shape bytes_per_line = 3 * width qt_image = QImage(annotated_image.data, width, height, bytes_per_line, QImage.Format_RGB888) self.image_label.setPixmap(QPixmap.fromImage(qt_image)) # 更新统计信息 self.update_stats(results[0]) def update_stats(self, result): """更新统计信息""" num_detections = len(result.boxes) classes = result.boxes.cls.tolist() if result.boxes is not None else [] stats_text = f""" 检测统计: ---------- 检测数量: {num_detections} 检测类别: {classes} FPS: {result.speed.get('inference', 0):.2f} """ self.stats_label.setText(stats_text) if __name__ == '__main__': app = QApplication(sys.argv) window = TrafficSignApp() window.show() sys.exit(app.exec_())

6.2 Gradio Web界面

python

import gradio as gr import cv2 import numpy as np from ultralytics import YOLO import pandas as pd class TrafficSignWebApp: def __init__(self): self.model = YOLO('models/yolov8s.pt') self.class_names = self.load_class_names() def load_class_names(self): """加载类别名称""" # 这里应该是你的类别文件路径 with open('data/traffic_sign.names', 'r') as f: classes = [line.strip() for line in f.readlines()] return classes def predict_image(self, image, conf_threshold): """图像预测""" # 转换图像 image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 推理 results = self.model(image_rgb, conf=conf_threshold) # 绘制结果 annotated_image = results[0].plot() # 创建检测结果表格 detections = [] if results[0].boxes is not None: for box in results[0].boxes: class_id = int(box.cls[0]) confidence = float(box.conf[0]) x1, y1, x2, y2 = box.xyxy[0].tolist() detections.append({ '类别': self.class_names[class_id], '置信度': f'{confidence:.3f}', '位置': f'({x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f})' }) df = pd.DataFrame(detections) if detections else pd.DataFrame() return annotated_image, df def predict_video(self, video_path, conf_threshold): """视频预测""" cap = cv2.VideoCapture(video_path) # 获取视频信息 fps = int(cap.get(cv2.CAP_PROP_FPS)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 创建输出视频 fourcc = cv2.VideoWriter_fourcc(*'mp4v') output_path = 'output_video.mp4' out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) frame_count = 0 total_detections = [] while cap.isOpened(): ret, frame = cap.read() if not ret: break # 推理 results = self.model(frame, conf=conf_threshold) # 绘制结果 annotated_frame = results[0].plot() out.write(annotated_frame) # 统计检测结果 if results[0].boxes is not None: for box in results[0].boxes: class_id = int(box.cls[0]) total_detections.append(self.class_names[class_id]) frame_count += 1 cap.release() out.release() # 生成统计信息 if total_detections: detection_counts = pd.Series(total_detections).value_counts() stats_df = pd.DataFrame({ '交通标志': detection_counts.index, '出现次数': detection_counts.values }) else: stats_df = pd.DataFrame() return output_path, stats_df def create_interface(self): """创建Gradio界面""" with gr.Blocks(title="交通标志识别系统", theme=gr.themes.Soft()) as app: gr.Markdown("# 🚦 交通标志识别系统") gr.Markdown("基于YOLOv8的深度学习交通标志检测系统") with gr.Row(): with gr.Column(): input_image = gr.Image(label="输入图像", type="numpy") conf_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.25, label="置信度阈值") image_btn = gr.Button("检测图像") with gr.Column(): output_image = gr.Image(label="检测结果") output_table = gr.Dataframe(label="检测结果详情", headers=["类别", "置信度", "位置"]) with gr.Row(): with gr.Column(): input_video = gr.Video(label="输入视频") video_btn = gr.Button("检测视频") with gr.Column(): output_video = gr.Video(label="检测结果视频") stats_table = gr.Dataframe(label="视频统计", headers=["交通标志", "出现次数"]) # 绑定事件 image_btn.click( fn=self.predict_image, inputs=[input_image, conf_slider], outputs=[output_image, output_table] ) video_btn.click( fn=self.predict_video, inputs=[input_video, conf_slider], outputs=[output_video, stats_table] ) # 添加示例 gr.Examples( examples=[ ["examples/sign_001.jpg", 0.3], ["examples/sign_002.jpg", 0.25], ["examples/traffic_video.mp4", 0.3] ], inputs=[input_image, conf_slider], outputs=[output_image, output_table], fn=self.predict_image, cache_examples=True ) return app if __name__ == '__main__': app = TrafficSignWebApp() interface = app.create_interface() interface.launch(server_name="0.0.0.0", server_port=7860, share=True)

7. 模型评估与对比

7.1 评估指标

python

import numpy as np from sklearn.metrics import precision_recall_curve, average_precision_score import matplotlib.pyplot as plt class ModelEvaluator: def __init__(self, model, test_loader, num_classes=43): self.model = model self.test_loader = test_loader self.num_classes = num_classes def evaluate(self): """全面评估模型""" metrics = { 'mAP': self.calculate_map(), 'precision': self.calculate_precision(), 'recall': self.calculate_recall(), 'f1_score': self.calculate_f1(), 'inference_time': self.measure_inference_time() } return metrics def calculate_map(self, iou_threshold=0.5): """计算mAP""" all_predictions = [] all_targets = [] self.model.eval() with torch.no_grad(): for images, targets in self.test_loader: outputs = self.model(images) all_predictions.extend(outputs) all_targets.extend(targets) # 计算每个类别的AP ap_scores = [] for class_id in range(self.num_classes): ap = self.calculate_ap_for_class( all_predictions, all_targets, class_id, iou_threshold ) ap_scores.append(ap) return np.mean(ap_scores) def plot_pr_curve(self): """绘制PR曲线""" fig, axes = plt.subplots(1, 2, figsize=(15, 6)) # 总体PR曲线 self.plot_overall_pr_curve(axes[0]) # 各类别PR曲线 self.plot_class_pr_curves(axes[1]) plt.tight_layout() return fig

7.2 YOLO版本对比实验

python

import pandas as pd import matplotlib.pyplot as plt class YOLOComparison: def __init__(self): self.results = {} def run_comparison(self): """运行对比实验""" yolo_versions = ['yolov5s', 'yolov6s', 'yolov7', 'yolov8s'] for version in yolo_versions: print(f"正在测试 {version}...") # 加载模型 model = YOLO(f'models/{version}.pt') # 评估模型 metrics = self.evaluate_model(model) self.results[version] = metrics print(f"{version} 结果: {metrics}") self.save_results() self.plot_comparison() def plot_comparison(self): """绘制对比图""" df = pd.DataFrame(self.results).T fig, axes = plt.subplots(2, 2, figsize=(12, 10)) # mAP对比 df['mAP'].plot(kind='bar', ax=axes[0, 0], title='mAP对比') axes[0, 0].set_ylabel('mAP@0.5') # 推理速度对比 df['inference_time'].plot(kind='bar', ax=axes[0, 1], title='推理速度对比') axes[0, 1].set_ylabel('FPS') # 模型大小对比 df['model_size'].plot(kind='bar', ax=axes[1, 0], title='模型大小对比') axes[1, 0].set_ylabel('MB') # F1分数对比 df['f1_score'].plot(kind='bar', ax=axes[1, 1], title='F1分数对比') axes[1, 1].set_ylabel('F1 Score') plt.tight_layout() plt.savefig('yolo_comparison.png', dpi=300, bbox_inches='tight') plt.show()

8. 部署与优化

8.1 模型量化

python

import torch.quantization class ModelQuantizer: def __init__(self, model): self.model = model def quantize_model(self): """量化模型""" # 设置为评估模式 self.model.eval() # 量化配置 quantization_config = torch.quantization.get_default_qconfig('fbgemm') self.model.qconfig = quantization_config # 准备量化 torch.quantization.prepare(self.model, inplace=True) # 校准（使用校准数据集） self.calibrate_model() # 转换量化模型 quantized_model = torch.quantization.convert(self.model, inplace=True) return quantized_model def calibrate_model(self, calibration_loader): """校准模型""" with torch.no_grad(): for images, _ in calibration_loader: _ = self.model(images)

8.2 TensorRT加速

python

import tensorrt as trt class TensorRTConverter: def __init__(self, onnx_model_path): self.onnx_path = onnx_model_path self.logger = trt.Logger(trt.Logger.WARNING) def convert_to_tensorrt(self, engine_path, max_batch_size=32): """转换为TensorRT引擎""" builder = trt.Builder(self.logger) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser = trt.OnnxParser(network, self.logger) # 解析ONNX模型 with open(self.onnx_path, 'rb') as f: if not parser.parse(f.read()): for error in range(parser.num_errors): print(parser.get_error(error)) return None # 构建配置 config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1GB config.set_flag(trt.BuilderFlag.FP16) # 使用FP16 # 构建引擎 engine = builder.build_engine(network, config) # 保存引擎 with open(engine_path, 'wb') as f: f.write(engine.serialize()) return engine