物体检测之非极大值抑制NMS原理及代码

非极大值抑制算法：

在深度学习物体检测中的重要用于消除多余的bbox、留下置信度最高的bbox。

大概步骤：

• 起始，设所有的框都没有被抑制，所有框按照 score 从大到小排序。
• 从第 0 个框（分数最高）开始遍历：对于每一个框，如果该框没有被抑制，就将所有与它 IoU 大于 thresh 的框设为抑制。
• 返回没被抑制的框。

Python代码：

import numpy as np
test = [[83, 54, 165, 163, 0.8], [67, 48, 118, 132, 0.5], [91, 38, 192, 171, 0.6]]
def nms(dets, thresh):
    dets = np.array(dets)
    # 获取每个框（x1,y1）（x2,y2）
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    # 获取置信度
    score = dets[:, 4]
    # 置信度高到低排序
    order = np.argsort(score)[::-1]
    # 计算面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    # 存放要保留框的下坐标
    keep = []
    # 依次计算直至全部计算完成
    while order.size > 0:
        i = order[0]
        keep.append(i)
        # 获取重叠部分的（xx1,yy1）,(xx2,yy2)
        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:]])
        # 计算重叠部分面积（没有重叠的为0）
        inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1)
        # 计算IoU
        IoU = inter / (areas[i] + areas[order[1:]] - inter)
        # 求符合要求的index
        index = np.where(IoU <= thresh)[0]
        # 第一个参与计算所以index + 1
        order = order[index + 1]
    return keep

print(nms(test,0.5))

参考链接：

1. https://zhuanlan.zhihu.com/p/40976906
2. https://www.jianshu.com/p/d2c7f6d9708f