深度学习目标检测:NMS(非极大值抑制)介绍与应用
引言:目标检测的核心挑战与NMS的诞生
在人工智能飞速发展的今天,计算机视觉作为其重要分支,正在深刻地改变着我们的生活。而目标检测(Object Detection)是计算机视觉领域中的一个核心任务,它的目标是在图像或视频中不仅识别出感兴趣的物体是什么(类别),还要确定它们在哪里(位置和大小),通常通过一个矩形框(Bounding Box)来表示。从自动驾驶、安防监控到工业质检和医疗影像分析,目标检测都有着极其广泛的应用。
早期的目标检测方法,如基于手工特征(HOG, SIFT等)和传统机器学习分类器(SVM等)的检测器,往往计算量大且性能受限。自2012年深度学习在ImageNet图像识别大赛中取得突破性进展以来,深度学习技术迅速被引入到目标检测领域,催生了R-CNN、Fast R-CNN、Faster R-CNN、YOLO、SSD等一系列划时代的检测算法。这些算法极大地提升了目标检测的精度和速度。
然而,无论是两阶段检测器(如Faster R-CNN,先生成候选区域再分类和回归)还是一阶段检测器(如YOLO、SSD,直接在预设位置上预测)或者Anchor-free检测器,它们的工作机制往往是生成大量的、可能是重叠的候选预测框。这是因为神经网络会在图像的多个位置、多个尺度上进行预测,对于同一个目标物体,可能会生成多个得分较高的候选框。这些候选框可能略有差异,有的更精确,有的略微偏离,但都指向同一个物体。
想象一下,如果我们不对这些冗余的预测框进行处理,最终的检测结果将会是同一个物体被许多个框包围,这显然不是我们想要的结果。我们希望每个真实物体只对应一个最准确、最具代表性的预测框。
为了解决这一问题,一个核心的后处理步骤被引入,这就是非极大值抑制(Non-Maximum Suppression,简称NMS)。NMS的目标是从大量的、重叠的候选框中,挑选出那些最有可能是真实目标的“最佳”预测框,同时抑制掉那些指向同一个目标的、得分较低的冗余框。它是几乎所有基于深度学习的目标检测器不可或缺的一部分。
本文将深入探讨NMS的原理、算法流程、关键参数以及它在深度学习目标检测中的具体应用。同时,我们也将讨论传统NMS的局限性以及为了克服这些局限性而提出的各种改进方法。
深度学习目标检测的输出与NMS的需求
在深入理解NMS之前,有必要简要回顾一下深度学习目标检测器的输出形式。大多数现代目标检测器,无论其内部架构如何复杂,最终都会输出一个预测列表,列表中包含:
- 边界框(Bounding Box): 表示物体在图像中的位置和大小,通常用矩形的左上角坐标 (x1, y1) 和右下角坐标 (x2, y2) 表示,或者用中心坐标 (cx, cy) 和宽度 (w) 和高度 (h) 表示。
- 类别预测(Class Prediction): 表示该边界框内包含的物体属于哪个类别(例如,人、车、猫等)。
- 置信度得分(Confidence Score): 表示该预测框包含一个物体的置信度,以及该物体属于预测类别的置信度。这个分数通常是一个介于0到1之间的值。
一个典型的检测器可能对一张图像生成数千甚至数万个这样的预测结果。例如,一个Anchor-based检测器会在图像的不同位置、不同长宽比和尺寸的Anchor上进行预测;即使是Anchor-free检测器,也可能在特征图的每个点上进行预测。
对于一个真实的物体,比如图像中的一只猫,模型可能会生成多个预测框:一个非常准确且得分很高,另一个稍微偏左得分也较高,还有一个稍微偏大得分中等,甚至还有一些低得分的框。如下图所示(脑补):
图像局部:一只猫
预测框1: [x1, y1, x2, y2], Class: Cat, Score: 0.98
预测框2: [x1', y1', x2', y2'], Class: Cat, Score: 0.95 (与预测框1高度重叠)
预测框3: [x1'', y1'', x2'', y2''], Class: Cat, Score: 0.85 (与预测框1、2重叠)
预测框4: [x''', y''', x''', y''''], Class: Cat, Score: 0.60 (与预测框1、2、3重叠)
...
如果我们直接输出所有得分高于某个阈值(例如0.5)的框,那么同一只猫就会被多个框圈住。这不仅在视觉上混乱,也影响了后续的任务,比如目标跟踪(tracking)或计数(counting)。NMS的使命就是从这些指向同一物体的预测框中,选出得分最高的那个,并剔除掉所有与它高度重叠的其他低得分框。
NMS(非极大值抑制)的原理与算法流程
非极大值抑制的核心思想非常直观:在一个局部区域内,只保留得分最高的那个边界框,同时抑制(移除)所有与其高度重叠且得分较低的边界框。这里的“高度重叠”是通过计算边界框之间的交并比(Intersection over Union,IoU)来衡量的。
1. 交并比(IoU):衡量边界框重叠程度
IoU是衡量两个边界框重叠程度的标准度量。对于任意两个边界框 A 和 B,它们的IoU定义为它们交集的面积除以它们并集的面积。
IoU(A, B) = Area(A ∩ B) / Area(A ∪ B)
其中:
* Area(A ∩ B) 是边界框 A 和 B 交集的面积。
* Area(A ∪ B) 是边界框 A 和 B 并集的面积。
并集的面积可以表示为:Area(A ∪ B) = Area(A) + Area(B) – Area(A ∩ B)。
IoU的值介于0到1之间。
* IoU = 0 表示两个框完全不重叠。
* IoU = 1 表示两个框完全重合。
* IoU 值越大,表示两个框重叠程度越高。
在NMS中,我们使用IoU来判断两个边界框是否指向同一个物体。如果两个框的IoU大于某个预设的阈值(称为NMS阈值),并且其中一个框的得分远高于另一个,那么得分低的那个框很可能是冗余的,应该被抑制。
2. NMS算法步骤(经典贪婪NMS)
经典的NMS算法是一个贪婪过程。其基本步骤如下:
假设我们有一个列表 B
,其中包含所有候选预测边界框,每个框 b
都有其坐标信息和对应的置信度得分 s
。还有一个预设的置信度阈值 score_threshold
和一个NMS阈值 nms_threshold
。
- 按得分过滤: 首先,移除所有置信度得分低于
score_threshold
的边界框。这一步是为了初步筛选掉那些不太可能是真实物体的预测。 - 按得分排序: 将剩余的所有边界框按它们的置信度得分进行降序排列。得到排序后的列表
B_sorted
。 - 初始化输出列表: 创建一个空的列表
D
,用于存储最终保留的边界框。 - 循环处理: 当
B_sorted
列表不为空时,执行以下操作:- a. 选取得分最高的框: 从
B_sorted
中选取第一个(即得分最高的)边界框m
。 - b. 添加到输出列表: 将边界框
m
从B_sorted
中移除,并添加到最终输出列表D
中。 - c. 抑制重叠框: 遍历
B_sorted
中剩余的所有边界框b_i
:- 计算
m
与b_i
之间的交并比IoU(m, b_i)
。 - 如果
IoU(m, b_i)
大于nms_threshold
,则认为b_i
是m
的冗余预测,将其从B_sorted
中移除。
- 计算
- a. 选取得分最高的框: 从
- 重复: 重复步骤4,直到
B_sorted
为空。 - 输出结果: 列表
D
中包含了经过NMS处理后的最终预测边界框集合。
举例说明:
假设我们有以下几个针对同一物体的预测框:
* 框 A: Score = 0.98
* 框 B: Score = 0.95, IoU(A, B) = 0.9
* 框 C: Score = 0.85, IoU(A, C) = 0.8
* 框 D: Score = 0.60, IoU(A, D) = 0.7
* 框 E: Score = 0.90, IoU(A, E) = 0.2 (这是另一个物体的框,或者是一个背景误检,与A重叠度低)
假设 score_threshold = 0.5
,nms_threshold = 0.5
。
- 所有框得分都大于0.5,保留。
- 按得分排序:A (0.98), B (0.95), E (0.90), C (0.85), D (0.60)。
D
= []- 第一次循环:
- a. 选取 A (得分 0.98)。
- b. 将 A 移入
D
。D
= [A],B_sorted
= [B, E, C, D]。 - c. 遍历剩余框:
- IoU(A, B) = 0.9 > 0.5,移除 B。
- IoU(A, E) = 0.2 <= 0.5,保留 E。
- IoU(A, C) = 0.8 > 0.5,移除 C。
- IoU(A, D) = 0.7 > 0.5,移除 D。
B_sorted
现在只剩下 [E]。
- 第二次循环:
- a. 选取 E (得分 0.90)。
- b. 将 E 移入
D
。D
= [A, E],B_sorted
= []。 - c.
B_sorted
为空,没有剩余框需要遍历。
B_sorted
为空,循环结束。- 最终输出列表
D
= [A, E]。
在这个例子中,框 A、B、C、D都指向同一个物体,其中 A 的得分最高。NMS保留了 A,并抑制了与其高度重叠(IoU > 0.5)的 B、C、D。框 E 因为与 A 重叠度不高(IoU <= 0.5),被认为是另一个独立的预测(可能指向另一个物体或背景),所以它也被保留下来进行下一轮处理。
3. 关键参数的影响
- 置信度阈值 (
score_threshold
): 决定了哪些低得分的预测框会被直接丢弃,不参与NMS过程。设置得过高可能漏掉一些真实但置信度稍低的物体;设置得过低会增加NMS的计算量,并可能保留更多低质量的误检。 - NMS阈值 (
nms_threshold
): 这是NMS算法的核心参数。- 设置得过高(例如0.7-0.9):会保留更多重叠的框。在物体紧密堆叠(如人群)的情况下,这可能有助于保留所有物体,但也可能导致同一个物体有多个框。
- 设置得过低(例如0.3-0.5):会更激进地抑制重叠框。在物体之间有明显间隔的情况下效果很好,但在物体紧密堆叠时,可能会错误地将相邻的两个不同物体的框抑制掉,导致漏检。
选择合适的NMS阈值是实际应用中的一个重要调优步骤,它需要在“避免冗余框”和“避免漏检紧密相邻物体”之间取得平衡。通常通过在验证集上实验来确定。
4. 计算效率考虑
经典的NMS算法在最坏情况下的时间复杂度是 O(N^2),其中N是候选框的数量。这是因为对于选定的最高得分框,需要与所有剩余的框计算IoU。虽然通过排序和逐步移除可以优化,但在实际应用中,特别是当N非常大时(例如数万个框),O(N^2) 的计算仍然可能成为瓶颈。
然而,在GPU上实现时,由于可以并行计算IoU,并且许多深度学习框架提供了高度优化的NMS实现,其实际运行速度通常可以满足实时性要求。一些实现会利用边界框在空间上的局部性,使用KD树或R树等空间数据结构来加速查找可能重叠的框,避免与所有框进行IoU计算,从而将复杂度降低到接近 O(N log N) 或 O(N)。
NMS在不同目标检测框架中的应用
NMS作为一个通用的后处理步骤,被广泛应用于各种深度学习目标检测框架:
- Faster R-CNN系列(Faster R-CNN, Mask R-CNN等): 在Region Proposal Network (RPN) 阶段,会生成大量的Anchor框,并预测它们是前景还是背景,以及进行初步的框回归。在RPN的输出中,会先进行一次NMS(通常称为RPN NMS)来筛选出少量(例如2000个)高质量的候选区域(Region Proposals)。然后这些候选区域被送入检测头(Detection Head)进行更精确的分类和回归。检测头会为每个候选区域输出更精细的框和类别的置信度。最后,对检测头输出的所有预测框(跨越所有候选区域)再进行一次NMS(通常称为Detection NMS或Final NMS),以获得最终的检测结果。这里通常对每个类别独立地进行NMS。
- YOLO系列(YOLOv1-YOLOv8, YOLO-R, PP-YOLO等): YOLO系列是一阶段检测器,它将图像划分为网格,并在每个网格上预测边界框、置信度和类别概率。这同样会生成大量的预测框(网格数 * 每个网格预测的框数)。对所有这些预测框进行NMS是后处理的关键一步,用于筛选最终的检测结果。早期YOLO版本直接在整个图像上应用NMS,而后续版本可能在一些中间步骤或对不同尺度的预测分别应用NMS,最后再合并。YOLO的NMS通常也是按类别独立进行的。
- SSD (Single Shot MultiBox Detector): SSD同样是一阶段检测器,在不同尺度的特征图上预设了大量的默认框(Default Boxes)。网络预测这些默认框的偏移量以及每个默认框包含各类物体的置信度。最终输出的预测框数量巨大。NMS是SSD中不可或缺的后处理步骤,用于从这些预测中选出最终结果。与Faster R-CNN和YOLO类似,SSD通常也是对每个类别独立执行NMS。
- RetinaNet (Focal Loss引入): RetinaNet也是一阶段Anchor-based检测器,通过特征金字塔网络(FPN)在多尺度上预测。它使用Focal Loss解决了前景背景类别不平衡问题,但同样需要NMS来处理输出的数千个预测框。
- Anchor-free检测器(如CenterNet, CornerNet, FCOS等): 这类检测器不依赖预设的Anchor框,而是预测物体的中心点、角点或其他关键点,并由此推断边界框。虽然它们生成候选框的方式与Anchor-based方法不同,但最终输出的预测框集合(例如,CenterNet在中心点预测大小,FCOS在每个点预测到边的距离)仍然可能存在重叠。因此,即使是Anchor-free检测器,在很多情况下也需要进行NMS或类似的后处理来去除冗余预测。例如,CenterNet可能会在热图的峰值点生成框,但即使是这样,也可能需要NMS来处理相邻峰值或不同尺度的预测。
总结来说,无论检测器架构如何变化,只要其输出是大量的、可能重叠的边界框预测,NMS(或其变体)就几乎是不可避免的后处理步骤。其基本作用始终是:从冗余预测中选出最佳代表。
传统NMS的局限性
尽管经典NMS简单有效且被广泛应用,但它也存在一些固有的局限性:
- 阈值敏感性: NMS阈值是一个超参数,其选择对性能影响很大,尤其是在不同数据集或不同场景下。一个固定的NMS阈值难以适应所有情况。例如,对于相距较远的物体,高阈值没问题;但对于拥挤场景中紧密相邻的同类物体,高阈值可能导致同一个物体出现多个框,而低阈值则可能错误地抑制掉一个真实物体(因为它与得分最高的邻居重叠度超过了低阈值)。
- 处理拥挤场景的困难: 这是传统NMS最主要的弱点。如果两个真实物体是同类,且它们靠得非常近,它们的边界框可能会高度重叠。在这种情况下,如果得分最高的那个框与相邻真实物体的框重叠度很高,传统NMS可能会将相邻真实物体的框误判为冗余而抑制掉,导致漏检。例如,在人群密集的场景中,相邻两个人的检测框可能IoU很高,得分较高的那个框会导致得分较低的那个框被移除,即使后者指向的是另一个真实的人。
- 硬阈值效应: NMS的抑制是“硬”的:如果IoU大于阈值,无论得分多高(除了最高的那个),都会被完全移除;如果IoU小于等于阈值,即使得分很低但与高分框重叠一点点,也不会被直接移除(除非它在后续迭代中成为最高分或者被其他高分框抑制)。这种硬性的判断方式不够灵活。
- 计算效率(在大规模数据集和模型输出框极多的情况下): 虽然GPU加速缓解了问题,但O(N^2)的理论复杂度在某些极端情况下仍然可能成为瓶颈,尤其是在需要处理数百万个候选框的场景。
这些局限性促使研究人员提出了各种改进和替代方案。
高级NMS变体与替代方法
为了克服传统NMS的局限性,特别是提高在拥挤场景下的性能和处理不同重叠程度的能力,研究人员提出了多种改进的NMS算法和NMS-free的检测方法。
1. Soft-NMS(软NMS)
Soft-NMS(由Bodla等人在2017年提出)是传统NMS的一个直接且有效的改进。它不再是硬性地移除重叠框,而是降低那些与高得分框重叠度较高的框的得分。这样,即使一个框与最高得分框高度重叠,如果其原始得分较高,经过分数衰减后仍有可能保留下来,尤其是在其指向的是一个独立的真实物体时。
Soft-NMS的算法流程与传统NMS类似,但在抑制重叠框的步骤有所不同:
- 按得分过滤和排序: 同传统NMS。
- 初始化输出列表: 同传统NMS。
- 循环处理: 当
B_sorted
列表不为空时:- a. 选取得分最高的边界框
m
。 - b. 将
m
移入D
。 - c. 遍历
B_sorted
中剩余的所有边界框b_i
:- 计算
IoU(m, b_i)
。 - 与传统NMS不同之处: 根据
IoU(m, b_i)
的值,对b_i
的得分s_i
进行衰减。常用的衰减函数有两种:- 线性衰减:
if IoU(m, b_i) > nms_threshold:
s_i = s_i * (1 - IoU(m, b_i))
else:
s_i = s_i
如果重叠度高,得分线性下降。 - 高斯衰减:
s_i = s_i * exp(-(IoU(m, b_i)^2 / sigma))
其中sigma
是一个可调参数。高斯衰减更加平滑,重叠度越高,衰减越剧烈。
- 线性衰减:
- (可选)如果衰减后的得分低于某个非常低的阈值,可以将该框移除以提高效率。
- 计算
- d. 将经过得分衰减处理后的
b_i
保留在B_sorted
中(只是分数变了),然后移除m
。
- a. 选取得分最高的边界框
- 重复: 重复步骤3,直到
B_sorted
为空。 - 最终过滤: 最后,从列表
D
中移除所有得分低于最终输出阈值的框。
Soft-NMS的优势在于,它将硬抑制变成了软抑制,对于紧密重叠的框,不再是“要么留下要么移除”,而是“得分高的留下,得分低的根据重叠度降低分数”。这在处理拥挤场景时表现更好,因为它允许两个紧密相邻的真实物体都有机会保留其对应的预测框,即使其中一个得分相对较低。缺点是计算量比传统NMS稍大(需要更新分数而不是直接移除),但O(N^2)的复杂度不变。
2. 其他学习型或改进方法
- IoU-Net等基于IoU预测的方法: 一些工作尝试预测边界框的IoU置信度,并用这个IoU置信度来代替或结合原始的分类置信度进行NMS。这样可以更好地衡量框的定位精度,有助于选出定位更准确的框。
- Learning-based NMS: 研究人员尝试用神经网络来替代传统的NMS过程。例如,NMSNet直接学习一个模型来预测哪些框应该被保留;或者使用图神经网络(GNN)来建模候选框之间的关系(如位置、大小、类别、重叠度),然后通过GNN来决定哪些框应该被保留或如何调整它们的得分。这些方法旨在通过学习数据中的复杂模式来克服传统NMS的局限性,特别是在极端拥挤或不规则重叠的情况下。它们的潜力很大,但也带来了模型的复杂性和训练的难度。
- 矩阵运算NMS: 为了提高计算效率,特别是利用GPU并行计算能力,许多深度学习框架(如TensorFlow, PyTorch)提供了基于矩阵运算优化的NMS实现。它们通常不是严格遵循上述贪婪算法,而是通过构建IoU矩阵和利用并行比较来快速筛选掉大量冗余框,效率远高于naive的Python循环实现。
- NMS-free检测器: 这是一个更根本的思路,即设计一种检测器架构,其输出本身就具有唯一性或冗余性很低,从而完全避免或大大简化NMS后处理。CenterNet就是一个典型的例子,它检测物体的中心点,理论上每个物体只有一个中心点(在整数像素位置上)。虽然实际实现中可能需要一些后处理来聚合或清理,但其目标是最小化或消除对NMS的需求。CornerNet和ExtremeNet等基于关键点的方法也有类似的思想。这些方法的兴起代表了目标检测后处理发展的一个重要方向。
尽管Soft-NMS和Learning-based方法提供了更好的理论性能,但在许多实际应用中,考虑到实现复杂度、推理速度和足够的性能,传统的贪婪NMS(尤其是GPU优化的版本)仍然是主流选择,特别是对于那些非极端拥挤的场景。Soft-NMS因其简单有效,也被广泛采用。
总结
NMS(非极大值抑制)是深度学习目标检测流程中一个至关重要的后处理步骤。它解决了深度学习模型为了召回率而生成大量冗余预测框的问题,通过计算边界框之间的交并比(IoU),并结合置信度得分,以贪婪的方式选取局部最优的边界框,并抑制与其高度重叠的低得分框。
本文详细阐述了传统NMS的原理、IoU的计算、完整的算法流程以及关键参数(置信度阈值和NMS阈值)的影响。我们看到,无论是两阶段还是单阶段的检测器,NMS都扮演着从海量预测中精选最终结果的角色。
然而,传统NMS在处理拥挤场景和对阈值敏感等方面存在局限性。这催生了Soft-NMS等改进算法,它们通过软化抑制机制来提高在复杂场景下的鲁棒性。此外,基于学习的NMS方法和NMS-free的检测器设计代表了未来后处理或避免后处理的发展方向。
理解NMS不仅有助于我们更好地理解目标检测的完整流程,也是在实际应用中对检测器性能进行调优、处理特定场景(如密集物体检测)的关键。尽管新的后处理技术不断涌现,经典NMS作为一种基础且高效的方法,在未来很长一段时间内仍将是目标检测领域不可或缺的一部分。对它的深入理解,是掌握现代目标检测技术的基石。