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恰恰就是鸣人,水门和自来也得分最高的那三个。
代码还是比较简单,就不解释啦。
加油,为了未来!为了瓜瓜!