当前位置： 首页 > news >正文

保姆级教程：用Python手把手实现YOLOv5中的NMS（附代码与可视化）

news 2026/6/10 21:04:41

保姆级教程：用Python手把手实现YOLOv5中的NMS（附代码与可视化）

在目标检测的实际应用中，模型往往会输出大量冗余的预测框。以YOLOv5为例，输入一张640×640的图像，模型会生成25200个初始预测框。这些框密集地覆盖在可能的目标位置，但最终我们只需要保留最准确的几个。这就是非极大抑制(Non-Maximum Suppression, NMS)技术的核心作用——像一位精明的编辑，从海量初稿中筛选出最优质的几篇。

本文将带你用Python从零实现NMS的完整流程，不仅会解释每个步骤的数学原理，还会通过Matplotlib动态可视化筛选过程。无论你是刚接触目标检测的学生，还是需要调试模型参数的工程师，都能通过这份"代码+可视化"的实操指南深入理解NMS的工作机制。我们会重点解决三个实际问题：如何计算框的重叠度(IoU)、如何处理多类别NMS、以及如何避免常见实现错误。

1. 环境准备与数据模拟

1.1 安装必要库

确保已安装以下Python库，这些将用于数值计算和可视化：

pip install numpy matplotlib opencv-python

1.2 模拟预测框数据

实际项目中，预测框来自模型输出。为方便演示，我们模拟20个随机框及其置信度：

import numpy as np np.random.seed(42) boxes = np.random.uniform(100, 500, size=(20, 4)) # 生成(x1,y1,x2,y2)格式的框 boxes[:, 2:] += boxes[:, :2] # 转换宽高为右下角坐标 scores = np.random.uniform(0.5, 0.95, size=20) # 置信度分数 class_ids = np.random.randint(0, 3, size=20) # 3个类别 print("示例框坐标:\n", boxes[:3]) print("对应置信度:", scores[:3]) print("类别ID:", class_ids[:3])

提示：实际YOLOv5输出需先转换坐标格式。原始输出为(cx,cy,w,h)需转为(x1,y1,x2,y2)

1.3 可视化初始框

使用Matplotlib绘制所有预测框：

import matplotlib.pyplot as plt import matplotlib.patches as patches def plot_boxes(boxes, scores, title): fig, ax = plt.subplots(figsize=(10,6)) ax.set_xlim(0,600); ax.set_ylim(0,600) for i, (x1,y1,x2,y2) in enumerate(boxes): ax.add_patch(patches.Rectangle( (x1,y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none', alpha=0.3)) ax.text(x1, y1, f"{scores[i]:.2f}", color='blue') ax.set_title(title) plt.gca().invert_yaxis() plt.show() plot_boxes(boxes, scores, "初始预测框 (含置信度)")

2. NMS核心算法实现

2.1 IoU计算：衡量框的重叠度

交并比(Intersection over Union)是NMS的核心指标，计算两个矩形A和B的重叠程度：

def calculate_iou(box1, box2): # 确定相交区域的坐标 x1_inter = max(box1[0], box2[0]) y1_inter = max(box1[1], box2[1]) x2_inter = min(box1[2], box2[2]) y2_inter = min(box1[3], box2[3]) # 计算相交区域面积 inter_area = max(0, x2_inter - x1_inter) * max(0, y2_inter - y1_inter) # 计算各自面积 area1 = (box1[2]-box1[0])*(box1[3]-box1[1]) area2 = (box2[2]-box2[0])*(box2[3]-box2[1]) # 避免除以零 iou = inter_area / (area1 + area2 - inter_area + 1e-6) return iou

测试两个示例框：

box_a = [100, 100, 200, 200] box_b = [150, 150, 250, 250] print(f"IoU值: {calculate_iou(box_a, box_b):.2f}")

2.2 基础NMS实现

按置信度排序后，迭代抑制重叠框：

def nms(boxes, scores, iou_threshold=0.5): # 按置信度降序排列 order = np.argsort(scores)[::-1] keep = [] while order.size > 0: # 选取当前最高分框 i = order[0] keep.append(i) # 计算与剩余框的IoU ious = np.array([calculate_iou(boxes[i], boxes[j]) for j in order[1:]]) # 保留IoU低于阈值的索引 inds = np.where(ious <= iou_threshold)[0] order = order[inds + 1] # +1补偿order[1:] return keep

可视化NMS过程：

def visualize_nms_steps(boxes, scores, iou_thresh=0.5): order = np.argsort(scores)[::-1] fig, ax = plt.subplots(figsize=(10,6)) ax.set_xlim(0,600); ax.set_ylim(0,600) for step in range(len(order)): if step == 0: color = 'g' # 当前最高分框 elif step in keep: color = 'b' # 保留框 else: color = 'r' # 被抑制框 idx = order[step] x1,y1,x2,y2 = boxes[idx] ax.add_patch(patches.Rectangle( (x1,y1), x2-x1, y2-y1, linewidth=1, edgecolor=color, facecolor='none', alpha=0.7)) ax.text(x1, y1, f"{scores[idx]:.2f}", color='black') ax.set_title(f"NMS结果 (IoU阈值={iou_thresh})") plt.gca().invert_yaxis() plt.show() keep = nms(boxes, scores) visualize_nms_steps(boxes, scores)

3. 多类别NMS处理

3.1 按类别分组处理

实际目标检测需要独立处理每个类别的预测框：

def multiclass_nms(boxes, scores, class_ids, iou_threshold=0.5): unique_classes = np.unique(class_ids) final_keep = [] for cls in unique_classes: # 获取当前类别的框和分数 cls_mask = (class_ids == cls) cls_boxes = boxes[cls_mask] cls_scores = scores[cls_mask] # 对该类别执行NMS keep = nms(cls_boxes, cls_scores, iou_threshold) # 记录原始索引 original_indices = np.where(cls_mask)[0] final_keep.extend(original_indices[keep]) return final_keep

3.2 多类别可视化

用不同颜色区分不同类别的保留框：

keep_multi = multiclass_nms(boxes, scores, class_ids) fig, ax = plt.subplots(figsize=(10,6)) ax.set_xlim(0,600); ax.set_ylim(0,600) colors = ['r', 'g', 'b'] # 不同类别颜色 for i in keep_multi: x1,y1,x2,y2 = boxes[i] ax.add_patch(patches.Rectangle( (x1,y1), x2-x1, y2-y1, linewidth=2, edgecolor=colors[class_ids[i]], facecolor='none')) ax.text(x1, y1, f"cls:{class_ids[i]} score:{scores[i]:.2f}", color=colors[class_ids[i]]) ax.set_title("多类别NMS结果") plt.gca().invert_yaxis() plt.show()

4. 高级优化与调试技巧

4.1 Soft-NMS实现

传统NMS的硬阈值可能丢失邻近目标，Soft-NMS通过分数衰减更柔和地处理重叠框：

def soft_nms(boxes, scores, iou_threshold=0.3, sigma=0.5, score_threshold=0.1): # 初始化保留列表 keep = [] # 复制分数避免修改原数据 new_scores = scores.copy() while True: # 获取当前最高分 max_score_index = np.argmax(new_scores) max_score = new_scores[max_score_index] if max_score < score_threshold: break keep.append(max_score_index) new_scores[max_score_index] = 0 # 抑制已选框 # 计算与所有其他框的IoU ious = np.array([calculate_iou(boxes[max_score_index], boxes[i]) for i in range(len(boxes))]) # 高斯权重衰减 decay_weights = np.exp(-(ious**2)/sigma) new_scores = new_scores * decay_weights return keep

4.2 性能优化技巧

处理YOLOv5的25200个框时，效率至关重要：

  1. 向量化IoU计算：替换循环为矩阵运算
def batched_iou(boxes1, boxes2): # 扩展维度用于广播计算 boxes1 = np.expand_dims(boxes1, 1) # [N,1,4] boxes2 = np.expand_dims(boxes2, 0) # [1,M,4] # 计算相交区域 inter_x1 = np.maximum(boxes1[...,0], boxes2[...,0]) inter_y1 = np.maximum(boxes1[...,1], boxes2[...,1]) inter_x2 = np.minimum(boxes1[...,2], boxes2[...,2]) inter_y2 = np.minimum(boxes1[...,3], boxes2[...,3]) inter_area = np.maximum(0, inter_x2 - inter_x1) * np.maximum(0, inter_y2 - inter_y1) # 计算各自面积 area1 = (boxes1[...,2]-boxes1[...,0])*(boxes1[...,3]-boxes1[...,1]) area2 = (boxes2[...,2]-boxes2[...,0])*(boxes2[...,3]-boxes2[...,1]) return inter_area / (area1 + area2 - inter_area + 1e-6)
  1. GPU加速：使用PyTorch或TensorFlow实现
import torch def nms_torch(boxes, scores, iou_threshold): # 转换数据为PyTorch张量 boxes = torch.tensor(boxes, device='cuda') scores = torch.tensor(scores, device='cuda') # 按分数降序排序 order = torch.argsort(scores, descending=True) keep = [] while order.numel() > 0: i = order[0] keep.append(i.item()) if order.numel() == 1: break # 计算IoU box_i = boxes[i].unsqueeze(0) other_boxes = boxes[order[1:]] ious = torchvision.ops.box_iou(box_i, other_boxes).squeeze() # 保留低重叠框 mask = ious <= iou_threshold order = order[1:][mask] return keep

4.3 常见问题排查

调试NMS时遇到的典型问题及解决方案：

  1. 坐标格式错误：

    • 症状：IoU计算异常导致大量误抑制
    • 检查：确认输入是(x1,y1,x2,y2)格式且x2>x1, y2>y1

  2. 阈值选择不当：

    • 现象：过高的阈值导致重复框，过低则丢失真实目标
    • 建议：从0.45开始调整，行人检测可能需要0.3-0.4

  3. 多类别处理异常：

    • 错误：不同类别的框相互抑制
    • 验证：确保按class_id分组处理

  4. 分数排序错误：

    • 影响：错误的框被优先保留
    • 调试：打印排序后的分数序列确认顺序正确

5. 完整YOLOv5 NMS集成

5.1 处理YOLOv5原始输出

YOLOv5输出为3个检测头的拼接结果，需先转换格式：

def process_yolov5_output(pred, conf_thres=0.25): # pred: [batch, num_boxes, 85] 格式 # 85 = x,y,w,h,obj_conf + 80类分数 # 过滤低置信度预测 obj_conf = pred[..., 4] class_scores = pred[..., 5:] scores = obj_conf.unsqueeze(-1) * class_scores # [n,80] max_scores, class_ids = torch.max(scores, dim=-1) keep = max_scores > conf_thres boxes = pred[..., :4][keep] scores = max_scores[keep] class_ids = class_ids[keep] # 转换cxcywh到xyxy x1 = boxes[..., 0] - boxes[..., 2]/2 y1 = boxes[..., 1] - boxes[..., 3]/2 x2 = boxes[..., 0] + boxes[..., 2]/2 y2 = boxes[..., 1] + boxes[..., 3]/2 boxes = torch.stack([x1,y1,x2,y2], dim=-1) return boxes, scores, class_ids

5.2 完整推理流程

整合预处理、NMS和后处理：

def yolov5_detection(model, img, iou_thres=0.45): # 模型推理 pred = model(img) # 处理原始输出 boxes, scores, class_ids = process_yolov5_output(pred[0]) # 多类别NMS keep = multiclass_nms(boxes, scores, class_ids, iou_thres) # 返回最终检测结果 return boxes[keep], scores[keep], class_ids[keep]

5.3 实际案例演示

加载预训练YOLOv5模型并运行完整流程：

import cv2 from PIL import Image # 加载示例图像 img = Image.open("zidane.jpg") img_np = np.array(img) # 运行检测 boxes, scores, class_ids = yolov5_detection(model, img_np) # 绘制结果 for box, score, cls_id in zip(boxes, scores, class_ids): x1,y1,x2,y2 = map(int, box) cv2.rectangle(img_np, (x1,y1), (x2,y2), (0,255,0), 2) cv2.putText(img_np, f"{class_names[cls_id]}:{score:.2f}", (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2) plt.imshow(img_np) plt.show()

在实现过程中发现，当处理密集目标时，适当降低IoU阈值至0.4能获得更好的召回率，但会增加一些重叠框。对于无人机拍摄的车辆检测项目，最终采用0.35的阈值配合Soft-NMS取得了最佳平衡。

http://www.zskr.cn/news/1327645.html

相关文章：

  • 别再只会拖控件了！FastReport 实战：手把手教你用代码搞定复杂报表（含分组、过滤、合计）
  • 绿化养护与绿植租摆公司推荐，华瑞环境服务优 - myqiye
  • QPSK、OQPSK、DQPSK傻傻分不清？一文讲透相位调制家族的区别与选型
  • 【城市生命线合集】300余份城市生命线、数字管廊、综合管线、智慧市政方案报告合集（PPT+WORD+PDF）
  • 从一块烧坏的板子说起：PCB电源平面设计的5个实战避坑指南（附嘉立创制程参数）
  • 高效硬件诊断工具：一站式解决AMD平台系统稳定性与性能优化问题
  • 告别等待！FF14国服玩家必备的副本动画跳过插件指南
  • DLSS Swapper完整指南：5分钟掌握游戏DLSS版本管理技巧
  • 2026年民间实体投资机构推荐 - 速递信息
  • 如何安全备份微信聊天记录？WeChatExporter开源工具全解析
  • Docker镜像导出(export/save)与导入(import/load)保姆级对比：别再混淆这4个命令了
  • 如何3步快速下载小红书无水印作品：终极小红书内容采集工具指南
  • Windows Defender深度移除技术架构解析：从内核服务到系统集成的完整解决方案
  • 大麦网抢票脚本：5分钟告别抢票焦虑的智能解决方案
  • Windows 11 LTSC 微软商店一键安装完整指南：3分钟解锁完整应用生态
  • Autovisor智慧树刷课终极指南：2025最全自动化学习解决方案
  • 用Arduino搞定M100和GM861S扫码枪，串口数据读取保姆级教程（附完整代码）
  • 图像生成模型‘杂交’指南：如何把VAE的稳定和GAN的清晰结合起来（VAEGAN实战）
  • 观察Taotoken在周末晚间高峰时段的API请求成功率
  • 移民公司推荐：如何选择可靠的移民服务机构 - 品牌排行榜
  • Figma中文汉化终极指南：3分钟告别英文界面困扰
  • 手把手教你用C语言给STM32写一个Modbus RTU从机库（含完整项目源码）
  • 海康摄像头码流设置详解：主码流、子码流怎么选？兼顾清晰度与流畅度的配置方案
  • 别再买交换机了！用Ubuntu 22.04和Netplan把旧电脑改造成软网桥（保姆级配置）
  • 脸部下垂怎么选合适的护肤品 CA逆时光 30天告别松弛纹 - 全网最美
  • 不是模型不行，是你没做好特征工程（附完整步骤）
  • 3分钟免费加速GitHub：告别龟速下载的终极解决方案
  • EasyExcel导出财务数据报表：手把手教你配置金额、百分比、小数位格式（含完整代码）
  • 思源宋体TTF实战秘籍：三步搞定专业中文字体配置
  • 手把手教你用L293D扩展板+Arduino Uno搭建一个简易CNC绘图仪（含28BYJ-48步进电机驱动教程）