Python实现NMS(非极大值抑制)通用模板

1. 代码

import numpy as np
 
def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    #dets:N*M,N是bbox的个数，M的前4位是对应的（x1,y1,x2,y2），第5位是对应的分数
    #thresh:0.3,0.5....
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]
 
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)#求每个bbox的面积
    order = scores.argsort()[::-1]#对分数进行倒排序
 
    keep = []#用来保存最后留下来的bbox
    while order.size > 0:
        i = order[0]#无条件保留每次迭代中置信度最高的bbox
        keep.append(i)
        #计算置信度最高的bbox和其他剩下bbox之间的交叉区域
        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:]])
        #计算置信度高的bbox和其他剩下bbox之间交叉区域的面积
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        #求交叉区域的面积占两者（置信度高的bbox和其他bbox）面积和的必烈
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        #保留ovr小于thresh的bbox，进入下一次迭代。
        inds = np.where(ovr <= thresh)[0]
        #因为ovr中的索引不包括order[0]所以要向后移动一位
        order = order[inds + 1]
 
    return keep

2. 参考文献
https://blog.youkuaiyun.com/Suan2014/article/details/79640461