ARTICLE DETAIL

资讯详情

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

YOLOv5目标检测分类计数与可视化实战:从原理到工程实现

YOLOv5目标检测分类计数与可视化实战:从原理到工程实现 1. 项目缘起从“看见”到“数清”的刚需在计算机视觉的日常开发里目标检测是基础中的基础。我们常常用YOLOv5这样的模型在图像或视频里框出一个个目标比如生产线上的零件、果园里的苹果、或者监控画面中的人。但很多时候光“框出来”还不够业务方紧接着就会问“那到底有多少个”——这就是目标计数。更进一步他们希望这个数字能直接、清晰地显示在结果图上方便报告、存档或者实时监控。这个从“检测”到“分类计数并可视化”的需求在工业质检、农业估产、智慧安防等领域几乎是标配。然而把YOLOv5的检测结果直接拿来计数新手很容易踩进几个坑里。最常见的就是重复计数同一个苹果因为模型在不同帧或略有差异的推理中给出了两个略有重叠的框就被数了两次。另一个头疼的问题是显示杂乱原始的检测结果图所有识别框和类别标签都堆在一起计数信息要么没有要么挤在角落根本看不清。网上能找到的很多教程要么只讲检测要么计数逻辑写得过于简单没考虑这些实际工程问题。所以今天我就结合自己多次在项目里折腾的经验来聊聊怎么用YOLOv5不仅实现高精度的目标检测更能完成准确的目标分类计数并且把结果以清晰、美观的方式叠加显示在原始图像上。我会重点拆解如何避免重复计数、如何设计计数信息的显示布局以及如何将整个流程封装成便于调用的模块。你会发现要实现这个功能核心代码可能也就百来行但里面的细节和思路才是真正决定项目成败的关键。2. 环境搭建与YOLOv5基础避开初学者的第一个坑工欲善其事必先利其器。第一步的环境准备很多人觉得简单照着官方README做就行但恰恰这里埋伏着几个导致后续各种诡异错误的“暗桩”。2.1 创建并激活独立的Python环境强烈建议不要用系统全局的Python环境。使用conda或venv创建一个独立环境能完美隔离不同项目间的依赖冲突。这里以conda为例# 创建一个名为yolo_count的新环境指定Python版本推荐3.8或3.9兼容性最好 conda create -n yolo_count python3.8 -y conda activate yolo_count2.2 克隆YOLOv5仓库与安装依赖YOLOv5的官方仓库更新频繁直接pip install yolov5并不是官方推荐的做法因为那样安装的是别人打包的库可能缺少训练、导出等脚本。正确做法是克隆整个仓库# 克隆官方仓库 git clone https://github.com/ultralytics/yolov5.git cd yolov5 # 安装requirements.txt中列出的依赖 pip install -r requirements.txt注意requirements.txt里默认的torch和torchvision通常是CPU版本。如果你有NVIDIA GPU并需要CUDA加速务必在安装完其他依赖后去 PyTorch官网 根据你的CUDA版本获取对应的安装命令。例如对于CUDA 11.3pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu1132.3 验证安装与初步测试安装完成后不要急着写自己的代码。先用官方提供的预训练模型和示例图片跑一个最简单的检测验证环境是否正常。import torch # 加载官方预训练模型这里用中等大小的yolov5s model torch.hub.load(ultralytics/yolov5, yolov5s, pretrainedTrue) # 进行推理 img https://ultralytics.com/images/zidane.jpg # 使用YOLOv5官方示例图片 results model(img) # 显示结果 results.show() # 会弹出窗口显示带检测框的图片 results.print() # 在控制台打印检测到的目标信息如果这一步能成功弹出图片并在控制台看到类似“person 0.89...”的输出恭喜你基础环境没问题了。这个results对象里就包含了我们后续实现计数和显示所需的所有原始数据。3. 核心逻辑拆解从原始输出到结构化计数数据YOLOv5模型推理后返回的results对象是个宝藏但它的数据结构需要理解。我们实现计数的第一步就是正确地从中提取并处理信息。3.1 理解Results对象的构成执行results.print()你可能会看到这样的输出image 1/1: 720x1280 2 persons, 1 tie, 1 handbag Speed: 10.2ms pre-process, 12.1ms inference, 1.2ms NMS per image at shape (1, 3, 384, 640)这只是摘要。详细数据藏在results.pandas().xyxy[0]里如果你更喜欢用Pandas DataFrame或者直接访问results.xyxy[0]返回一个Torch Tensor。我更倾向于直接使用Tensor因为速度更快也便于后续的数值计算。results.xyxy[0]的形状通常是[N, 6]其中N是检测到的目标数量6列分别代表x_min,y_min,x_max,y_max: 边界框的左上角和右下角坐标。confidence: 检测置信度范围0-1。class: 目标类别索引整数对应模型训练时data.yaml里names列表的顺序。3.2 实现基于IOU的去重计数这是避免重复计数的关键。同一个物理目标由于模型不确定性或视频相邻帧间目标移动很小可能会产生多个高度重叠的检测框。我们需要用非极大值抑制NMS的思想来进行去重。虽然YOLOv5在推理时已经做过一次NMS但其阈值iou_thres参数默认0.45可能对于你的特定场景还不够严格或者你需要对最终结果再做一次聚合。这里我们实现一个基于类内IOU的合并与计数函数。IOUIntersection over Union交并比计算两个框的重叠程度。import torch def count_objects(detections, iou_threshold0.5): 对检测结果进行去重计数。 Args: detections (torch.Tensor): [N, 6] 格式为 [x1, y1, x2, y2, conf, cls] iou_threshold (float): IOU阈值高于此值则认为两个框是同一目标。 Returns: dict: 键为类别索引int值为该类别的计数int。 torch.Tensor: 去重后的检测框 [M, 6], M N。 if detections.shape[0] 0: return {}, detections # 按置信度降序排序优先保留置信度高的框 detections detections[detections[:, 4].argsort(descendingTrue)] keep [] # 保留的框的索引 counted_detections [] # 用于存储最终保留的框 # 获取所有类别 unique_classes detections[:, 5].unique() count_dict {int(cls): 0 for cls in unique_classes} for cls in unique_classes: cls_mask (detections[:, 5] cls) cls_detections detections[cls_mask] while len(cls_detections) 0: # 取出当前置信度最高的框因为已排序第一个就是 current_box cls_detections[0] keep.append(torch.where((detections[:, 5] cls) (detections[:, 4] current_box[4]))[0][0].item()) counted_detections.append(current_box) count_dict[int(cls)] 1 # 计数1 if len(cls_detections) 1: break # 计算当前框与剩余同类框的IOU other_boxes cls_detections[1:] ious calculate_iou(current_box.unsqueeze(0), other_boxes) # 找出IOU过低的框即不同的目标保留下来进行下一轮循环 low_iou_mask ious iou_threshold cls_detections other_boxes[low_iou_mask] counted_detections torch.stack(counted_detections) if counted_detections else torch.empty((0, 6)) return count_dict, counted_detections def calculate_iou(box1, box2): 计算两组框之间的IOU。 box1: [1, 4] (x1, y1, x2, y2) box2: [N, 4] # 计算交集区域的坐标 inter_x1 torch.max(box1[:, 0], box2[:, 0]) inter_y1 torch.max(box1[:, 1], box2[:, 1]) inter_x2 torch.min(box1[:, 2], box2[:, 2]) inter_y2 torch.min(box1[:, 3], box2[:, 3]) # 计算交集面积 inter_area torch.clamp(inter_x2 - inter_x1, min0) * torch.clamp(inter_y2 - inter_y1, min0) # 计算各自面积 area_box1 (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1]) area_box2 (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1]) # 计算并集面积和IOU union_area area_box1 area_box2 - inter_area iou inter_area / (union_area 1e-6) # 加一个极小值防止除零 return iou.squeeze()这个count_objects函数做了几件事1按类别分别处理2在同类框里根据IOU阈值判断是否属于同一物体3只保留每个“物体簇”中置信度最高的那个框作为代表并计数一次。通过调整iou_threshold你可以控制去重的严格程度。对于静止图片阈值可以设高一点如0.6对于视频连续帧可能需要根据帧率适当调低。3.3 整合类别名称映射计数字典的键是类别索引我们需要把它转换成人类可读的类别名。YOLOv5模型加载后其类别名存储在model.names中这是一个列表。例如COCO预训练模型的model.names[0]是“person”。def get_count_with_names(count_dict, model): 将计数字典中的类别索引转换为类别名称。 names model.names result {} for cls_idx, count in count_dict.items(): cls_name names[int(cls_idx)] result[cls_name] count return result现在我们已经有了一个干净、去重后的检测框张量counted_detections以及一个按类别名称统计的计数字典count_with_names。接下来就是如何把这些信息美观地画到图上了。4. 结果可视化在图像上清晰呈现计数信息可视化部分的目标是生成一张既包含所有检测框又突出显示统计结果的图片。直接使用results.show()或results.render()得到的图片计数信息并不明显。我们需要自定义绘制逻辑。4.1 使用OpenCV进行高级绘制我们将使用OpenCV已在requirements.txt中安装来绘制。主要步骤是1将YOLOv5的结果图片RGB格式转换为OpenCV格式BGR2绘制去重后的检测框和标签3在图像的固定位置如左上角绘制一个半透明的信息板显示各类别的计数。import cv2 import numpy as np def visualize_detections_with_count(original_img, detections, count_dict, model, conf_threshold0.25): 在图像上绘制检测框并叠加计数信息板。 Args: original_img (numpy.ndarray): 原始RGB图像通常来自results.imgs[0]。 detections (torch.Tensor): 去重后的检测框 [M, 6]。 count_dict (dict): 类别名称到计数的映射。 model: YOLOv5模型对象用于获取颜色和名称。 conf_threshold: 置信度阈值低于此值的框不绘制。 Returns: numpy.ndarray: 绘制好的BGR图像可供保存或显示。 # 复制图像避免修改原图 img_draw original_img.copy() # 如果原图是RGBYOLOv5输出通常是的转换为BGR供OpenCV使用 if img_draw.shape[2] 3: img_draw cv2.cvtColor(img_draw, cv2.COLOR_RGB2BGR) h, w img_draw.shape[:2] names model.names colors [[np.random.randint(0, 255) for _ in range(3)] for _ in names] # 1. 绘制检测框和标签 for *xyxy, conf, cls in detections: if conf conf_threshold: continue label f{names[int(cls)]} {conf:.2f} # 为每个类别分配固定颜色确保一致性 color colors[int(cls)] # 画矩形框 cv2.rectangle(img_draw, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), color, 2) # 计算文本背景大小 (text_width, text_height), baseline cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2) # 画文本背景 cv2.rectangle(img_draw, (int(xyxy[0]), int(xyxy[1]) - text_height - baseline - 5), (int(xyxy[0]) text_width, int(xyxy[1])), color, -1) # -1表示填充 # 写文本 cv2.putText(img_draw, label, (int(xyxy[0]), int(xyxy[1]) - baseline - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2) # 2. 在左上角绘制计数信息板 # 信息板参数 panel_x, panel_y 20, 20 panel_width, panel_height 220, 50 len(count_dict) * 30 # 创建一个半透明矩形区域 overlay img_draw.copy() cv2.rectangle(overlay, (panel_x, panel_y), (panel_x panel_width, panel_y panel_height), (50, 50, 50), -1) alpha 0.6 # 透明度 img_draw cv2.addWeighted(overlay, alpha, img_draw, 1 - alpha, 0) # 绘制信息板边框和标题 cv2.rectangle(img_draw, (panel_x, panel_y), (panel_x panel_width, panel_y panel_height), (200, 200, 200), 2) cv2.putText(img_draw, Detection Counts:, (panel_x 10, panel_y 30), cv2.FONT_HERSHEY_DUPLEX, 0.7, (255, 255, 255), 2) # 绘制每一类别的计数 y_offset panel_y 60 for idx, (cls_name, count) in enumerate(count_dict.items()): color colors[list(names.values()).index(cls_name)] if cls_name in names.values() else (255, 255, 255) count_text f{cls_name}: {count} cv2.putText(img_draw, count_text, (panel_x 20, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) y_offset 30 # 3. 可选在图像底部中央绘制总计数 total_count sum(count_dict.values()) total_text fTotal Objects: {total_count} (total_width, total_height), _ cv2.getTextSize(total_text, cv2.FONT_HERSHEY_DUPLEX, 1.0, 3) cv2.putText(img_draw, total_text, (w // 2 - total_width // 2, h - 30), cv2.FONT_HERSHEY_DUPLEX, 1.0, (0, 200, 255), 3) return img_draw这个可视化函数做了三件事首先用不同颜色绘制每个检测框及其类别、置信度标签其次在左上角创建了一个半透明的信息面板清晰列出每个类别的具体数量最后在图片底部中央醒目地显示了目标总数。这样的布局信息层次分明无论是用于屏幕显示还是生成报告图片都非常清晰。4.2 封装成完整流程函数现在我们把前面的所有步骤整合成一个函数实现“输入图片路径输出带计数的结果图”的完整流程。def detect_count_and_visualize(model, img_path, iou_thres0.5, conf_thres0.25): 端到端的检测、计数与可视化流程。 # 1. 推理 results model(img_path) # 获取原始检测张量 (xyxy格式) detections results.xyxy[0] # 2. 去重计数 count_dict_idx, filtered_dets count_objects(detections, iou_thresholdiou_thres) # 3. 转换类别索引为名称 count_dict_name get_count_with_names(count_dict_idx, model) # 4. 可视化 # results.imgs[0] 是推理用的原始图像经过预处理和resize的 # 为了获得原始尺寸的绘制我们通常直接用原始图像或results.render()[0] # 这里使用results.render()返回的带原始框的图片列表第一张 rendered_imgs results.render() # 返回一个列表每个元素是一个带框的RGB图像 if rendered_imgs: orig_img_with_boxes rendered_imgs[0] else: # 如果render()没返回则用原始推理图像 orig_img_with_boxes results.imgs[0] final_img visualize_detections_with_count(orig_img_with_boxes, filtered_dets, count_dict_name, model, conf_thresholdconf_thres) # 5. 返回结果 return final_img, count_dict_name, filtered_dets # 使用示例 model torch.hub.load(ultralytics/yolov5, yolov5s, pretrainedTrue) result_img, counts, boxes detect_count_and_visualize(model, your_image.jpg) # 显示图片 cv2.imshow(Result, result_img) cv2.waitKey(0) cv2.destroyAllWindows() # 保存图片 cv2.imwrite(result_with_count.jpg, result_img) print(Counts:, counts)5. 进阶优化与实战踩坑点把基础流程跑通只是第一步要让这个功能在实际项目中稳定可靠还需要考虑更多细节。5.1 处理视频流与实时计数对于视频文件或摄像头实时流核心逻辑不变但需要处理帧与帧之间计数结果的平滑问题。直接对每一帧独立计数会导致数字频繁跳动。一个常见的优化是使用滑动窗口或简单滤波。from collections import deque import time class SmoothCounter: def __init__(self, window_size5): self.window_size window_size self.history deque(maxlenwindow_size) # 存储最近N帧的计数字典 def update(self, current_count_dict): 更新历史记录并返回平滑后的计数取最近N帧的平均值或众数。 self.history.append(current_count_dict) if len(self.history) self.window_size: return current_count_dict # 平滑策略取众数出现次数最多的值 smoothed {} all_classes set() for d in self.history: all_classes.update(d.keys()) for cls in all_classes: counts [d.get(cls, 0) for d in self.history] # 取出现次数最多的值作为平滑结果 from collections import Counter smoothed[cls] Counter(counts).most_common(1)[0][0] return smoothed # 在视频处理循环中使用 counter SmoothCounter(window_size10) cap cv2.VideoCapture(0) # 打开摄像头 while True: ret, frame cap.read() if not ret: break # 将BGR帧转换为RGBYOLOv5期望RGB输入 rgb_frame cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 推理、计数使用前面定义的函数但传入numpy数组 results model(rgb_frame) detections results.xyxy[0] count_dict_idx, filtered_dets count_objects(detections) count_dict_name get_count_with_names(count_dict_idx, model) # 平滑计数 smoothed_counts counter.update(count_dict_name) # 使用平滑后的计数进行可视化 # ... (可视化代码传入smoothed_counts)5.2 针对特定场景的调参经验置信度阈值conf_thres默认0.25。在干净背景下检测大目标可以提高到0.5以上以减少误检在复杂、小目标场景下可能需要降低到0.1甚至0.05但同时要配合更严格的后续过滤如根据目标尺寸。NMS IOU阈值iou_thres默认0.45。这个值控制检测阶段框合并的激进程度。如果你发现同一个目标被拆分成多个小框比如一个行人被分成头和身体两个框可以适当提高这个值如0.6让它们更容易合并。反之如果不同目标靠得太近被合并成一个框就降低这个值。计数去重IOU阈值这是我们自定义函数里的iou_threshold。它应该大于或等于推理时的NMS IOU阈值。通常设置在0.5-0.7之间。对于静态图片用0.6对于视频可以结合目标跟踪如使用ByteTrack、DeepSORT来跨帧关联目标实现更准确的计数而不是每帧独立去重。5.3 常见问题与排查“训练map总是0”的启示虽然本篇重点在推理和可视化但标题相关热词中提到了“yolov5训练map总是0”。这提醒我们计数的前提是检测模型要足够准确。如果你的模型训练时mAP平均精度为0计数自然无从谈起。这里快速过一下可能的原因数据标注问题检查标注文件YOLO格式的.txt文件内容是否正确。坐标是否归一化0-1之间类别索引是否从0开始且连续标注框是否完全包含了目标数据路径配置data.yaml文件里的train和val路径是否正确建议使用绝对路径。类别不匹配data.yaml中的names列表是否与标注文件中的类别索引对应nc类别数设置是否正确学习率过高过高的初始学习率可能导致训练发散。尝试使用更小的学习率如--lr 0.001。模型复杂度与数据量不匹配数据量很少时使用过大的模型如yolov5x容易过拟合。从小模型yolov5n或yolov5s开始尝试。检查训练输出关注训练日志中损失值box_loss, obj_loss, cls_loss的变化趋势。它们应该随着训练轮数逐渐下降并趋于平稳。如果损失值一开始就是NaN或者异常大基本可以断定是数据或配置问题。5.4 部署与性能考量当需要将这套系统部署到边缘设备如RK3568、RV1106等时不能直接使用Python脚本。你需要模型导出使用YOLOv5提供的export.py脚本将PyTorch模型导出为ONNX、TensorRT或OpenVINO等格式。python export.py --weights yolov5s.pt --include onnx --img 640 --batch 1推理引擎重写在C或相应的边缘计算框架中加载导出的模型并重新实现前处理图像缩放、归一化、推理、后处理解码输出、NMS以及我们上面讨论的计数和显示逻辑。后处理中的框解码和NMS是性能关键点务必用高效的方式实现。显示部分在嵌入式设备上可能没有GUI。可视化结果可以保存为图片或者通过RTSP流推送到网络在远程客户端查看。6. 完整代码示例与使用指南最后我将提供一个整合了所有功能的脚本并附上详细的注释和使用说明。你可以将此脚本保存为yolov5_count_display.py并根据你的需求修改。 YOLOv5目标分类计数与图像显示 - 完整实现脚本 作者一个爱折腾的开发者 功能加载YOLOv5模型对输入图像/视频进行目标检测、去重计数并将结果可视化保存。 import torch import cv2 import numpy as np from pathlib import Path import argparse def count_objects(detections, iou_threshold0.5): 基于IOU的去重计数函数。 # ... (函数体同上文此处省略以节省篇幅实际使用时请复制完整函数) pass def calculate_iou(box1, box2): 计算IOU。 # ... (函数体同上文) pass def get_count_with_names(count_dict, model): 将类别索引转换为名称。 # ... (函数体同上文) pass def visualize_detections_with_count(original_img, detections, count_dict, model, conf_threshold0.25): 可视化检测结果与计数信息。 # ... (函数体同上文) pass def process_image(model, img_path, output_dir./output, iou_thres0.5, conf_thres0.25): 处理单张图片。 Path(output_dir).mkdir(parentsTrue, exist_okTrue) results model(img_path) detections results.xyxy[0] count_dict_idx, filtered_dets count_objects(detections, iou_thresholdiou_thres) count_dict_name get_count_with_names(count_dict_idx, model) # 使用渲染后的图像进行可视化 rendered_imgs results.render() if rendered_imgs: orig_img_with_boxes rendered_imgs[0] else: orig_img_with_boxes results.imgs[0] final_img visualize_detections_with_count(orig_img_with_boxes, filtered_dets, count_dict_name, model, conf_thresholdconf_thres) # 生成输出文件名 img_name Path(img_path).stem output_path Path(output_dir) / f{img_name}_result.jpg cv2.imwrite(str(output_path), final_img) print(f结果已保存至: {output_path}) print(f检测计数: {count_dict_name}) return final_img, count_dict_name def process_video(model, video_path, output_dir./output, iou_thres0.5, conf_thres0.25, show_videoFalse): 处理视频文件。 # ... (视频处理逻辑结合SmoothCounter类) pass def main(): parser argparse.ArgumentParser(descriptionYOLOv5目标分类计数与显示) parser.add_argument(--source, typestr, default./data/images, help输入源可以是图片路径、图片文件夹、视频文件或0摄像头) parser.add_argument(--weights, typestr, defaultyolov5s.pt, help模型权重路径如 yolov5s.pt) parser.add_argument(--output, typestr, default./output, help结果输出目录) parser.add_argument(--iou-thres, typefloat, default0.5, help计数去重的IOU阈值) parser.add_argument(--conf-thres, typefloat, default0.25, help检测置信度阈值) args parser.parse_args() # 加载模型 print(f加载模型: {args.weights}) model torch.hub.load(ultralytics/yolov5, custom, pathargs.weights, force_reloadFalse) model.conf args.conf_thres # 设置模型置信度阈值 model.iou 0.45 # 设置模型NMS IOU阈值 source_path Path(args.source) if source_path.is_file(): if source_path.suffix.lower() in [.jpg, .jpeg, .png, .bmp]: process_image(model, str(source_path), args.output, args.iou_thres, args.conf_thres) elif source_path.suffix.lower() in [.mp4, .avi, .mov]: process_video(model, str(source_path), args.output, args.iou_thres, args.conf_thres, show_videoTrue) elif source_path.is_dir(): for img_file in source_path.glob(*.[jp][pn]g): process_image(model, str(img_file), args.output, args.iou_thres, args.conf_thres) elif args.source 0: process_video(model, 0, args.output, args.iou_thres, args.conf_thres, show_videoTrue) else: print(错误不支持的输入源。) if __name__ __main__: main()使用指南确保已安装好环境见第2部分。将上述完整脚本保存。准备一张测试图片例如test.jpg。在终端运行python yolov5_count_display.py --source ./test.jpg --weights yolov5s.pt --output ./results查看./results文件夹下的结果图片和控制台输出的计数信息。通过这个完整的流程你应该能够将YOLOv5变成一个强大的目标检测与计数工具。记住核心在于理解数据流从原始输出到结构化计数和可视化设计清晰传达信息。在实际项目中你可能还需要根据具体场景调整去重逻辑、优化显示样式甚至集成到更大的系统流水线中。希望这些从实际项目中总结出的细节和代码能帮你少走弯路。
返回列表