YOLOv5 铁路信号机视频自动识别与仿真系统
基于 YOLOv5 目标检测 + DeepSort 目标跟踪 + PyQt5 可视化界面的铁路信号机智能识别系统
目录
项目简介
本项目是一款面向铁路行车安全场景的智能辅助系统，核心目标是对列车运行过程中前方信号机进行实时检测、跟踪与筛选，并通过 PyQt5 可视化界面将结果直观呈现给司机，以缓解司机视觉疲劳、降低误判风险、提升行车安全性。
系统主要具备以下三大功能：
- 图片检测：对单张铁路场景图片进行信号机目标检测
- 视频检测：对铁路行车录像进行离线信号机检测与跟踪
- 实时检测：接入摄像头进行实时信号机检测与跟踪
支持检测的信号机类别如下：
| 信号机种类 | 检测类别名称 | 说明 |
|---|---|---|
| 红色 | red |
停车，禁止越过信号机 |
| 绿色 | green |
允许越过该信号机 |
| 单黄色 | s_yellow |
减速，或指示正线停车 |
| 双黄色 | d_yellow |
指示侧线停车 |
| 月白色 | white |
引导信号，或允许越过该信号机调车 |
| 蓝色 | blue |
禁止越过该信号机调车 |
算法原理
1. YOLOv5 目标检测
YOLOv5（You Only Look Once v5）是一种单阶段目标检测算法，其核心思想是将目标检测问题转化为回归问题，在一次前向传播中同时预测目标的边界框和类别概率。
本系统基于 YOLOv5s（small）模型，并针对铁路信号机检测场景进行了以下改进：
- Mosaic9 数据增强：在原有 Mosaic4（4 张图拼接）基础上扩展为 9 张图拼接，大幅增加训练样本的背景多样性，提升模型对小目标的感知能力。
- 带边缘扩展的 Copy-Paste 数据增强：将目标实例从原图中抠出并粘贴到其他背景图像上，同时对目标边缘进行适当扩展，增强模型对信号机目标的特征学习能力。
- 小目标检测层：在原有 3 个检测头（P3/P4/P5）基础上增加更高分辨率的 P2 检测头，专门用于捕捉小目标信号机，显著提升远距离信号机的检出率。
2. DeepSort 目标跟踪
DeepSort 是一种基于"检测后跟踪"（Tracking-by-Detection）范式的多目标跟踪算法，其工作流程如下：
- 卡尔曼滤波预测：利用卡尔曼滤波器对每个目标的下一帧位置进行预测，缩小匹配搜索范围
- 特征提取：使用深度卷积网络（ReID 模型）提取每个目标的外观特征向量
- 级联匹配：结合马氏距离（运动信息）和余弦距离（外观特征）进行匈牙利算法匹配
- 轨迹管理：对未匹配目标建立新轨迹，对已消失轨迹进行删除
在本项目中，DeepSort 用于对 YOLOv5 检测出的信号机进行实时跟踪，获取同一信号机在连续帧中的轨迹信息，为后续的多信号机筛选提供数据基础。
3. 多信号机筛选（MultiSignalFilter）
在实际铁路场景中，司机视野内可能同时出现多个信号机，但只有指示本列车运行的信号机才是有效目标。MultiSignalFilter 模块基于信号机连续轨迹的运动特征，对多个候选信号机进行筛选，输出与列车运行相关的有效信号机检测结果。
环境要求
| 环境 | 版本要求 |
|---|---|
| 操作系统 | Windows 10/11 |
| Python | 3.8（推荐） |
| PyTorch | ≥ 1.12.0 |
| CUDA | 11.x（推荐，需与 PyTorch 版本匹配） |
| cuDNN | 与 CUDA 版本对应 |
| 显卡 | NVIDIA GPU（显存 ≥ 4GB，推荐 6GB 以上） |
说明：CPU 模式可运行但速度较慢，强烈建议使用 GPU 环境。
安装步骤
第一步：安装 CUDA 和 cuDNN
前往 NVIDIA 官网下载并安装与显卡驱动匹配的 CUDA Toolkit，并配置对应版本的 cuDNN。
第二步：安装 PyTorch
根据 CUDA 版本从 PyTorch 历史版本页面 选择对应命令安装，例如 CUDA 11.6：
bash
pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 --extra-index-url https://download.pytorch.org/whl/cu116第三步：克隆项目并安装依赖
bash
# 克隆项目代码
git clone <your-repo-url>
cd yolov5-5.0
# 安装项目依赖
pip install -r requirements.txt主要依赖说明：
torch / torchvision：深度学习框架opencv-python：图像处理numpy / scipy / matplotlib：数值计算与可视化PyQt5：图形界面框架PyYAML：配置文件解析tensorboard：训练日志可视化tqdm：进度条显示openpyxl / pandas：数据导出处理easydict / Cython：DeepSort 依赖
运行步骤
启动图形界面（推荐）
bash
python mainwindow.py启动后将进入系统主界面，界面包含以下功能区域：
- 参数设置模块：调整检测置信度、IoU 阈值、推理图片尺寸等参数
- 功能选择模块：切换图片检测 / 视频检测 / 实时检测模式
- 结果显示模块：展示检测结果、检测类别与置信度
- 视频播放模块：实时播放检测视频流
- 轨迹显示模块：动态展示 DeepSort 跟踪的实时轨迹
- 用时显示模块 ：显示单帧检测耗时
命令行推理（可选）
bash
# 对单张图片进行检测
python detect.py --source data/images/bus.jpg --weights weights/train_signal500.pt --conf 0.4
# 对视频进行检测
python detect.py --source path/to/video.mp4 --weights weights/train_signal500.pt --conf 0.4
# 使用摄像头实时检测
python detect.py --source 0 --weights weights/train_signal500.pt --conf 0.4
模型训练（可选）
bash
python train.py --data data/train_signal.yaml --cfg models/yolov5s_improved.yaml --weights weights/yolov5s.pt --epochs 100 --batch-size 16 --img-size 640注意事项
- GPU 显存不足 ：若训练时出现
CUDA out of memory错误，请适当减小--batch-size或--img-size参数。 - PyTorch 与 CUDA 版本匹配：请务必确认 PyTorch 版本与 CUDA 版本兼容，否则 GPU 加速将无法生效。
- 权重文件 ：
weights/目录下已提供预训练权重，其中train_signal500.pt为最优模型（mAP@0.5:0.95 = 0.8204），推荐优先使用。 - 摄像头权限：使用实时检测功能时，需确保系统已授予摄像头访问权限。
- 视频格式 ：视频检测支持常见格式（
.mp4、.avi、.mov等），建议使用 H.264 编码的 MP4 文件。 - 路径中不要包含中文：部分第三方库对中文路径支持不佳，建议将项目放置在纯英文路径下。
- 数据集格式 ：自定义训练时，数据集需采用 YOLO 格式标注（
class x_center y_center width height，归一化到 0-1）。 - DeepSort 权重 ：DeepSort 的 ReID 权重文件位于
deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7，请确保该文件存在。
文件夹结构
yolov5-5.0/
├── data/ # 数据配置目录
│ ├── images/ # 示例图片（bus.jpg、zidane.jpg）
│ ├── scripts/ # 数据集下载脚本（COCO、VOC、Argoverse）
│ ├── coco.yaml # COCO 数据集配置
│ ├── voc.yaml # VOC 数据集配置
│ ├── road_signal.yaml # 公路信号机数据集配置
│ ├── train_signal.yaml # 铁路信号机数据集配置
│ ├── train_signal500.yaml # 铁路信号机（500张）数据集配置
│ ├── hyp.scratch.yaml # 从头训练的超参数配置
│ └── hyp.finetune.yaml # 微调训练的超参数配置
│
├── models/ # 模型定义目录
│ ├── hub/ # 各类模型变体配置（YOLOv3/5/6系列）
│ ├── common.py # 通用模块（卷积、残差块等基础组件定义）
│ ├── yolo.py # YOLO 模型主体构建逻辑
│ ├── export.py # 模型导出工具（ONNX、TorchScript等）
│ ├── experimental.py # 实验性功能（模型加载辅助）
│ ├── yolov5s.yaml # YOLOv5s 模型结构配置
│ └── yolov5s_improved.yaml # YOLOv5s 改进版模型结构配置（含小目标检测层）
│
├── utils/ # 工具函数目录
│ ├── aws/ # AWS 云部署脚本
│ ├── google_app_engine/ # Google App Engine 部署配置
│ ├── wandb_logging/ # Weights & Biases 日志集成
│ ├── datasets.py # 数据集加载与预处理
│ ├── general.py # 通用工具函数（NMS、坐标转换等）
│ ├── loss.py # 损失函数定义
│ ├── metrics.py # 评估指标计算（mAP等）
│ ├── plots.py # 可视化工具（绘制检测框、结果图等）
│ ├── torch_utils.py # PyTorch 工具函数（设备选择、模型EMA等）
│ ├── autoanchor.py # 自适应锚框计算
│ └── activations.py # 激活函数定义
│
├── deep_sort_pytorch/ # DeepSort 目标跟踪模块
│ ├── configs/
│ │ └── deep_sort.yaml # DeepSort 跟踪参数配置
│ ├── deep_sort/
│ │ ├── deep/ # ReID 特征提取网络
│ │ │ ├── checkpoint/
│ │ │ │ └── ckpt.t7 # ReID 预训练权重
│ │ │ ├── model.py # ReID 网络结构定义
│ │ │ └── feature_extractor.py# 特征提取器封装
│ │ ├── sort/ # Sort 跟踪核心算法
│ │ │ ├── kalman_filter.py # 卡尔曼滤波器
│ │ │ ├── linear_assignment.py# 匈牙利匹配算法
│ │ │ ├── nn_matching.py # 最近邻匹配
│ │ │ ├── track.py # 轨迹类定义
│ │ │ └── tracker.py # 跟踪器主逻辑
│ │ └── deep_sort.py # DeepSort 整体封装
│ └── utils/ # DeepSort 辅助工具
│
├── Datasets_Processing/ # 数据集处理工具
│ ├── datasets_statistic.py # 数据集统计分析
│ ├── reformat.py # 数据集格式转换
│ ├── rename.py # 文件批量重命名
│ └── SuperfluousFilesSearch.py # 冗余文件查找
│
├── weights/ # 模型权重目录
│ ├── yolov5s.pt # YOLOv5s 预训练权重
│ ├── coco128.pt # COCO128 训练权重
│ ├── road_signal.pt # 公路信号机训练权重
│ ├── train_signal.pt # 铁路信号机训练权重
│ └── train_signal500.pt # 铁路信号机（500张）最优权重
│
├── log/ # 运行日志目录
│ ├── position_log.xlsx # 信号机位置日志
│ ├── time_log.xlsx # 检测时间日志
│ └── position_log.png # 位置日志可视化
│
├── resource/ # UI 资源目录（图标、背景图等）
├── UiMainwindow.py # 主窗口 UI 代码（由 .ui 文件生成）
├── UiMainwindow.ui # 主窗口 Qt Designer 文件
├── mainwindow.py # 主程序入口（启动图形界面）
├── model_load.py # 模型加载封装
├── multi_signal_filter.py # 多信号机筛选模块
├── my_window_effect.py # 窗口特效（毛玻璃效果等）
├── resource.qrc # Qt 资源描述文件
├── resource_rc.py # Qt 资源编译后的 Python 文件
├── signal.yaml # 信号机类别配置
├── detect.py # 推理检测脚本
├── train.py # 模型训练脚本
├── test.py # 模型测试/评估脚本
└── requirements.txt # 项目依赖清单模型权重参考
| 模型文件 | 训练数据集 | 输入尺寸 | mAP@0.5:0.95 |
|---|---|---|---|
coco128.pt |
COCO128 | 640×640 | 0.7996 |
road_signal.pt |
road_signal | 640×640 | 0.7516 |
train_signal.pt |
train_signal | 640×640 | 0.7805 |
train_signal500.pt |
train_signal500 | 640×640 | 0.8204 |