目标检测 NMS 的 Python 实现（附算法动图）

在这篇文章中，我将给出一份带运行示例的NMS Python脚本，并对算法和代码进行详细解说。相信大家看完这篇文章后，能够轻松地掌握NMS的底层原理。

如果你对目标检测的基础知识不太熟，欢迎先阅读我的上篇文章：目标检测简介

示例脚本（包括可视化的代码）链接：https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/nms

算法介绍

在目标检测算法中，为了尽量不漏掉物体，会输出大量的检测结果（每一条结果由检测概率与检测框组成）。这些检测结果很可能有重复的，即会有多个框标出了同一个物体。下图就是一个例子，算法正确识别出了两辆车，但却有多个检测框标出了同一辆车。

我们需要一个算法来过滤多余的检测框。最常用的算法就是NMS(Non-Maximum Suppresion, 非极大值抑制)。该算法的思路很简单：只保留局部概率最大的检测框，与其重合的其他检测框都会被舍去。

算法的伪代码如下：

输入所有检测概率、检测框
当还有检测框没有处理：
    从剩下的框里挑出检测概率最大的框 bbox_a
    遍历所有没有处理的框bbox_i：
        如果 bbox_i != bbox_a 且 bbox_i 与 bbox_a 重合：
            舍去 bbox_i
    把 bbox_a 输出成一个检测结果

当然，这个算法的描述还不够准确：究竟该怎么定义两个检测框是“重合”呢？如果两个检测框有交集就说它们重合是行不通的，因为图片中很可能会有挨得很近的物体，它们的检测框就是相交的。因此，我们一般用IoU（交并比）来描述两个框的重合程度，如果IoU超出某个阈值，就说这两个框是“重合”的。

IoU的计算很简单，算出两个矩形的「交面积」，再用两个矩形面积之和减去「交面积」就可以得到「并面积」，「交面积」比「并面积」就是IoU。

在NMS之前，一般还会先做一步预处理：对于预测概率小于某个阈值的检测框，我们不信任它们，会在进行NMS之前就把它们舍去。在代码实现时这部分逻辑常常会放到NMS的函数里。这样，整个NMS的流程是这样的：

上述的NMS针对是识别一种物体的任务，在推广到多分类NMS时，只要把输入的检测概率换成有物体的概率乘上概率最大的类别的概率即可。

代码实现

理论上这段代码是兼容任何格式的张量的。经测试，NumPy, PyTorch 的张量都可以正确运行。

先看一下IoU的实现代码：

def iou(b1: Tuple[int, int, int, int], b2: Tuple[int, int, int, int]) -> float:
    intersection = box_intersection(b1, b2)
    inter_area = area(intersection)
    union_area = area(b1) + area(b2) - inter_area
    return inter_area / union_area

所有的检测框用一个长度为4的元组表示。box_intersection用于计算两个框的相交框，area用于计算框的面积。这段代码和之前描述得一样，先获取相交区域，计算交面积，再计算并面积，最后算除法。

box_intersection的实现如下：

def box_intersection(
        b1: Tuple[int, int, int, int],
        b2: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]:
    x11, y11, x12, y12 = b1
    x21, y21, x22, y22 = b2

    xl = max(x11, x21)
    xr = min(x12, x22)
    yt = max(y11, y21)
    yb = min(y12, y22)
    return (xl, yt, xr, yb)

算相交区域，可以理解成分别求出x, y方向上的相交部分，再把两部分合成一个框。而求直线的相交，就是取最大的左端点和最小的右端点。

area的实现如下：

def area(box: Tuple[int, int, int, int]) -> float:
    x1, y1, x2, y2 = box
    width = max(x2 - x1, 0)
    height = max(y2 - y1, 0)
    return width * height

如果两个框不相交，x2 - x1和y2 - y1会出现小于0的情况。因此，要保证它们最小为0，再算面积就不会有错了。

有了iou，就可以实现了NMS了。我的NMS函数定义如下：

def nms(predicts: np.ndarray,
        score_thresh: float = 0.6,
        iou_thresh: float = 0.3):
    """Non-Maximum Suppression

    Args:
        predicts (np.ndarray): Tensor of shape [n, 5]. The second demesion
            includes 1 probability and 4 numbers x, y, w, h denoting a bounding
            box.
        score_thresh (float): The boxes with probability lower than
            score_threash will be discarded.
        iou_thresh (float): The threshold determining whether two boxes are
            "overlapped".

    Returns:
        (np.ndarray, List[int]): The filtered predictions and the indices of
            remaining boxes.
    """

NMS算法需要输入检测结果predicts，输出过滤后的检测结果filtered_predicts。此外，NMS算法有两个输入属性score_thresh， iou_thresh，分别表示被选定的检测框最小需要的概率、判断两个框是否重合的IoU阈值。为了方便其他的计算（比如多分类NMS），我还输出了一个索引数组indices，表示被选中的框的索引。这方面的逻辑可以根据自己的项目要求进行优化，没有统一的规定。

NMS的实现和开始的伪代码几乎一模一样：

def nms(predicts: np.ndarray,
        score_thresh: float = 0.6,
        iou_thresh: float = 0.3):
    n_remainder = len(predicts)
    vis = [False] * n_remainder

    # Filter predicts with low probability
    for i, predict in enumerate(predicts):
        if predict[0] < score_thresh:
            vis[i] = True
            n_remainder -= 1

    # NMS
    output_predicts = []
    output_indices = []
    while n_remainder > 0:
        max_pro = -1
        max_index = 0
        # Find argmax
        for i, p in enumerate(predicts):
            if not vis[i]:
                if max_pro < p[0]:
                    max_index = i
                    max_pro = p[0]

        # Append output
        max_p = predicts[max_index]
        output_predicts.append(max_p)
        output_indices.append(max_index)

        # Suppress
        for i, p in enumerate(predicts):
            if not vis[i] and i != max_index:
                if iou(p[1:5], max_p[1:5]) > iou_thresh:
                    vis[i] = True
                    n_remainder -= 1
        vis[max_index] = True
        n_remainder -= 1

    return output_predicts, output_indices

一开始，先声明一些辅助的变量。n_remainder表示还有多少个框没被访问，vis[i]表示第i个框是否被访问。

n_remainder = len(predicts)
vis = [False] * n_remainder

一开始，先过滤那些概率过小的框：

for i, predict in enumerate(predicts):
    if predict[0] < score_thresh:
        vis[i] = True
        n_remainder -= 1

之后，正式进入NMS，先准备好输出的列表：

# NMS
output_predicts = []
output_indices = []
    while n_remainder > 0:

在还有没访问的框时，先找出剩下的框中概率最大的那个框：

while n_remainder > 0:
    max_pro = -1
    max_index = 0
    # Find argmax
    for i, p in enumerate(predicts):
        if not vis[i]:
            if max_pro < p[0]:
                max_index = i
                max_pro = p[0]

之后，抑制掉和概率最大框“重合”的框。

while n_remainder > 0:
    # Find argmax
    ...

    max_p = predicts[max_index]
    # Suppress
    for i, p in enumerate(predicts):
        if not vis[i] and i != max_index:
            if iou(p[1:5], max_p[1:5]) > iou_thresh:
                vis[i] = True
                n_remainder -= 1

最后，把这个结果添加进输出，并维护好vis和n_remainder。

while n_remainder > 0:
    # Find argmax
    ...

    # Suppress
    ...

    # Append output
    output_predicts.append(max_p)
    output_indices.append(max_index)
    vis[max_index] = True
    n_remainder -= 1

这里可以把vis[max_index] = True那两行移到抑制操作的前面，这样判断里就不用加i != max_index了。我这样写是为了强调一下，判断重合的时候不需要判断自己。

实现完了，返回结果。

return output_predicts, output_indices

单元测试

做单元测试的时候，最好是找一份现有的实现用做对照。为了让整个测试过程更贴合实际一些，我用MMDetection的YOLOv3跑了一个NMS前和NMS后的检测框结果，并把NMS前的检测框输入进了自己实现的NMS，比较了两份NMS的输出。以下是具体的测试过程。

照着MMDetection的官方文档，下载好YOLOv3模型，用下面的代码即可对MMDetection的demo图片进行推理并保存结果：

from mmdet.apis import inference_detector, init_detector, show_result_pyplot
from mmdet.models.detectors import BaseDetector

# Choose to use a config and initialize the detector
config = 'configs/yolo/yolov3_mobilenetv2_320_300e_coco.py'
# Setup a checkpoint file to load
checkpoint = 'checkpoints/yolov3_mobilenetv2_320_300e_coco_20210719_215349-d18dff72.pth'
# initialize the detector
model: BaseDetector = init_detector(config, checkpoint, device='cuda:0')

img = 'demo/demo.jpg'
result = inference_detector(model, img)

model.show_result(img, result, score_thr=0.3, out_file='work_dirs/1.jpg')

为了得到NMS前的输入，我在mmdet/models/dense_head/YOLOV3Head.get_bboxes里面插入了一段输出结果的代码。

...

save_dict = {
    'bboxes': bboxes,
    'cls_scores': scores,
    'prob': objectness
}
torch.save(save_dict, 'work_dirs/bboxes.pt')

det_bboxes, det_labels = multiclass_nms(
...

得到NMS前的输入是

把这份数据输入进我们的NMS中，得到的可视化结果如下：

这跟开始那份输出一模一样。看来我们实现的这份NMS完全正确。

如果你对测试过程、可视化、多分类NMS感兴趣，欢迎直接阅读我的项目源码。

总结

在这篇文章中，我给出了一份NMS的简洁而靠近底层的Python实现，并对NMS算法进行了介绍。通过阅读这篇文章，相信大家已经完全理解了NMS的原理，并且能够用任何一种语言实现NMS。一般的NMS开源实现支持的参数更多，代码会更复杂一些，但它们的核心和我的这份实现是一样的。

这份NMS的实现还有很大的改进空间。比如每轮求概率最大的框时，可以先排好序，或者用优先队列，这样均摊下来每次获取概率最大的框的复杂度是O(logn)。但是后面判断重复的框一定有一个O(n)的计算，这部分的优化并不显著。大家有余力可以参考成熟的CV项目里NMS是怎么高效用C++实现的。