基于YOLOv5的目标检测与分类计数实战:从模型部署到可视化优化
1. 项目概述:从检测到计数的实用化跨越
在计算机视觉的落地应用里,目标检测早已不是新鲜事。YOLOv5以其出色的速度和精度平衡,成为了许多开发者和研究者的首选工具。但很多时候,仅仅在图像上画出一个个框(Bounding Box)是远远不够的。业务方真正关心的是:“这一帧里,到底有多少个A类物品,多少个B类物品?” 这就是从“检测”到“分类计数”的跨越。我最近完成的一个项目,核心需求正是基于YOLOv5的检测结果,实时统计不同类别的目标数量,并将这些统计信息清晰、美观地叠加显示在原始图像上。这听起来简单,但在实际实现中,如何避免重复计数、如何设计高效的计数逻辑、如何让显示信息既直观又不遮挡关键区域,每一步都藏着细节。
这个项目的价值在于,它将YOLOv5的算法能力转化为了可直接用于生产监控、库存盘点、流量统计等场景的实用功能。比如,在仓库管理中,你可以用摄像头拍下货架,程序不仅能识别出不同种类的商品,还能直接告诉你每种商品的数量,极大提升了效率。为了实现它,我们需要深入YOLOv5的输出结构,设计一个稳健的计数系统,并利用OpenCV等工具完成最后的视觉化呈现。整个过程涉及模型推理、后处理逻辑、计数算法和图像渲染等多个环节的衔接。
2. 核心思路与方案设计
2.1 为什么选择YOLOv5作为基础
在众多目标检测模型中选定YOLOv5,是经过一番考量的。首先,它的生态极其成熟。从官方的PyTorch实现到各种尺寸的预训练模型(n, s, m, l, x),再到详尽的文档和活跃的社区,这大大降低了我们的入门和调试门槛。其次,YOLOv5的推理速度非常快,即使在资源受限的边缘设备(如你提到的RK3568)上,经过适当优化也能达到实时或准实时的性能,这对于需要连续处理视频流进行计数的场景至关重要。最后,它的输出格式相对规整,为我们后续提取检测框和类别信息提供了便利。
当然,YOLOv5并非没有挑战。其中一个常见问题就是“重复计数”,即同一个物体被多个重叠的框检测到。这在后续计数环节是致命的,会直接导致数量统计翻倍。因此,我们的方案设计必须将“去重”作为核心环节之一。
2.2 整体流程架构设计
整个项目的流程可以清晰地划分为四个阶段,它们像流水线一样环环相扣:
- 图像输入与预处理 :读取图像或视频帧,并按照YOLOv5模型的要求进行缩放、归一化等操作。这一步通常由模型本身的预处理函数完成。
- 模型推理与原始输出获取 :将预处理后的图像送入加载好的YOLOv5模型,得到原始的检测输出。这个输出包含了大量候选框的信息。
- 后处理与计数逻辑 :这是项目的核心。我们需要对原始输出进行非极大值抑制(NMS)来消除重叠框,然后根据置信度阈值过滤掉不可靠的检测。最后,设计计数逻辑,遍历所有保留下来的检测框,按类别进行累加。这里的计数逻辑需要仔细设计,确保一物一计。
- 结果可视化与输出 :将原始图像、绘制了检测框的图像以及统计得到的分类计数结果,通过OpenCV等库合成到一张最终图像上,并显示或保存。
这个架构的关键在于第三阶段——后处理与计数。它决定了计数的准确性。
2.3 关键技术选型:NMS与计数逻辑
非极大值抑制(NMS)
:这是解决“重复计数”的第一道,也是最重要的一道防线。YOLOv5本身在训练和推理时都会使用NMS。其原理很简单:对于同一类别的所有检测框,根据置信度排序,选中置信度最高的框,然后计算它与其余所有框的交并比(IoU)。如果IoU超过设定的阈值(通常为0.45),就认为它们检测的是同一个物体,将那些框抑制(删除)。这个过程迭代进行,直到处理完所有框。在
detect.py
或自己编写代码时,可以通过
torchvision.ops.nms
或YOLOv5内置的
non_max_suppression
函数方便地调用。
计数逻辑的设计
:经过NMS后,理论上每个物体只剩下一个最可靠的框。此时,计数就变得直观:我们只需要遍历所有检测结果,创建一个字典(例如Python的
defaultdict(int)
或
Counter
),以类别ID或类别名称为键,每遇到一个该类别的检测框,就将对应的值加1。但是,在实际动态视频流中,还需要考虑帧间关联以避免闪烁或瞬时漏检导致的计数波动,不过这属于更高级的追踪范畴,本项目聚焦于单帧图像的静态计数。
3. 环境搭建与YOLOv5模型部署
3.1 基础环境配置步骤
工欲善其事,必先利其器。一个稳定、兼容的环境是项目成功的基石。我强烈建议使用Anaconda来管理Python环境,它能有效解决包依赖冲突的问题。
# 1. 创建并激活一个新的conda环境(以Python 3.8为例,这是与PyTorch和YOLOv5兼容性较好的版本)
conda create -n yolov5_counting python=3.8
conda activate yolov5_counting
# 2. 安装PyTorch。请务必根据你的CUDA版本(如果有GPU)去PyTorch官网获取正确的安装命令。
# 例如,对于CUDA 11.3:
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
# 如果没有GPU,则安装CPU版本:
# pip install torch torchvision torchaudio
# 3. 克隆YOLOv5官方仓库
git clone https://github.com/ultralytics/yolov5
cd yolov5
# 4. 安装YOLOv5的依赖包
pip install -r requirements.txt
注意 :
requirements.txt里包含的opencv-python、matplotlib、pandas等是可视化所必需的。如果安装缓慢,可以考虑使用国内镜像源,例如pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple。
3.2 模型下载与验证
YOLOv5提供了多种预训练模型,从轻量级的YOLOv5n到高精度的YOLOv5x。对于计数任务,如果目标物体不是特别小或密集,YOLOv5s通常是一个在精度和速度上平衡的不错选择。
import torch
# 加载预训练模型(这里以YOLOv5s为例)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
# 或者使用本地克隆的仓库
# model = torch.hub.load('./yolov5', 'custom', path='yolov5s.pt', source='local')
# 将模型设置为评估模式(这对推理很重要)
model.eval()
# 验证模型是否加载成功:用一张示例图片测试
img = 'https://ultralytics.com/images/zidane.jpg'
results = model(img)
results.print() # 打印检测结果
results.show() # 显示带检测框的图片
如果运行成功,你会看到终端打印出检测到的物体类别、坐标和置信度,同时弹出一张画好框的图片。这说明你的YOLOv5基础环境已经跑通了。
3.3 关于训练自己数据集的补充
热搜词里提到了“yolov5训练自己的数据集”。如果你的计数对象不在COCO数据集(YOLOv5默认预训练模型所针对的80类)中,这一步就是必须的。流程大致如下:
-
数据准备
:收集图片,使用LabelImg等工具标注,生成YOLO格式的txt文件(每个txt文件对应一张图,内容为
class_id x_center y_center width height,均归一化到0-1)。 -
配置文件
:在
yolov5/data/目录下创建自己数据集的yaml文件,定义路径、类别数和类别名。 -
选择与修改模型配置
:在
yolov5/models/目录下选择模型配置文件(如yolov5s.yaml),将其中的nc(类别数)修改为你数据集的类别数。 -
开始训练
:运行
python train.py --img 640 --batch 16 --epochs 100 --data your_data.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt。训练好的模型会保存在runs/train/exp/weights/best.pt。
实操心得 :训练自己的模型时,数据集的质量和数量是关键。确保标注准确,且各类别样本数量相对均衡,否则模型会偏向于数量多的类别。如果遇到“yolov5训练map总是0”的问题,首先检查数据标注格式是否正确,路径是否配置准确,然后可以尝试减小学习率(
--lr),或者检查数据集是否真的包含可检测的目标。
4. 核心代码实现:分类计数与图像叠加
4.1 解析模型输出与后处理
当我们调用
model(img)
进行推理后,返回的
results
对象包含了丰富的信息。我们需要从中提取用于计数的结构化数据。
import cv2
from collections import Counter
import numpy as np
def count_and_display(image_path, model, conf_thres=0.25, iou_thres=0.45):
"""
对单张图像进行目标检测、分类计数并可视化结果。
参数:
image_path: 输入图像路径
model: 加载好的YOLOv5模型
conf_thres: 置信度阈值,低于此值的检测将被忽略
iou_thres: NMS所用的IoU阈值
"""
# 1. 读取图像
img_original = cv2.imread(image_path)
if img_original is None:
print(f"错误:无法读取图像 {image_path}")
return
img_rgb = cv2.cvtColor(img_original, cv2.COLOR_BGR2RGB) # YOLOv5期望RGB输入
# 2. 模型推理
results = model(img_rgb, size=640) # 可以调整size,但需与训练时一致
# 3. 解析结果
# results.pandas().xyxy[0] 返回一个Pandas DataFrame,包含检测框信息
detections_df = results.pandas().xyxy[0]
# 列包括: xmin, ymin, xmax, ymax, confidence, class, name
# 4. 应用置信度阈值过滤
filtered_detections = detections_df[detections_df['confidence'] > conf_thres]
# 5. 分类计数
# 使用Counter统计类别名称的出现次数
class_counter = Counter(filtered_detections['name'])
# 打印计数结果
print("分类计数结果:")
for class_name, count in class_counter.items():
print(f" {class_name}: {count}")
# 6. 可视化
# 将results的渲染结果(带框的RGB图像)转换回BGR供OpenCV显示
img_with_boxes = results.render()[0] # render()返回一个列表
img_with_boxes_bgr = cv2.cvtColor(img_with_boxes, cv2.COLOR_RGB2BGR)
# 7. 在图像上绘制计数信息
display_img = img_with_boxes_bgr.copy()
# 设置文本参数
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.7
thickness = 2
# 选择一块区域显示计数,例如左上角
y_offset = 30
for i, (class_name, count) in enumerate(class_counter.items()):
text = f"{class_name}: {count}"
# 获取文本大小,用于背景框
(text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, thickness)
# 绘制半透明背景矩形,使文字更清晰
cv2.rectangle(display_img,
(10, y_offset + i*30 - text_height - 5),
(10 + text_width, y_offset + i*30 + 5),
(0, 0, 0), -1) # 黑色背景
# 绘制白色文字
cv2.putText(display_img, text, (10, y_offset + i*30),
font, font_scale, (255, 255, 255), thickness)
# 8. 显示图像
cv2.imshow('YOLOv5 Detection & Counting', display_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
# 可选:保存结果图像
# cv2.imwrite('result.jpg', display_img)
return class_counter, display_img
# 使用示例
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
model.eval()
count_result, result_image = count_and_display('your_image.jpg', model)
这段代码完成了从推理到显示的全过程。关键点在于
results.pandas().xyxy[0]
,它直接将检测结果转换成了易于操作的DataFrame。
Counter
工具让分类计数变得异常简洁。
4.2 计数逻辑的深度剖析与优化
上面的计数逻辑在静态图片上工作良好,但它基于一个假设:NMS已经完美地去除了所有重复框。然而在复杂场景下,特别是目标密集、遮挡严重时,NMS可能无法完全消除所有重复,或者可能错误地抑制了并排的两个相似物体。
更健壮的计数逻辑可以考虑以下几点:
- 基于跟踪的帧间去重 :对于视频流,单纯统计单帧数量会导致数字剧烈跳动。更优的做法是引入目标跟踪(如ByteTrack, DeepSORT),为每个检测到的物体分配一个唯一ID。计数时,我们统计的是出现在画面中的独立ID数量,而不是每一帧的检测框数量。这能有效平滑计数结果,并解决短时遮挡导致的“消失-重现”计数问题。
-
区域计数
:有时我们只关心特定区域内的物体数量,比如传送带上的零件、十字路口特定车道的车辆。这需要在计数前增加一个区域判断。我们可以定义一个多边形区域(ROI),只统计中心点落在该区域内的检测框。
import cv2 def is_point_in_polygon(point, polygon): """判断点是否在多边形内(使用射线法)""" # polygon: [(x1,y1), (x2,y2), ...] x, y = point inside = False n = len(polygon) p1x, p1y = polygon[0] for i in range(1, n + 1): p2x, p2y = polygon[i % n] if y > min(p1y, p2y): if y <= max(p1y, p2y): if x <= max(p1x, p2x): if p1y != p2y: xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x if p1x == p2x or x <= xinters: inside = not inside p1x, p1y = p2x, p2y return inside # 在计数循环中 roi_polygon = [(100,100), (500,100), (500,400), (100,400)] # 矩形区域 for index, row in filtered_detections.iterrows(): # 计算检测框中心点 x_center = (row['xmin'] + row['xmax']) / 2 y_center = (row['ymin'] + row['ymax']) / 2 if is_point_in_polygon((x_center, y_center), roi_polygon): class_counter[row['name']] += 1 - 自定义去重规则 :如果某些特定类别的物体即使经过NMS后仍然容易产生重复框,可以针对这些类别设置更严格的去重规则。例如,对于“瓶子”这类形状规则、可能密集排列的物体,可以在NMS后,再对同一类别的框进行一次自定义的IoU或中心点距离检查。
4.3 可视化效果的增强
将计数信息显示在图像上,美观和清晰同样重要。前面的代码使用了简单的文本叠加,这里可以进一步优化:
- 动态位置 :将计数信息放在图像顶部或底部,并确保不遮挡重要检测目标。可以计算图像中检测框的分布,自动选择空白区域放置统计面板。
-
样式美化
:使用不同的背景色、边框和字体来增强可读性。OpenCV的
cv2.rectangle可以绘制圆角矩形(通过绘制多个短线段模拟),cv2.putText可以调整字体样式。 -
进度条或图表
:对于数量监控场景,可以用进度条的形式显示当前数量与设定阈值的对比,更加直观。
# 示例:绘制一个简单的水平进度条表示某类物体的数量 max_count = 20 current_count = class_counter.get('person', 0) bar_width = 200 bar_height = 20 bar_x, bar_y = 10, 50 # 绘制背景条 cv2.rectangle(display_img, (bar_x, bar_y), (bar_x + bar_width, bar_y + bar_height), (100, 100, 100), -1) # 绘制前景条(根据比例) fill_width = int((current_count / max_count) * bar_width) if max_count > 0 else 0 fill_width = min(fill_width, bar_width) cv2.rectangle(display_img, (bar_x, bar_y), (bar_x + fill_width, bar_y + bar_height), (0, 255, 0), -1) # 绘制文本 count_text = f"Person: {current_count}/{max_count}" cv2.putText(display_img, count_text, (bar_x, bar_y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
5. 高级应用与性能调优
5.1 处理视频流与实时计数
将单张图片的处理逻辑扩展到视频流或摄像头输入,是实现实时计数的关键。核心是使用循环读取每一帧,并复用检测与计数函数。
import cv2
import time
def real_time_counting(camera_index=0, model=None):
"""
实时摄像头视频流目标计数。
"""
cap = cv2.VideoCapture(camera_index)
if not cap.isOpened():
print("无法打开摄像头")
return
prev_time = 0
fps = 0
while True:
ret, frame = cap.read()
if not ret:
print("无法获取帧")
break
# 计算FPS
current_time = time.time()
fps = 1 / (current_time - prev_time) if prev_time > 0 else 0
prev_time = current_time
# 将BGR帧转换为RGB
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# 推理
results = model(frame_rgb, size=640) # 可根据性能调整size,如320更快
# 解析与计数
detections_df = results.pandas().xyxy[0]
conf_thres = 0.25
filtered_detections = detections_df[detections_df['confidence'] > conf_thres]
class_counter = Counter(filtered_detections['name'])
# 渲染检测框
rendered_frame = results.render()[0]
display_frame = cv2.cvtColor(rendered_frame, cv2.COLOR_RGB2BGR)
# 在帧上叠加计数信息和FPS
y_offset = 30
# 显示FPS
cv2.putText(display_frame, f"FPS: {int(fps)}", (10, y_offset),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
# 显示计数
for i, (cls_name, count) in enumerate(class_counter.items()):
text = f"{cls_name}: {count}"
cv2.putText(display_frame, text, (10, y_offset + 30 + i*25),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
cv2.imshow('Real-time YOLOv5 Counting', display_frame)
# 按'q'退出
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
注意事项 :实时处理时,性能是关键。如果帧率过低,可以尝试:1) 使用更小的YOLOv5模型(如
yolov5n);2) 减小推理图像尺寸(size参数,如从640降到320);3) 跳过部分帧进行处理(如每2帧处理1帧);4) 使用GPU进行推理加速。
5.2 模型优化与加速技巧
当部署在资源受限的边缘设备(如RK3568、RV1106)时,模型优化必不可少。
-
模型量化
:将模型权重从浮点数(FP32)转换为整数(INT8),可以显著减少模型大小并提升推理速度,对精度影响通常可控。PyTorch提供了
torch.quantization工具。# 这是一个简化的量化示例,实际过程更复杂,需要准备校准数据 model_fp32 = ... # 你的FP32模型 model_fp32.eval() model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm') # x86后端 model_int8 = torch.quantization.prepare(model_fp32, inplace=False) # ... 用校准数据运行模型 ... model_int8 = torch.quantization.convert(model_int8) # 保存和加载量化模型 torch.save(model_int8.state_dict(), 'quantized_model.pth') -
TorchScript或ONNX导出
:将PyTorch模型转换为TorchScript或ONNX格式,可以利用PyTorch或ONNX Runtime的优化推理引擎,有时能获得更好的性能。YOLOv5官方提供了
export.py脚本支持导出多种格式。python export.py --weights yolov5s.pt --include torchscript onnx - 特定硬件SDK :对于RK3568、RV1106这类芯片,厂商通常会提供完整的AI工具链(如RKNN Toolkit),可以将ONNX或PyTorch模型转换并优化为能在其NPU上高效运行的专有格式。这一步通常能带来最大的性能提升。
5.3 计数结果的持久化与输出
对于生产环境,我们不仅需要显示,还需要记录。可以将每帧的计数结果保存到文件或数据库中。
import csv
from datetime import datetime
def log_counts_to_csv(count_dict, timestamp, filename='count_log.csv'):
"""
将计数结果和时间戳记录到CSV文件。
"""
file_exists = os.path.isfile(filename)
with open(filename, 'a', newline='') as csvfile:
fieldnames = ['timestamp'] + list(count_dict.keys())
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
if not file_exists:
writer.writeheader()
row_data = {'timestamp': timestamp}
row_data.update(count_dict)
writer.writerow(row_data)
# 在处理每一帧或每隔一段时间后调用
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
log_counts_to_csv(dict(class_counter), current_time)
你也可以将结果输出到网络接口,供其他系统调用,或者触发某些动作(如数量超过阈值时报警)。
6. 常见问题排查与实战心得
6.1 高频问题速查表
| 问题现象 | 可能原因 | 排查步骤与解决方案 |
|---|---|---|
| 检测不到任何目标 |
1. 置信度阈值(
conf_thres
)设置过高。
2. 图像预处理尺寸与训练时不匹配。 3. 模型未正确加载或处于训练模式。 4. 目标物体不在模型识别范围内(未训练此类)。 |
1. 逐步调低
conf_thres
(如从0.5调到0.1)观察。
2. 确保推理时
size
参数与模型训练时一致(默认640)。
3. 检查模型加载代码,确认调用
model.eval()
。
4. 使用模型支持的类别列表(如COCO的80类)进行验证,或训练自己的模型。 |
| 计数结果远多于实际数量(重复计数) |
1. NMS的IoU阈值(
iou_thres
)设置过高,抑制不够。
2. 置信度阈值过低,保留了太多噪声检测。 3. 目标密集,NMS算法本身局限。 |
1. 适当降低
iou_thres
(如从0.45调到0.3)。
2. 适当提高
conf_thres
,过滤低质量检测框。
3. 考虑使用更先进的NMS变体,如Soft-NMS,或引入自定义后处理(如小目标合并)。 |
| 计数结果少于实际数量(漏检) |
1. 置信度阈值(
conf_thres
)设置过高。
2. 目标尺寸太小或太模糊。 3. 模型在该类别上精度不足。 |
1. 适当降低
conf_thres
。
2. 尝试增大输入图像尺寸(
size
),或使用专门针对小目标优化的模型。
3. 收集更多该类别数据对模型进行微调。 |
| 推理速度非常慢 |
1. 使用了过大的模型(如YOLOv5x)。
2. 在CPU上运行。 3. 输入图像尺寸过大。 |
1. 换用更小的模型(如YOLOv5n或YOLOv5s)。
2. 检查CUDA是否可用,确保模型
.to(device)
到GPU。
3. 减小推理时的
size
参数(如从640降到320),会损失一定精度。
|
| 可视化时文字不清晰或位置不当 |
1. 字体大小(
font_scale
)或粗细(
thickness
)不合适。
2. 文字颜色与背景色对比度低。 3. 文字位置遮挡了关键目标。 |
1. 根据图像分辨率调整
font_scale
和
thickness
。
2. 为文字添加背景色块(如半透明黑色矩形)。 3. 动态计算图像中检测框的分布,将统计面板放在空白区域。 |
6.2 调试技巧:可视化中间结果
当计数不准时,不要只盯着最终数字。将中间过程可视化,能帮你快速定位问题所在。
# 在NMS前后,绘制检测框进行对比
def visualize_detections(image, detections_df, title="Detections"):
img_viz = image.copy()
for idx, row in detections_df.iterrows():
x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
label = f"{row['name']} {row['confidence']:.2f}"
cv2.rectangle(img_viz, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(img_viz, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
cv2.imshow(title, img_viz)
cv2.waitKey(0)
# 在计数函数中,NMS前和NMS后分别调用此函数
# visualize_detections(img_original, detections_df, "Before NMS")
# visualize_detections(img_original, filtered_detections, "After NMS & Confidence Filter")
通过对比NMS前后的框,你可以清楚地看到是否有大量重叠框未被抑制,从而判断是NMS参数问题还是模型本身输出质量的问题。
6.3 关于超参数调优的体会
热搜词里提到了“yolov5超参数”。在计数任务中,直接影响结果的超参数主要有两个:
-
conf_thres(置信度阈值) :这是过滤低质量检测的第一道关卡。设得太高,会漏掉一些模糊但真实的目标;设得太低,会引入大量噪声和误检,干扰计数。我的经验是从默认值0.25开始,在验证集上观察精确率(Precision)和召回率(Recall)的平衡,或者直接观察计数结果的准确性来调整。 -
iou_thres(NMS的IoU阈值) :这个参数控制去重的严格程度。对于通常分得很开的物体(如行人),0.45的默认值很好。但对于堆叠在一起的同类物体(如仓库里的箱子),可能需要调低(如0.3),以防止一个框抑制掉旁边紧挨着的另一个真实物体。反之,如果同一物体总是产生多个紧贴的框,可能需要调高IoU阈值来更激进地合并。
调参没有银弹,最好的方法就是用一批有标注的测试图像,写一个脚本批量运行不同参数组合,然后统计计数的准确率,选择表现最好的那一组。这个过程虽然繁琐,但对于提升最终系统的可靠性至关重要。
更多推荐
所有评论(0)