NMS(非极大值抑制)python实现

一、原理

1）背景

基于深度学习的目标检测算法都会从一张图像中得到属于某个类别的多个矩形框，并且会得到每个矩形框属于这个类别的概率。

2）存在的问题
在检测结果中，往往会出现很多矩形框检测到的都是同一个目标的情况，但其实每个目标只需要一个最准确的矩形框 就可以了。

3）NMS 的目标
每个物体只保留最准确（置信度最高）的一个矩形框，其余的全删除。

二、流程

所谓非极大值抑制：依靠分类器得到多个候选框，以及关于候选框中属于类别的概率值，根据分类器得到的类别分类概率做排序，具体算法流程如下：

（1）将所有框的得分排序，选中最高分及其对应的框 
（2）遍历其余的框，如果和当前最高分框的重叠面积(IOU)大于一定阈值，我们就将框删除。（为什么要删除，是因为超过设定阈值，认为两个框的里面的物体属于同一个类别，比如都属于狗这个类别。我们只需要留下一个类别的可能性框图即可。）
（3）从未处理的框中继续选一个得分最高的，重复上述过程。

三、程序实现

import numpy as np

def nms(dets, thresh):
    #x[:,n]就是取所有集合的第n个数据
    x1 = dets[:, 0] #xmin
    y1 = dets[:, 1] #ymin
    x2 = dets[:, 2] #xmax
    y2 = dets[:, 3] #ymax
    scores = dets[:, 4] #confidence
    
    areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 每个boundingbox的面积
    order = scores.argsort()[::-1] # boundingbox的置信度排序
    keep = [] # 用来保存最后留下来的boundingbox
    while order.size > 0:     
        i = order[0] # 置信度最高的boundingbox的index
        keep.append(i) # 添加本次置信度最高的boundingbox的index
        
        # 当前bbox和剩下bbox之间的交叉区域
        # 选择大于x1,y1和小于x2,y2的区域
        # np.maximum(array1, array2)：逐位比较array1和array2，并输出两者的最大值。
        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 + 某区域面积 - 交叉区域面积)
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        #保留交集小于一定阈值的boundingbox
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
        
    return keep


if __name__ == '__main__':
    dets = np.array([
        [204, 102, 358, 250, 0.5],
        [257, 118, 380, 250, 0.7],
        [280, 135, 400, 250, 0.6],
        [255, 118, 360, 235, 0.7]])
    thresh = 0.7
    res = nms(dets, thresh)
    print(res)