YOLO 目标检测后处理:NMS 算法详解
在工业质检线上,一台 AI 相机正高速扫描经过的电路板。模型瞬间识别出数十个'焊点缺陷'候选框——可明明只有一个异常区域,为何系统报出了七八次?这种'一物多检'的混乱不仅让操作员困惑,更可能触发误剔除机制,造成产线停机。这正是目标检测落地时最典型的痛点之一。
问题的根源不在模型本身,而在于输出端的处理逻辑。YOLO 这类单阶段检测器为了保证高召回率,天生会生成大量重叠预测框。如果没有一套高效的'仲裁机制',再精准的模型也会输出一团乱码。这时候,非极大值抑制(Non-Maximum Suppression, NMS)就扮演了关键角色:它像一位冷静的裁判,在一堆指向同一目标的候选框中,只留下最有说服力的那个。
从冗余到清晰:NMS 如何重塑检测输出
想象一下,YOLO 模型在推理时就像一群密集的探头同时工作。每个网格单元都试图捕捉视野中的物体,相邻区域对同一个目标产生响应几乎是必然现象。比如一只猫横跨多个网格,不同锚框从略微偏移的角度框出它的轮廓,最终输出可能是十几个高度重叠的边界框,置信度从 0.98 到 0.6 不等。
如果不加干预,这样的输出根本无法用于实际决策。下游系统无法判断这是'一只猫被多次确认',还是'发现了七只紧密排列的猫'。NMS 要解决的就是这个语义模糊性问题。
其核心策略非常直观:先按置信度给所有框排序,然后逐个审查。取当前最高分的框 A 作为基准,计算其余每个框与 A 的交并比(IoU)。如果某个框 B 与 A 的重叠面积超过设定阈值(例如 0.5),就认为它们在'竞争'同一个目标,既然 B 得分更低,理应被淘汰。这个过程不断迭代,直到所有框都被评估完毕。
这里有个工程细节常被忽视:NMS 必须按类别独立执行。假设画面中有一人一车靠得很近,他们的检测框 IoU 可能高达 0.6。若不做类别区分,人框和车框会互相残杀,最终只剩一个幸存者。正确的做法是分别对'人'类框和'车'类框做两套 NMS,确保跨类别抑制不会发生。
import numpy as np
def nms(boxes, scores, iou_threshold=0.5):
"""
执行非极大值抑制 (NMS)
参数:
boxes: numpy 数组,形状为 (N, 4),每行表示 [x1, y1, x2, y2]
scores: numpy 数组,形状为 (N,),表示每个框的置信度得分
iou_threshold: float,IoU 阈值,超过此值则视为重复
返回:
keep_indices: 保留下来的框的索引列表
"""
sorted_indices = np.argsort(scores)[::-1]
keep_boxes = []
while len(sorted_indices) > 0:
current_idx = sorted_indices[0]
keep_boxes.append(current_idx)
if len(sorted_indices) == 1:
break
others = sorted_indices[1:]
ious = compute_iou(boxes[current_idx], boxes[others])
sorted_indices = others[ious <= iou_threshold]
return keep_boxes
def compute_iou(box1, boxes):
x1, y1, x2, y2 = box1
x1s, y1s, x2s, y2s = boxes[:, 0], boxes[:, 1], boxes[:, ], boxes[:, ]
inter_x1 = np.maximum(x1, x1s)
inter_y1 = np.maximum(y1, y1s)
inter_x2 = np.minimum(x2, x2s)
inter_y2 = np.minimum(y2, y2s)
inter_w = np.maximum(, inter_x2 - inter_x1)
inter_h = np.maximum(, inter_y2 - inter_y1)
inter_area = inter_w * inter_h
area1 = (x2 - x1) * (y2 - y1)
areas = (x2s - x1s) * (y2s - y1s)
union_area = area1 + areas - inter_area
iou = inter_area / np.maximum(union_area, )
iou

