基于YOLOv5的目标分类计数与可视化系统实战指南

基于YOLOv5的目标分类计数与可视化系统实战指南 1. 项目概述与核心价值最近在做一个工业质检的小项目需要实时统计流水线上不同缺陷类型的数量并把结果直接“画”在监控画面上。这听起来是个很常见的需求对吧但真上手做你会发现从“检测出来”到“清晰、准确、实时地显示统计结果”中间有不少坑要踩。比如同一个目标在连续帧里被重复计数怎么办统计数字在复杂的背景上怎么显示才清晰易读动态更新的计数如何与视频流流畅结合我最初尝试用OpenCV简单写点逻辑很快就遇到了性能瓶颈和显示混乱的问题。后来我决定基于YOLOv5来构建这个系统。YOLOv5大家应该不陌生它凭借出色的速度和精度平衡在目标检测领域几乎是“标配”了。但官方仓库的演示更多是展示检测框和标签对于“分类计数”这个场景我们需要在它的输出基础上做二次加工和可视化。这不仅仅是调用一个detect.py那么简单它涉及到推理结果的解析、计数逻辑的设计、以及高性能的可视化渲染。这个项目就是把我趟过的路、踩过的坑以及最终稳定运行的方案梳理出来。无论你是做安防的人流量统计、交通车流分析还是像我一样的工业视觉应用这套把YOLOv5升级为“带计数器的可视化监控系统”的思路都能直接拿来用。2. 系统整体设计与思路拆解2.1 为什么选择YOLOv5作为核心检测器在开始动手之前我们先聊聊选型。目标检测框架那么多SSD、Faster R-CNN、YOLO系列等等为什么偏偏是YOLOv5对于计数显示这个场景核心诉求就三点快、准、稳。快速度优先我们的系统往往需要处理视频流实时性或准实时性是硬指标。YOLOv5的“You Only Look Once”单阶段检测架构天生就比两阶段检测器快。特别是其提供的s、m、l、x不同尺寸的模型让我们可以根据硬件算力比如你用的是边缘设备RK3568还是服务器GPU灵活选择在速度和精度间取得最佳平衡。如果要在RV1106这类NPU上部署YOLOv5的工程化支持和社区转换工具也相对成熟。准精度足够YOLOv5在COCO等通用数据集上的表现有目共睹更重要的是它非常容易训练自己的数据集。项目里我们通常要检测特定类别的目标如“划痕”、“漏焊”、“行人”、“轿车”YOLOv5的数据准备格式YOLO格式的txt标注文件简单明了训练脚本封装得好即便是新手也能相对顺利地完成模型迭代避免出现“训练map总是0”这种挫败感这个问题我们后面会专门讲如何排查。稳生态完善YOLOv5的PyTorch实现代码结构清晰推理接口detect.py或直接调用模型都非常方便。它输出的结果张量包含了我们计数所需的所有信息边界框坐标、置信度、类别ID。这为我们后续的计数逻辑处理提供了干净的数据源头。所以选择YOLOv5不是盲目跟风而是它在速度、精度、易用性这个“不可能三角”中为我们计数显示这个应用场景找到了一个非常可靠的立足点。2.2 计数与显示系统的核心模块设计有了强大的检测器我们还需要围绕它构建一个完整的处理流水线。整个系统可以分解为四个核心模块它们像流水线一样协同工作视频/图像输入模块负责读取数据源。可以是本地视频文件、RTSP网络流、USB摄像头或者单张图片。这个模块需要稳定能处理丢帧、断流等情况。YOLOv5推理模块核心检测环节。加载训练好的权重.pt文件对每一帧图像进行推理输出原始检测结果。智能计数处理模块这是项目的“大脑”也是区别于单纯检测的关键。它需要做三件事结果解析从YOLOv5的输出中提取出每一帧里所有检测到的目标的类别、位置和置信度。去重与跟踪防重复计数这是最大的坑点之一。如果一个目标在视频中静止或移动缓慢它会在连续多帧中被检测到。如果每帧都简单累加计数会爆炸式增长。因此必须引入防重复计数逻辑。我采用了一种基于IOU交并比和位置变化的简单跟踪方法效果很好后面会详细说。计数逻辑为每一个需要统计的类别如“person” “car”维护一个计数器。只有当某个目标被确认为“新出现”或“首次进入计数区域”时对应的计数器才加一。可视化渲染输出模块这是项目的“脸面”。负责把检测框、类别标签、以及最重要的——实时更新的分类计数——清晰美观地绘制到图像上并显示或保存。这里会用到OpenCV的绘图函数并涉及如何布局计数信息才不会遮挡关键画面。这个设计思路清晰地将检测、数据处理、显示解耦使得每一部分都可以独立优化和调试。3. 环境搭建与YOLOv5基础配置3.1 创建并配置Python环境我强烈建议使用Anaconda或Miniconda来管理Python环境避免包冲突。下面是一步步的操作# 1. 创建新的conda环境指定Python版本3.8是一个兼容性很好的选择 conda create -n yolov5_counting python3.8 # 2. 激活环境 conda activate yolov5_counting # 3. 安装PyTorch。请务必去PyTorch官网根据你的CUDA版本选择正确的安装命令。 # 例如如果你有CUDA 11.3可以安装 pip install torch1.12.1cu113 torchvision0.13.1cu113 torchaudio0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113 # 如果没有GPU则安装CPU版本 # pip install torch torchvision torchaudio注意PyTorch版本与CUDA版本的匹配至关重要不匹配会导致无法利用GPU甚至安装失败。使用nvidia-smi查看CUDA版本。3.2 克隆与安装YOLOv5YOLOv5的官方仓库更新活跃为了稳定性我建议克隆一个特定版本的分支。# 克隆YOLOv5仓库这里以v6.1版本为例相对稳定 git clone -b v6.1 https://github.com/ultralytics/yolov5.git cd yolov5 # 安装项目依赖 pip install -r requirements.txt安装完成后你可以快速测试一下YOLOv5是否正常工作python detect.py --source data/images/bus.jpg --weights yolov5s.pt --conf 0.25如果一切顺利会在runs/detect/exp目录下生成一张带有检测结果的bus.jpg图片。这证明YOLOv5的基础推理功能没问题了。3.3 准备或训练你自己的检测模型如果你要检测通用目标人、车等可以直接使用YOLOv5预训练的yolov5s.pt等模型。但工业场景通常需要训练自己的数据集。数据准备将你的图片和标注文件YOLO格式每个txt文件对应一张图片每行内容为class_id x_center y_center width height坐标已归一化按如下结构放置custom_dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/创建数据集配置文件创建一个custom_data.yaml文件内容如下# 数据集路径 path: /path/to/your/custom_dataset train: images/train val: images/val # 类别数量 nc: 2 # 例如2个类别划痕、污渍 # 类别名称列表 names: [scratch, stain]开始训练python train.py --img 640 --batch 16 --epochs 100 --data custom_data.yaml --weights yolov5s.pt--img 640: 输入图像尺寸。YOLOv5训练时会自动缩放到此尺寸。--batch 16: 批大小根据你的GPU内存调整。--epochs 100: 训练轮数。--data: 指向你刚创建的custom_data.yaml。--weights: 使用预训练的yolov5s.pt进行迁移学习这比从零训练快得多效果也好。实操心得解决“训练map总是0”这是新手常遇到的问题。map平均精度为0通常意味着模型根本没学到东西。请按以下顺序排查检查数据标注确保你的标注文件txt内容格式正确坐标值在0-1之间。可以用python -c “with open(‘label.txt’ ‘r’) as f: print(f.read())”快速查看。检查数据集配置确保custom_data.yaml中的path、train、val路径绝对正确并且nc类别数和names与实际完全匹配。检查图像和标签是否对应确保images/train里的每张图片在labels/train里都有一个同名的txt文件扩展名不同。降低学习率在train.py命令后添加--hyp data/hyps/hyp.scratch-low.yaml使用更低的学习率开始训练防止初期震荡。可视化训练数据使用python train.py --data custom_data.yaml --weights yolov5s.pt --epochs 1只跑一个epoch然后检查生成的train_batch*.jpg图片看标注框是否正确地画在了图片上。这是最直接的验证方法。训练完成后最好的模型权重会保存在runs/train/exp/weights/best.pt中这就是我们后续计数系统要用的模型。4. 核心计数逻辑的实现与防重设计4.1 解析YOLOv5的推理结果YOLOv5的推理输出是一个PyTorch张量其形状通常为[N, 6]其中N是当前帧检测到的目标数量6代表[x1, y1, x2, y2, confidence, class_id]左上角和右下角坐标、置信度、类别ID。我们的计数逻辑就从解析这个结果开始。import cv2 import torch from models.experimental import attempt_load from utils.general import non_max_suppression, scale_coords # 1. 加载模型 device torch.device(cuda:0 if torch.cuda.is_available() else cpu) model attempt_load(runs/train/exp/weights/best.pt, devicedevice) # 或 ‘yolov5s.pt’ model.eval() # 2. 预处理图像函数 def preprocess_image(img, img_size640): # 将OpenCV的BGR图像转换为RGB并调整大小、归一化等 # 这里省略具体细节YOLOv5的utils.datasets模块有现成函数 pass # 3. 推理与解析 def detect_and_parse(frame, model, device): img preprocess_image(frame).to(device) with torch.no_grad(): pred model(img)[0] # 原始预测 # 应用非极大值抑制 (NMS)过滤重叠框 pred non_max_suppression(pred, conf_thres0.25, iou_thres0.45)[0] detections [] if pred is not None: # 将坐标缩放回原始图像尺寸 pred[:, :4] scale_coords(img.shape[2:], pred[:, :4], frame.shape).round() for *xyxy, conf, cls in pred: # xyxy: 边界框坐标 # conf: 置信度 # cls: 类别ID (整数) detections.append({ bbox: [int(x) for x in xyxy], confidence: float(conf), class_id: int(cls), class_name: model.names[int(cls)] # 获取类别名称 }) return detections这样每一帧的检测结果detections就是一个列表里面包含了每个目标的详细信息。4.2 实现防重复计数的跟踪器简单每帧累加计数会导致数量虚高。我们需要一个跟踪器来关联连续帧中的同一个目标。这里实现一个轻量化的基于IOU的跟踪器。class SimpleTracker: def __init__(self, iou_threshold0.5, max_frames_to_skip10): 初始化跟踪器。 :param iou_threshold: 判断是否为同一目标的IOU阈值 :param max_frames_to_skip: 目标最大丢失帧数超过则视为消失 self.iou_threshold iou_threshold self.max_frames_to_skip max_frames_to_skip self.tracked_objects {} # 格式: {track_id: {bbox: [x1,y1,x2,y2], ‘class_id’: cid, ‘frames_skipped’: 0}} self.next_id 0 self.class_counters {} # 分类计数器: {class_name: count} def _calculate_iou(self, box1, box2): 计算两个矩形框的IOU交并比 x1 max(box1[0], box2[0]) y1 max(box1[1], box2[1]) x2 min(box1[2], box2[2]) y2 min(box1[3], box2[3]) inter_area max(0, x2 - x1) * max(0, y2 - y1) box1_area (box1[2] - box1[0]) * (box1[3] - box1[1]) box2_area (box2[2] - box2[0]) * (box2[3] - box2[1]) union_area box1_area box2_area - inter_area return inter_area / union_area if union_area 0 else 0 def update(self, detections): 更新跟踪器状态。 :param detections: 当前帧的检测结果列表每个元素是包含‘bbox’ ‘class_id’等的字典。 :return: 带有track_id的检测结果列表以及更新后的计数。 current_objects {} matched_new_detections [] # 为每个跟踪目标增加丢失帧数 for track_id, obj in self.tracked_objects.items(): obj[frames_skipped] 1 # 遍历当前帧的所有检测结果 for det in detections: best_iou self.iou_threshold best_match_id None # 在现有跟踪目标中寻找匹配同一类别且IOU最大 for track_id, obj in self.tracked_objects.items(): if obj[class_id] ! det[class_id]: continue iou self._calculate_iou(obj[bbox], det[bbox]) if iou best_iou: best_iou iou best_match_id track_id if best_match_id is not None: # 匹配成功更新该目标的位置并重置丢失帧数 self.tracked_objects[best_match_id][bbox] det[bbox] self.tracked_objects[best_match_id][frames_skipped] 0 det[track_id] best_match_id current_objects[best_match_id] self.tracked_objects[best_match_id] matched_new_detections.append(det) else: # 没有匹配到视为新目标分配新ID new_id self.next_id self.next_id 1 det[track_id] new_id new_obj {bbox: det[bbox], ‘class_id’: det[‘class_id’], ‘frames_skipped’: 0} current_objects[new_id] new_obj matched_new_detections.append(det) # **关键新目标出现分类计数器加一** class_name det[class_name] self.class_counters[class_name] self.class_counters.get(class_name, 0) 1 # 清理丢失超过阈值的跟踪目标 self.tracked_objects {} for track_id, obj in current_objects.items(): if obj[frames_skipped] self.max_frames_to_skip: self.tracked_objects[track_id] obj return matched_new_detections, self.class_counters.copy()这个SimpleTracker的工作原理是维护一个tracked_objects字典来记录每个被跟踪目标的最新位置和类别。每一帧它将新的检测框与已有的跟踪目标进行IOU匹配。如果匹配成功IOU大于阈值且类别相同则更新该目标的位置如果匹配失败则认为是新目标为其分配新ID并且触发对应类别的计数器加一。同时跟踪器会记录目标连续未出现的帧数超过max_frames_to_skip就将其移除防止跟踪列表无限膨胀。注意事项这个简易跟踪器适用于目标运动平缓、遮挡不严重的场景。对于高速运动或严重遮挡可能需要更复杂的算法如Kalman滤波匈牙利匹配的SORT/DeepSORT但复杂度也会大大增加。对于很多工业计数场景如传送带上的产品这个简易版已经足够稳健。5. 可视化渲染将计数结果清晰显示在图像上5.1 绘制检测框与标签这是基础工作YOLOv5的utils.plots模块提供了很好的函数我们可以借鉴并自定义。import cv2 import numpy as np def plot_one_box(bbox, img, colorNone, labelNone, line_thickness2): 在图像上绘制一个边界框和标签 tl line_thickness or round(0.002 * (img.shape[0] img.shape[1]) / 2) 1 color color or [random.randint(0, 255) for _ in range(3)] c1, c2 (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])) cv2.rectangle(img, c1, c2, color, thicknesstl, lineTypecv2.LINE_AA) if label: tf max(tl - 1, 1) # 字体粗细 t_size cv2.getTextSize(label, 0, fontScaletl / 3, thicknesstf)[0] c2 c1[0] t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # 填充的标签背景 cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thicknesstf, lineTypecv2.LINE_AA) return img5.2 设计并绘制分类计数信息面板这是本项目的亮点。我们需要把动态更新的分类计数以清晰、直观、不碍事的方式显示出来。我推荐两种主流布局顶部/底部条形统计栏在图像上方或下方绘制一个半透明的色块里面用文字显示各个类别的累计数量。适合类别不多的情况。侧边栏统计面板在图像右侧开辟一个垂直区域以列表形式显示类别和计数。适合画面宽度充足且不想遮挡上下内容的情况。这里以顶部条形统计栏为例def draw_counting_info(image, class_counters, tracker): 在图像顶部绘制计数信息栏。 :param image: 原始图像 (numpy array) :param class_counters: 分类计数器字典 {‘class_name’: count} :param tracker: 跟踪器实例用于显示当前跟踪目标数可选 h, w image.shape[:2] # 1. 创建一个顶部区域例如高度为60像素 info_bar_height 60 info_bar np.zeros((info_bar_height, w, 3), dtypenp.uint8) # 可以设置一个半透明的深色背景 info_bar[:] (40, 40, 40) # 深灰色 # 2. 定义显示的文本 font cv2.FONT_HERSHEY_SIMPLEX font_scale 0.6 thickness 2 color (255, 255, 255) # 白色文字 # 构建显示字符串 count_texts [] for cls_name, count in sorted(class_counters.items()): count_texts.append(f{cls_name}: {count}) # 可以加上总计数或跟踪目标数 total_count sum(class_counters.values()) count_texts.append(fTotal: {total_count}) if tracker: count_texts.append(fTracking: {len(tracker.tracked_objects)}) # 3. 将文本绘制到信息栏上 y_offset 30 # 文字基线起始纵坐标 for i, text in enumerate(count_texts): text_size cv2.getTextSize(text, font, font_scale, thickness)[0] x_offset 10 i * 150 # 横向排列每个类别间隔150像素 cv2.putText(info_bar, text, (x_offset, y_offset), font, font_scale, color, thickness) # 4. 将信息栏与原始图像拼接 final_image np.vstack([info_bar, image]) return final_image5.3 整合与主循环最后我们将所有模块串联起来形成完整的处理流程。def main(video_source0): # video_source可以是0摄像头 视频文件路径 或RTSP流地址 # 初始化 cap cv2.VideoCapture(video_source) tracker SimpleTracker(iou_threshold0.5) while cap.isOpened(): ret, frame cap.read() if not ret: print(视频流结束或读取失败。) break # 1. YOLOv5推理与解析 detections detect_and_parse(frame, model, device) # 2. 更新跟踪器并获取计数 tracked_dets, current_counts tracker.update(detections) # 3. 在原始帧上绘制检测框和跟踪ID for det in tracked_dets: label f{det[class_name]} {det[confidence]:.2f} ID:{det.get(track_id, -1)} # 为不同类别分配固定颜色便于观察 colors {scratch: (0, 0, 255), ‘stain’: (0, 255, 0)} # 红绿 color colors.get(det[class_name], (255, 0, 0)) plot_one_box(det[bbox], frame, colorcolor, labellabel) # 4. 绘制计数信息栏 frame_with_info draw_counting_info(frame, current_counts, tracker) # 5. 显示结果 cv2.imshow(YOLOv5 Object Counting, frame_with_info) # 按‘q’退出 if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows() if __name__ __main__: main(your_video.mp4) # 替换为你的视频源运行这个脚本你就能看到一个实时视频窗口里面不仅有YOLOv5检测到的目标框和类别标签还有在画面顶部实时更新的分类计数。同一个目标在画面中移动时其ID保持不变只有当它首次出现时计数器才会增加完美解决了重复计数的问题。6. 性能优化与高级功能拓展6.1 推理速度优化技巧当处理高分辨率视频或多路视频流时推理速度可能成为瓶颈。以下是一些立竿见影的优化手段模型轻量化优先使用yolov5s.pt甚至yolov5n.pt如果有这类小模型。在精度可接受的范围内模型越小越快。调整推理尺寸在detect_and_parse函数中preprocess_image的img_size参数直接影响速度。默认640x640可以尝试降低到480x480甚至320x320速度会显著提升但小目标检测能力会下降。启用半精度推理PyTorch支持FP16半精度推理能在支持Tensor Core的GPU上大幅提升速度且精度损失很小。model attempt_load(weights, devicedevice) model.half() # 转换为半精度 # 预处理时图像数据也需要转换为半精度 img img.half() if device.type ! ‘cpu’ else img.float()批处理如果有多帧图片可以同时处理使用批处理能更充分利用GPU。但对于实时视频流通常是一帧一帧处理。6.2 实现区域计数ROI Counting很多时候我们只关心特定区域的计数比如只统计进入某个警戒区的人数。这需要增加一个区域判断逻辑。def is_in_roi(bbox, roi_polygon): 判断目标中心点是否在指定多边形区域内。 :param bbox: [x1, y1, x2, y2] :param roi_polygon: 由多个点组成的列表表示多边形区域例如 [(x1,y1), (x2,y2), ...] :return: Boolean from shapely.geometry import Point, Polygon center_x (bbox[0] bbox[2]) / 2 center_y (bbox[1] bbox[3]) / 2 point Point(center_x, center_y) polygon Polygon(roi_polygon) return polygon.contains(point) # 在跟踪器update函数中只有当目标在ROI内且是新目标时才计数 if best_match_id is None: if is_in_roi(det[bbox], predefined_roi): # ... 分配新ID并增加计数器 ...同时可以在可视化时将这个ROI区域画出来方便调试。# 绘制ROI区域 roi_points np.array(predefined_roi, np.int32).reshape((-1, 1, 2)) cv2.polylines(frame, [roi_points], isClosedTrue, color(0, 255, 255), thickness2)6.3 结果记录与输出除了实时显示我们通常还需要将计数结果保存下来用于生成报表或进一步分析。日志记录可以使用Python的logging模块或简单写文件定期如每秒或当计数变化时将class_counters写入文本或CSV文件。import csv import time def log_counts(counters, filenamecount_log.csv): with open(filename, a, newline) as f: writer csv.writer(f) timestamp time.strftime(%Y-%m-%d %H:%M:%S) row [timestamp] [counters.get(cls, 0) for cls in [scratch, ‘stain’]] # 按固定顺序 writer.writerow(row)视频保存使用cv2.VideoWriter将带计数结果的画面保存为新视频。fourcc cv2.VideoWriter_fourcc(*mp4v) # 或 ‘XVID’ out cv2.VideoWriter(output_with_counting.mp4, fourcc, 20.0, (frame_width, frame_height_with_bar)) # 在主循环中将每一帧frame_with_info写入out out.write(frame_with_info)7. 常见问题排查与调试技巧在实际部署中你可能会遇到以下问题。这里是我的“踩坑”记录计数不准忽多忽少可能原因跟踪器的iou_threshold设置不当。阈值太高同一个目标轻微移动就可能被判定为新目标阈值太低两个靠近的不同目标可能被误认为同一个。解决根据你的场景调整iou_threshold。对于静止或慢速目标可以设高些如0.7对于快速运动目标设低些如0.3。在SimpleTracker的update方法里打印匹配的IOU值观察分布。可能原因max_frames_to_skip太小。目标被短暂遮挡如被人走过挡住后跟踪器立即将其删除当它再次出现时被当作新目标计数。解决适当增加这个值例如从10调到30给目标更长的“消失容忍时间”。画面卡顿延迟高可能原因推理速度慢。使用time.time()测量detect_and_parse函数的耗时。解决应用6.1节的优化技巧。首先尝试将模型换成更小的版本如s-n这是最有效的方法。其次尝试降低推理分辨率。可能原因OpenCV的imshow在高分辨率下本身有性能开销。解决可以降低显示窗口的分辨率或者减少imshow的调用频率如每处理2帧显示1帧。计数面板显示混乱或重叠可能原因类别名称过长或类别数量太多导致文本超出画面宽度。解决在draw_counting_info函数中动态计算文本位置或者改用侧边栏垂直排列。也可以考虑只显示数量超过0的类别。在边缘设备如RK3568上部署缓慢说明在ARM架构的边缘设备上直接运行PyTorch版的YOLOv5通常很慢。需要将PyTorch模型转换为该平台专用的推理引擎格式如RKNN for Rockchip, NCNN for 其他平台。建议这属于模型部署优化范畴是一个独立的大课题。基本流程是PyTorch (.pt) - ONNX (.onnx) - 目标平台模型.rknn/.param.bin。你需要参考对应芯片厂商的官方文档和转换工具。转换后计数和显示的逻辑Python部分基本可以复用只需替换模型加载和推理的代码为调用对应的推理引擎API。如何调试和可视化中间结果打印关键数据在tracker.update()函数中打印detections的数量、匹配情况、class_counters的变化。绘制跟踪轨迹对于每个track_id可以将其历史中心点保存下来并在帧上连成线直观看到跟踪效果。保存问题帧当计数出现异常时可以将当前帧和前后几帧保存为图片分析检测框和跟踪ID的变化这是定位问题最直接的方法。这个基于YOLOv5的目标分类计数与显示系统从核心的防重复计数逻辑到友好的可视化界面基本覆盖了实际应用中的主要需求。它不是一个僵化的代码而是一个框架你可以根据自己的具体场景调整参数、优化性能、增加功能如越界报警、数据持久化。希望这份详细的梳理能帮你快速搭建起属于自己的智能计数系统。