NMS的理解及Python实现

NMS的定义

全称为non maximum suppression，翻译成中文是非极大抑制。怎么理解呢？
这里上一张图。

如上图所示，绿色的框识别出了鸣人， 红色的识别出了水门， 蓝色的则识别出了自来也。
但是，我们对于每个对象只需要一个框就够了，在这里，无论是鸣人，水门还是自来也都有三个框，我们应该保留哪个框呢？当然是得分（预测概率）最高的哪个。
NMS就是干这个活的，也就是，不是得分最高的（极大值）的框都可以滚蛋了。

NMS实现

步骤1：按照得分的顺序将所有bounding box进行排序。
步骤2：最高的那个肯定留下，然后计算它与其他所有bounding box（得分都低于他）的IoU（交并比，不解释了）。若某个bounding box的ioU小于设定的阈值，那就留下他，不然就删掉他。
步骤3：余下的继续执行1~2步
这里举个例子：
给定上图的3个人物的3个bounding box及得分数组[x1, y1, x2, y2, score]。
第一行三个是鸣人，第二行三个是水门，第三行三个是自来也。

dets = [[218, 322, 385, 491, 0.98],[247, 312, 419, 461, 0.83],[237, 344, 407, 510, 0.92],
            [757, 218, 937, 394, 0.96],[768, 198, 962, 364, 0.85],[740, 240, 906, 414, 0.83],
            [1101, 84, 1302, 303, 0.82], [1110, 67, 1331, 260, 0.97], [1123, 42, 1362, 220, 0.85]]

我们按照上述步骤可以发现，得分最高的是鸣人的某个框(0.98)，那我们保留，并计算该框与其他所有框的IoU。这样会形成8个IoU值（原本9个，拿出一个和其他8个比）。遍历该IoU数组，假设阈值为0.5，如果ioU小于我们的阈值，我们肯定留下它，否则就干掉它。这里注意，虽然需要和8个bounding box对比，但是实际上鸣人的框与自来也还有水门的框的IoU是0，所以一定会保留水门和自来也的框，也就是该轮操作至少剩下6个框。

overlap [0.         0.         0.62888921 0.         0.         0.
 0.55520685 0.        ]

结果我们看到，果然，我们要保留这6个0，而IoU为0.62和0.55的那俩框其实就是鸣人的除最高得分的另外两个框，怎么办？删除就好了。
剩下的就是水门和自来也啦。

Python代码实现

import numpy as np

def py_cpu_nms(dets, thres):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]
    print("order", order)

    # needs to save
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2-xx1+1)
        h = np.maximum(0.0, yy2-yy1+1)

        inter = w * h
        overlap = inter / (area[i] + area[order[1:]] - inter)
        print("overlap", overlap)

        inds = np.where(overlap <= thres)[0]
        print("ind", inds)
        order = order[inds+1]
        print("order", order)

    return keep

if __name__ == "__main__":
    dets = [[218, 322, 385, 491, 0.98],[247, 312, 419, 461, 0.83],[237, 344, 407, 510, 0.92],
            [757, 218, 937, 394, 0.96],[768, 198, 962, 364, 0.85],[740, 240, 906, 414, 0.83],
            [1101, 84, 1302, 303, 0.82], [1110, 67, 1331, 260, 0.97], [1123, 42, 1362, 220, 0.85]]
    # [757, 218, 937, 394, 0.96],[768, 198, 962, 364, 0.85],[740, 240, 906, 414, 0.83]
    dets = np.array(dets)
    print(py_cpu_nms(dets, 0.5))

代码运行过程如下：

order [0 7 3 2 8 4 5 1 6]
overlap [0.         0.         0.62888921 0.         0.         0.
 0.55520685 0.        ]
ind [0 1 3 4 5 7]
order [7 3 8 4 5 6]
overlap [0.         0.59778611 0.         0.         0.64035466]
ind [0 2 3]
order [3 4 5]
overlap [0.63086943 0.61164895]
ind []
order []
result [0, 7, 3]

按照我们刚才输入的顺序，我们发现，最终留下的bounding box恰恰就是鸣人，水门和自来也得分最高的那三个。
代码还是比较简单，就不解释啦。

加油，为了未来！为了瓜瓜！