ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

基于YOLO+ByteTrack的电力机车视频检测与跟踪方案解析

基于YOLO+ByteTrack的电力机车视频检测与跟踪方案解析 这次我们来看一个轨道交通视频分析场景EP2K“至冬铁路”电力机车牵引列车通过某区域。表面上看这是行车记录或摄影素材的一句描述但如果把它放到工程视角里它就是一个非常典型的“视频目标识别 目标跟踪 批量处理”落地场景。我们需要回答的问题包括这段视频里电力机车是什么时候出现的、从哪个方向通过、画面里有没有其他干扰目标、能不能把每次通过都自动记下来。这篇文章不会去讨论铁路运营本身的业务而是围绕“如何搭建一套能处理这类行车视频的本地检测与跟踪方案”来展开。核心关注点放在四件事上视频抽帧与素材准备、目标检测与机车识别、连续帧跟踪与通过判定、批量任务与接口服务。这样你拿到一段类似场景的视频就能按同样的流程跑通一条可复用的分析链路。先给结论这不是一个开箱即用的单一软件而是一套需要自己组装的技术栈。常见做法是用 YOLO 系列模型做机车检测用 ByteTrack 或 DeepSORT 做跟踪用 FFmpeg 做视频抽帧再用 FastAPI 把推理封装成接口。视频分辨率、模型输入尺寸、显存大小都会直接影响效果第一次测试时建议用小分辨率、小模型、单视频跑通再逐步放大。1. 核心能力速览在开始部署前先把这条技术链路的整体能力列出来方便你判断它适不适合自己的环境。能力项说明项目类型视频目标检测与跟踪方案适配行车画面分析输入素材电力机车、列车通过区域的行车视频、监控视频或照片主要功能机车目标检测、目标跟踪、通过区域判定、批量视频处理、推理结果可视化核心模型YOLO 系列等通用目标检测模型YOLO 权重需按自己的素材类别训练或微调显存需求取决于模型版本、输入分辨率、推理批大小需以实际环境测试为准是否支持 CPU支持但 CPU 推理速度会明显低于 GPU适合少量图片验证支持平台Windows、Linux 均可建议 Linux 服务器做长期批量任务启动方式命令行启动、Python 脚本启动、FastAPI 接口服务启动是否支持 API支持可用 FastAPI/Flask 封装检测与跟踪服务是否支持批量任务支持可对视频目录做循环批量处理适合场景铁路摄影素材归档、机车通过频次统计、行车画面自动标注、模型效果验证表格里的“显存需求”和“是否支持 CPU”都是通用能力说明不要看作固定数值。实际部署时模型权重文件、视频分辨率、检测间隔都会改变资源占用。2. 适用场景与使用边界这类方案最适合的落地场景有三类。第一类是铁路摄影爱好者的素材管理比如你拍了一批电力机车通过某区域的长视频希望自动把“有车通过”的片段截取出来省去手动拉进度条。第二类是特定区域的通过频次统计比如统计一天内有多少次列车经过可以做简单的场景计数。第三类是行车视频的自动标注把检测到的机车位置、时间、帧号写进 CSV 或 JSON方便后续检索。但它不适合做实时安全控制。如果要用在铁路运行安全相关场景必须由专业系统完成不能把这段 AI 识别结果作为调度、防护或告警的唯一依据。原因很简单目标检测模型的漏检率和误检率都不可能做到零单模型判定的可靠性无法满足安全等级要求。因此本文所有内容仅限技术验证、素材整理、数据统计等辅助用途。使用边界也要注意。视频素材如果来自公开网络、朋友拍摄或自行拍摄需要确认原始素材的使用授权。涉及铁路沿线监控、未公开车站、非公开运行线路的画面不要随意采集、传播或用于商用。如果你处理的是包含人物、车牌、人脸的视频还要格外注意隐私合规。建议在项目目录里单独建一份 README记录素材来源、授权情况和使用范围生产环境里这比模型精度更重要。3. 环境准备与前置条件整体技术方案基于 Python因此环境准备以 Python 生态为主。建议操作系统优先选 Linux比如 Ubuntu 20.04 或 22.04。如果只有 Windows也可以跑通但批量视频处理和长时间运行时的稳定性不如 Linux。硬件方面如果你有 NVIDIA 显卡优先用 GPU 推理没有独立显卡用 CPU 也能完成小规模验证只是速度会慢很多。Python 版本建议 3.9 到 3.11。PyTorch 各版本对 Python 版本要求不同使用 3.9 或 3.10 兼容性比较稳。CUDA 和 cuDNN 按你安装的 PyTorch 版本来配。不要直接装系统最新版 CUDA先查 PyTorch 官方安装命令需要的 CUDA 版本。磁盘空间方面视频素材、模型权重、输出结果会占不少空间。一段 1080p 视频抽帧后每秒按 2 到 5 帧存储一小时素材就可能产生几千张图片建议单独挂一个大目录。FFmpeg 是必装的Windows 用户可以用包管理工具安装也可以直接下载静态编译版把可执行文件放到 PATH 中。准备阶段先执行三条命令确认环境状态。# 查看显卡驱动和 CUDA 版本 nvidia-smi # 查看 Python 版本 python --version # 查看 FFmpeg 是否可用 ffmpeg -version如果 nvidia-smi 显示不了而你有独立显卡先排查驱动。FFmpeg 如果提示找不到命令需要先安装或配置 PATH。4. 安装部署与启动方式下面给一套通用的安装流程。基本思路是先建虚拟环境再装深度学习框架最后装目标检测和视频处理相关依赖。创建项目目录和虚拟环境mkdir -p ep2k-project cd ep2k-project python -m venv venv # Linux / macOS source venv/bin/activate # Windows venv\Scripts\activate安装依赖。先把 PyTorch 按官方命令安装这里不写死 CUDA 版本你需要根据自己机器的驱动和官方索引选择对应命令。安装完成后再安装其他依赖。pip install ultralytics opencv-python numpy tqdm pip install fastapi uvicorn python-multipart pip install onnxruntime其中 ultralytics 会自动拉取 YOLO 相关依赖opencv-python 负责图像读写FastAPI 和 Uvicorn 用来封装推理接口。onnxruntime 不是必须的如果你想用 ONNX 格式推理可以提前装上。模型权重文件需要从官方渠道下载但这里不指定具体版本。检测类别如果是“电力机车”而通用模型本身只包含常见的 80 类目标不保证能识别出“EP2K 电力机车”。所以你要么使用通用模型的检测结果作为候选框再按区域过滤要么准备少量标注数据微调一个专用的机车检测模型。后者的精度会明显更高。一个可选的做法是先跑通官方预训练权重看看视频里能不能输出有效检测框再决定是否做微调。第一次跑不要追求精确类别先验证链路通不通。准备一个启动脚本把视频处理流程变成可重复执行的任务# processor.py import cv2 from pathlib import Path def extract_frames(video_path, output_dir, fps2): output_dir Path(output_dir) output_dir.mkdir(parentsTrue, exist_okTrue) cap cv2.VideoCapture(str(video_path)) video_fps cap.get(cv2.CAP_PROP_FPS) interval max(1, int(video_fps / fps)) frame_idx 0 saved_idx 0 while True: ret, frame cap.read() if not ret: break if frame_idx % interval 0: out_path output_dir / fframe_{saved_idx:06d}.jpg cv2.imwrite(str(out_path), frame) saved_idx 1 frame_idx 1 cap.release() print(fsaved {saved_idx} frames to {output_dir}) if __name__ __main__: extract_frames(test.mp4, output_frames, fps2)视频处理任务不像模型推理那样一次就能跑完建议先用短视频测试确认抽帧逻辑和输出目录正确后再跑完整素材。5. 功能测试与效果验证这一段是整个方案的验证核心。拿“EP2K 电力机车牵引列车通过某区域”这类视频来说建议按以下顺序逐项测试。5.1 视频抽帧测试测试目的确认视频能正常解码抽帧结果不花屏、不丢帧。操作步骤python processor.py输入视频为 test.mp4输出到 output_frames 目录。判断标准是输出帧数是否符合预期。比如视频时长 30 秒抽帧频率 2 fps预期约 60 张。如果输出远少于预期检查视频编码格式有些监控视频需要用 FFmpeg 先转码。ffmpeg -i test.mp4 -c:v libx264 -pix_fmt yuv420p test_h264.mp45.2 单张图片目标检测测试测试目的验证模型能否在画面中检测到列车或机车目标。用 YOLO 做单张图片检测from ultralytics import YOLO model YOLO(yolo.pt) # 替换为实际权重路径 results model.predict(sourceoutput_frames/frame_000000.jpg, conf0.25) for r in results: boxes r.boxes if boxes is not None: print(detected:, len(boxes)) r.save(filenameresult.jpg)如果没有检测框先调低 conf 阈值到 0.1确认是漏检还是阈值问题。如果画面里有很多干扰物比如电线杆、树木、建筑就要考虑只保留画面下方的轨道区域减少误检。5.3 连续视频检测与跟踪测试单张图片检测只能给出位置不能区分“这是同一辆车还是两辆车”。所以要加跟踪器。ByteTrack 是当前比较常用的多目标跟踪方案结合检测结果可以实现稳定跟踪。使用 ByteTrack 的通用流程如下import cv2 import numpy as np from ultralytics import YOLO from bytetracker import BYTETracker model YOLO(yolo.pt) tracker BYTETracker() cap cv2.VideoCapture(test.mp4) track_id_set set() while True: ret, frame cap.read() if not ret: break results model.predict(frame, conf0.3, verboseFalse) dets [] for r in results: if r.boxes is None: continue for box in r.boxes: x1, y1, x2, y2 box.xyxy[0].tolist() score box.conf.item() dets.append([x1, y1, x2, y2, score]) if len(dets) 0: tracks tracker.update(np.array(dets), frame.shape[:2], (frame.shape[1], frame.shape[0])) for t in tracks: track_id t.track_id track_id_set.add(track_id) cap.release() print(unique trains detected:, len(track_id_set))这里简单统计了去重后的 track_id 数量。如果同一辆机车从画面左侧进入、右侧离开跟踪器会保持同一个 ID因此最终数量更接近“通过次数”。判断成功的标准是同一辆列车跨越多帧时 ID 不跳变两辆不同列车同时出现时 ID 不互相混淆。如果 ID 频繁切换说明检测框不稳定可以调低检测阈值或提高跟踪器参数。5.4 批量视频处理测试单段视频跑通后可以处理整个目录的视频。video_dir./videos for video in $video_dir/*.mp4; do echo Processing $video python process_video.py $video done建议在 process_video.py 里记录每个视频的处理状态。例如输出 JSON 结果文件列出视频名、检测到列车的帧数、unique track id 数量、输出视频路径。这样即使中间某个视频出错也能快速定位到具体文件。{ video: EP2K_eastbound_001.mp4, frames: 1200, detected_frames: 320, unique_tracks: 2, output_video: output/EP2K_eastbound_001_annotated.mp4 }5.5 效果验证方法不要只看一两次结果就下结论。建议准备 3 段不同场景的视频比如白天、傍晚、逆光分别统计漏检和误检情况。漏检指有列车通过但模型没有检测框误检指没有列车但模型标出了目标。如果漏检率比较高优先补充训练数据或降低阈值。如果误检率比较高建议在后处理中增加轨道区域过滤。6. 接口 API 与批量任务如果只是想本地跑一两个视频不需要接口。但如果你想把检测能力接到自己的工具链里比如做一个内部素材管理系统就需要用 API 封装。6.1 推理接口启动用 FastAPI 起一个轻量服务接收视频文件或图片返回检测结果。# api.py import tempfile from fastapi import FastAPI, File, UploadFile from ultralytics import YOLO import cv2 app FastAPI() model YOLO(yolo.pt) app.post(/detect) async def detect(file: UploadFile File(...)): content await file.read() with tempfile.NamedTemporaryFile(suffix.jpg) as tmp: tmp.write(content) tmp.flush() result model.predict(tmp.name, conf0.3, verboseFalse) boxes [] for r in result: if r.boxes is not None: for box in r.boxes: boxes.append({ xyxy: box.xyxy[0].tolist(), conf: box.conf.item(), cls: int(box.cls.item()) }) return {count: len(boxes), boxes: boxes} if __name__ __main__: import uvicorn uvicorn.run(app, host127.0.0.1, port8000)启动服务python api.py服务启动后访问 http://127.0.0.1:8000/docs 可以打开 Swagger 文档直接测试 /detect 接口。这样比手动敲 curl 更方便。6.2 接口调用测试上传一张抽帧图片验证接口是否正常。curl -X POST http://127.0.0.1:8000/detect \ -F fileoutput_frames/frame_000000.jpg如果返回 JSON 里的 count 大于 0说明接口链路已经跑通。注意默认服务只监听 127.0.0.1外部机器访问不了。如果要在局域网内测试需要改成 0.0.0.0但要注意访问控制不要把接口裸奔到公网。6.3 批量任务目录设计批量处理不一定要用队列系统对中小规模任务目录设计就够用。建议采用如下目录结构ep2k-project/ ├── videos/ # 输入视频 ├── frames/ # 抽帧结果 ├── outputs/ # 检测结果、标注视频、JSON 日志 ├── models/ # 模型权重 └── logs/ # 任务日志处理逻辑可以按“扫描 input 目录 - 处理 - 写结果 - 移动视频到 done 目录”的顺序来。这样就算任务中断也能通过文件名判断哪些视频处理完。mkdir -p videos done outputs logs6.4 失败重试建议批量任务容易出现两类问题视频编码不兼容、模型推理进程崩溃。建议在代码里加上异常捕获和重试逻辑。每处理一个视频都写日志不把异常信息只输出到控制台。import logging logging.basicConfig(filenamelogs/process.log, levellogging.INFO) for video_path in video_list: try: process_video(video_path) logging.info(fOK {video_path}) except Exception as e: logging.error(fFAIL {video_path}: {e})如果某个视频反复失败可以把该文件单独抽出来转码后再处理。不要因为一个坏视频阻塞整个批量任务。7. 资源占用与性能观察运行视频检测任务时资源占用主要来自三个部分视频解码、模型推理、跟踪器计算。视频解码的 CPU 占用通常比较高。1080p 视频在 CPU 上解码时即使模型用 GPU也能看到 CPU 占用明显上升。建议先用 FFmpeg 把视频压缩或抽帧减少推理阶段反复解码的压力。模型推理的显存占用取决于模型输入尺寸和 batch size。越大的输入尺寸精度可能越高但显存占用也会上升。第一次测试建议用 640x640 输入batch size 设为 1。观察显存占用的方法是nvidia-smi -l 2每两秒刷新一次。重点关注进程的显存占用和 GPU 利用率。如果显存接近上限先尝试降低 batch size再把输入分辨率从 640 降到 416 或 320。不要一上来就追求高精度参数先保证任务能稳定跑完。如果使用 CPU 推理速度取决于 CPU 核心数和模型体积。小模型在 CPU 上处理单张图片可能还能等待但处理整段视频会非常慢。实际使用中CPU 更适合做单张图片测试和接口联调不适合长时间批量视频任务。批量处理时要注意输出磁盘占用。标注视频如果每一帧都写盘会拖慢整体流程。更高效的做法是检测跟踪完成后只把包含目标的片段裁剪保存其他帧不写盘。8. 常见问题与排查方法问题现象可能原因排查方式解决方案Python 安装依赖失败网络源不稳定或 Python 版本过旧检查 pip 版本和源切换国内镜像源或升级 Python 到 3.10模型加载报错权重文件路径错误或文件损坏检查 models 目录重新下载权重文件显卡识别不到驱动未安装或 CUDA 版本不对执行 nvidia-smi安装匹配的显卡驱动按 PyTorch 要求装 CUDAGPU 推理比 CPU 还慢模型过小或输入分辨率过低处理开销小对比 CPU/GPU 单张耗时扩大 batch size 或使用更大模型视频抽帧结果全部黑屏视频编码格式不支持用 FFmpeg 查看编码信息转码为 H.264检测不到列车阈值太高或模型不认识机车降低 conf 到 0.1观察检测框收集数据微调专用模型跟踪 ID 频繁跳变检测框不稳定查看单帧检测框位置降低检测阈值使用 ByteTrack 参数调优API 启动端口被占用8000 端口被其他程序占用检查端口监听情况换端口启动或停止占用进程批量任务卡住视频文件损坏或循环等待查看日志最后一行加异常捕获跳过异常视频输出结果文件太大每帧都写标注视频检查输出目录只保存包含目标的片段9. 最佳实践与使用建议第一次跑这套流程不要直接处理几十个视频。先选一段 10 秒短视频把抽帧、检测、跟踪、接口四个环节全部跑通再扩大范围。这样排错成本最低也最容易确认问题出在哪个环节。项目目录要分清楚。输入素材、临时抽帧、检测输出、模型权重、日志文件不要混在一起。一个最简单的做法就是前面提到的五目录结构videos、frames、outputs、models、logs。清理结果时只需清空 outputs 和 logs不会误删原始素材。批量任务一定要加日志和失败重试。不要相信“这次应该没问题”视频编码、内存占用、显存波动都可能让任务中断。日志里至少要记录每个视频的开始时间、结束时间、是否成功、检测目标数量。接口服务要限制访问范围。如果只是本机使用监听 127.0.0.1 就够了。如果部署到服务器建议放在内网不要直接暴露公网。API 调用频繁时可以在服务里加一个简单的并发限制避免多用户同时上传大视频把显存打满。涉及行车视频、铁路场景时合规问题优先级最高。不要用未授权的视频素材不要处理涉密或敏感区域画面。如果视频里出现工作人员、乘客、车牌等信息要考虑模糊处理或放弃使用。发布结果前再确认一次素材授权和展示边界。模型精度方面如果要用到生产环境最值得投入的是数据标注。通用模型能给出一些检测框但针对“电力机车牵引列车”这个目标最好手工标注几百张不同角度、不同光照的图片微调一次专用模型。这样漏检和误检会明显下降。10. 总结与下一步这套方案最值得尝试的点是它把视频处理的几个高频环节串在了一起抽帧、检测、跟踪、批量、API。对一个铁路视频分析场景来说早期不需要很强的模型能力先把“能跑通”这件事做好再考虑“精度更高”。最先应该验证的功能是单张图片检测。如果模型在抽帧图片上找不到目标后面所有跟踪和批量逻辑都没有意义。先跑通这一步再接入跟踪器最后再封装接口。最容易踩的坑是环境问题。CUDA 版本和 PyTorch 不匹配权重文件下载不完整FFmpeg 路径没有配置好都会让启动变得非常曲折。建议把所有环境检查命令写成一个 shell 脚本运行一次就确认基础环境是否正常。下一步可以继续扩展的方向包括训练一个专用机车检测模型加入 OCR 识别车号把“通过区域”画成 ROI 多边形只在区域内部判定目标把结果写入数据库按时间段统计通过频次。从“能检测”到“能统计”再到“能辅助业务”每一步都有明确的技术收益。
返回列表