ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyTorch实时车流量统计:YOLO检测、跟踪与TensorRT加速

PyTorch实时车流量统计:YOLO检测、跟踪与TensorRT加速 简介本项目是一套基于深度学习的高速公路车流量实时统计实践资料适合具备一定Python基础、希望上手计算机视觉目标检测的开发者与学习者。资源围绕车辆检测与计数展开涵盖数据处理、模型训练、测试评估与视频流部署等环节可帮助读者从零搭建一套可运行的交通监控统计流程。压缩包共30个文件包括Python脚本、YAML配置、Shell脚本、图片与视频示例等其中py文件覆盖网络定义、训练/测试及演示逻辑yml用于框架参数配置视频文件可用于实时检测效果验证包体约72.54MB整体结构清晰。该资源已有225人学习下载内容包含数据集示例、预训练模型引导、训练与测试脚本、部署参考代码等尤其适合学习YOLO、SSD等检测模型在真实交通场景中的落地应用。1. 高速路实时车流量统计把摄像头画面变成分钟级数字高速路摄像头传回的视频里车辆挨着车灯穿行真正要回答的问题只有一个这一分钟到底过去多少辆车。把这个问题拆成技术动作就是检测每辆车、给每辆车分配唯一ID并跟踪最后在一条虚拟线上做跨线统计。深度学习框架在这里承担的不只是物体识别的计算还决定了整套系统能不能跑到实时从选择PyTorch还是TensorFlow到YOLO这类检测网络怎样被加载、加速、部署每一步都与吞吐量挂钩。这篇文章面向准备用现有服务器或边缘设备落地交通检测的工程师按真实业务链路把模型选型、视频流接入、跟踪计数和性能验证串起来直接能用。2. 深度学习框架选型车辆检测里的PyTorch优势与YOLO对比2.1 动态图与部署生态让PyTorch更贴近实时场景讨论实时车流统计时第一件事不是搭界面而是确定用哪个深度学习框架承载模型推理。我在高速视频场景下见到的工程里PyTorch出现频率远高于TensorFlow。原因有两条其一PyTorch默认采用动态图构建模型在排查检测输出异常时可以直接打印中间张量比如把特征图尺寸打出来对比输入比例不用反复追踪静态图的节点定义其二PyTorch与深度学习框架生态里的目标检测库Ultralytics YOLO、MMDetection原生耦合避免在多个框架之间写转换层。TorchScript和ONNX的导出路径也比较直后面再做TensorRT加速很顺畅。实时统计对推理延迟很敏感PyTorch 2.0以后用torch.compile对卷积模型做图优化配合CUDA环境一般能获得20%左右的加速。如果项目里已经有一部分C代码LibTorch可以通过C API加载同一个TorchScript模型这是很多交通监控系统最终会走的路。下面这张表是我平时给团队做选型时的判断依据。深度学习框架动态排错导出部署实时推理常用后段适合场景PyTorch强TorchScript / ONNXTensorRT检测模型经常迭代开发排错比重高TensorFlow弱SavedModel / ONNXTFLite已有TF服务团队熟悉约定俗成ONNX Runtime无直接运行ONNXCUDA / TensorRT跨框架部署需要统一推理入口2.2 车辆检测模型选型YOLOv8并不是唯一答案模型的选择比框架更影响实时统计上限。高速路车辆类别少但目标尺寸跨度大近处车辆占据大半画面远处百米外的车辆只有十几个像素。YOLOv8n是当前最稳妥的起步方案参数量约3.2M在RTX 3060上处理640分辨率图片能跑到80FPS以上远远满足车流量统计需求。YOLOv8s虽然精度更高但算力开销接近翻倍。如果同时需要识别车辆颜色、车型或者后续嵌入车牌识别需要换用更大的backbone或者多目标头网络这时实时性要从模型切分角度重新考虑。还有一类做法是直接用YOLOv5的--half半精度模式配合固定输入形状把帧率再拉高一截。对于纯粹的车流量统计Fast R-CNN这类两阶段模型不建议使用虽然mAP好看但单帧推理在消费级显卡上也要100ms以上留给跟踪和计数的预算不够。更轻量的NanoDet-M在CPU上也能跑到15FPS适合完全没有GPU的边缘盒子但它的训练资料少出现漏检时能调的空间也小。模型输入尺寸精度特点RTX 3060单帧耗时YOLOv8n640×640小目标稍弱速度最快约12msYOLOv8s640×640精度与速度平衡约22msYOLOv5s640×640大车效果好小车一般约18msNanoDet-M320×320可跑CPU约40ms2.3 用PyTorch跑起一个高速场景检测模型的最小命令下面的命令从空环境开始安装最小依赖并加载YOLO预训练权重。所有车辆类别用COCO里的car、bus、truck三个标签过滤。conda create -n traffic python3.10 -y conda activate traffic pip install torch torchvision pip install ultralytics opencv-pythonimport torch from ultralytics import YOLO model YOLO(yolov8n.pt) results model.predict( road_frame.jpg, imgsz640, conf0.35, classes[2, 5, 7], ) for box in results[0].boxes: print(box.xyxy.tolist(), box.conf.item(), int(box.cls.item()))conf0.35是检测置信度阈值卡车反光导致的误检通常集中在0.3到0.4之间设成0.4容易漏掉远处小目标0.35是稳妥起点。classes固定为COCO里的car2、bus5、truck7把person、bicycle等类别排除后模型不用输出无用的框。这里注意model.predict(road_frame.jpg)会自动做预处理但在实时循环中我会手动接管因为自动流程每次都要读文件并重建内存连续推流会产生小毛刺。下一章就把这部分拆开做。3. 视频流接入与预处理从RTSP摄像头到PyTorch的输入张量3.1 RTSP拉流与帧率控制的取舍摄像头输出的通常是H.264编码的RTSP流。写代码时最直接的方法是OpenCV的VideoCapture但要在实时统计中保持帧率平稳必须限制解码后的帧处理速度。常见做法是让解码线程和推理线程解耦解码线程把最新帧丢进deque只保留最近几帧推理线程每隔固定时间取一帧比如每150ms取一次也就是约6FPS的处理速度。为什么不用摄像头的25FPS因为目标检测模型在中等GPU上处理一帧需要20到40ms跟踪、计数和叠加显示还要额外花时间。如果硬压到全帧率系统延迟会逐渐变大。另一个坑是RTSP长期运行一定会断线没有重试逻辑的服务跑半天就卡死。import cv2 from collections import deque URL rtsp://192.0.2.10:554/channel1 cap cv2.VideoCapture(URL) cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) frame_queue deque(maxlen3) while True: ok, frame cap.read() if not ok: cap.release() cap cv2.VideoCapture(URL) continue frame_queue.append(frame)deque(maxlen3)的作用是缓存区满了以后新帧进入会自动丢弃最老帧这样处理慢的时候不会让内存无限制增长。CAP_PROP_BUFFERSIZE设为3可以降低OpenCV内部缓存引入的延迟让推理线程看到更接近当前时间的画面。如果画面出现花屏先检查网络丢包不要急着改代码重试逻辑。3.2 帧预处理与张量归一化模型训练时图像被letterbox到640×640推理时也必须保持同样变换否则检测框坐标会偏移。高速路摄像头一般是1920×1080直接resize会把车辆压扁所以先等比缩放再在两侧填充灰边。下面是完整预处理函数import cv2 import numpy as np import torch def frame_to_tensor(frame: np.ndarray, size: int 640) - torch.Tensor: h, w frame.shape[:2] scale min(size / w, size / h) nw, nh int(w * scale), int(h * scale) resized cv2.resize(frame, (nw, nh), interpolationcv2.INTER_LINEAR) canvas np.full((size, size, 3), 114, dtypenp.uint8) x_off (size - nw) // 2 y_off (size - nh) // 2 canvas[y_off:y_off nh, x_off:x_off nw] resized rgb cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB) tensor torch.from_numpy(rgb).permute(2, 0, 1).float().div(255.0) return tensor.unsqueeze(0)size通常取640和训练输入保持一致。填充色114对应灰度中高值不是随意选的YOLO官方预处理里用的就是这个值。permute(2, 0, 1)把OpenCV默认的HWC排列换成PyTorch需要的CHW排列div(255.0)把像素归一化到0到1区间。模型返回的框坐标是letterbox画布上的坐标要还原回原图需要记录scale、x_off、y_off在画目标框时做一次逆变换。这个细节在延迟分析中经常被忽略但错一点点叠加在视频上的线与实际车底位置就会偏移。3.3 自研数据标注和增强的最小方案预训练模型能覆盖高速路常见车辆但国内高速很多货箱形态特殊黄昏逆光场景也常让模型把车灯当成两个目标。最省力的做法是录一小时高峰期视频抽帧筛选500张用labelImg标注car、bus、truck三类。标注文件是YOLO格式的txt之后训练命令如下yolo detect train \ datatraffic.yaml \ modelyolov8n.pt \ epochs40 \ imgsz640 \ batch16 \ mosaic0.5 \ close_mosaic5traffic.yaml里声明路径和三个类别名。mosaic0.5表示一半训练批次使用拼接增强高速路背景比较统一mosaic开太高会让模型学到边框拼接痕迹。close_mosaic5表示最后5个epoch关闭mosaic让模型回到真实样本分布上做收敛。增强项推荐参数原因horizontal flip0.5车辆左右对称增加双向车道样本mosaic0.5增加小目标多样性过高会破坏语义mixup0.2降低遮挡样本过拟合再高会让框抖动hsv_h / hsv_s0.01 / 0.5适应不同色温车灯和路面颜色数据不均衡时不要只靠过采样。货车多就专门收集货车片段用save_crop把检测结果里的车框截出来人工筛掉误检后再加进训练集比单纯调loss权重更直接。4. 实时统计核心跟踪ID与虚拟线计数的完整逻辑4.1 为什么要跟踪车辆而不是直接数检测框如果每一帧都把检测框和虚拟线做相交判断那么同一辆车经过线的几百毫秒会触发好几次计数。高速路车距较小时误计会被放大。正确做法是引入跟踪器每一帧检测框与已有轨迹关联分配稳定ID。跨过一条线时只记录ID首次跨线的时间点之后即使该ID还在线附近也不再计数。当前常见选择是ByteTrack或Deep SORT。ByteTrack速度更快适合车流密集的快速路Deep SORT在有明显遮挡时更稳因为它同时使用运动和外观特征。这里我以Deep SORT为例因为它能对长时间遮挡的车辆保持更稳定的ID避免同一辆车因为短暂被卡车遮住就生成两个新ID造成双倍计数。4.2 Deep SORT车辆ID关联的关键参数Deep SORT的核心是把每个检测框与已有轨迹做级联匹配然后通过匈牙利算法做最终分配。工程里我们不需要重写匹配逻辑直接用封装好的库。from deep_sort_realtime.deepsort_tracker import DeepSort tracker DeepSort( max_age30, n_init3, nn_budget100, max_cosine_distance0.3, max_iou_distance0.7, ) detections [] for box in results[0].boxes: detections.append(( [int(x) for x in box.xyxy[0].tolist()], box.conf.item(), int(box.cls.item()), )) tracks tracker.update_tracks(detections, frameframe) for track in tracks: if not track.is_confirmed(): continue tid track.track_id ltrb track.to_ltrb()检测格式必须是[x1, y1, x2, y2, conf, class_id]坐标直接用模型输出不要提前缩放回原图否则和跟踪器内部的坐标换算对不上。max_age非常重要它决定轨迹在失去匹配后还能存活多少帧。车辆在高速路被大型货车遮挡的时间通常不会超过1秒按25FPS算就是25帧所以30帧是安全值。设太长车辆驶出画面后ID还会挂着计数会莫名往后延迟。参数推荐值影响max_age30遮挡多少帧后仍保持轨迹过大则ID悬挂太久n_init3连续匹配到几帧后确认轨迹过滤闪烁误检nn_budget100外观特征库容量控制内存占用max_cosine_distance0.3外观匹配阈值越严格越难恢复遮挡max_iou_distance0.7位置重合度阈值控制目标移动距离上限提示max_age设置超过30时车辆离开画面后ID仍会留在轨迹里虚拟线计数逻辑会把它当成还在等待穿越导致统计滞后。建议和虚拟线跨越判定共用同一个轨迹坐标字典。4.3 虚拟线计数与车流密度估算虚拟线一般设在车道末端沿水平方向画一条line_y 480。这里给出最简实现保留每个ID最近两个中心点判断是否与虚拟线相交。注意交点不是检测框底边而是车辆中心的固定点否则不同程度的车高会带来位置偏差。from collections import deque track_centers {} counted_ids set() line_y 480 count_60s 0 def update_traffic(tracks): global count_60s for track in tracks: if not track.is_confirmed(): continue ltrb track.to_ltrb() cx int((ltrb[0] ltrb[2]) / 2) cy int((ltrb[1] ltrb[3]) / 2) tid track.track_id if tid not in track_centers: track_centers[tid] deque(maxlen2) prev None if track_centers[tid]: prev track_centers[tid][-1] track_centers[tid].append((cx, cy)) if prev is None: continue if (prev[1] line_y cy) or (cy line_y prev[1]): if tid not in counted_ids: counted_ids.add(tid) count_60s 1prev是上一帧的中心点。跨线判断用当前点和上一帧点是否位于线的两侧能处理车辆在连续帧之间跳过线的情况。counted_ids在这个函数里不需要立刻清空但每分钟定时器触发时要记录count_60s并清空集合否则同一个ID只在第一次跨线时计数后面的车辆不受影响。车流密度估算可以在同一段代码里完成用每分钟跨线车辆数除以车道数和区间长度得到每公里车辆数。这个数值比单纯的车流量更能反映拥堵程度放到监控大屏上时也更容易被业务人员理解。5. 用TensorRT压缩模型后再校验计数误差5.1 用TensorRT固化和量化PyTorch模型模型在PyTorch里跑只能算完成前半段。要稳定跑到25FPS以上常见做法是先把模型导出为ONNX再用TensorRT生成engine。固定输入尺寸能让TensorRT做更好的显存规划同时也可以把精度降到FP16。FP16在车辆检测任务里损失很小可以默认开。yolo export modelyolov8n.pt formatonnx dynamicFalse imgsz640 trtexec --onnxyolov8n.onnx --saveEngineyolov8n.trt --fp16dynamicFalse很关键动态shape会引入额外优化开销固定640×640后推理速度最稳定。--fp16是半精度不建议直接用--int8INT8需要准备校准数据集校准集不足时检测精度会明显下滑高速路小目标尤其明显。导出的engine可以通过YOLO(yolov8n.trt)加载接口与PyTorch模型一致改造成本很低。5.2 用固定视频段校验车流量统计误差计数逻辑改完之后最有效的验证不是跑单张图片而是录制一段15分钟同时包含高峰期和稀疏期的原始视频记录人工计数和系统输出。比较指标有两个跨线计数误差率以及ID切换次数。若误差超过5%优先检查检测置信度和虚拟线位置是否偏移。python traffic_check.py --video highway_clip.mp4 --engine yolov8n.trt \ --line-y 480 --ground-truth 217常见的校验结果视频段人工计数系统计数误差率14:00-14:0578802.56%18:00-18:051391507.91%黄昏时段误差率往往会飙高原因是低亮度让模型给出的置信度下降部分车辆被滤掉。这时可以先把输入帧的gamma提高1.2再进模型或者把置信度从0.35降到0.30。两种方法都试一遍哪个误差低就留哪个。相比Python层反复调阈值先把输入亮度对齐训练集分布往往见效更快。本文还有配套的精品资源点击获取
返回列表