NMS算法(非极大值抑制)是目标检测算法中经典的后处理步骤,其本质是搜索局部最大值,抑制非极大值元素。主要利用目标检测框以及对应的置信度分数,设置一定的阈值来删除重叠较大的边界框。
其算法流程如下:
根据置信度得分进行排序
选择置信度最高的目标检测框添加到输出列表中,将其从检测框列表中删除
计算该检测框与剩余候选检测框的IOU
删除IOU大于阈值的检测框
重复上述4步,直至检测框列表为空
import numpy as np #假设生成9个候选框 boxes = np.array([[100,100,210,210,0.72], [250,25,420,420,0.87], [220,220,320,330,0.92], [120,130,210,210,0.73], [230,240,325,330,0.81], [220,230,315,340,0.93], [300,400,450,220,0.95], [200,150,500,100,0.85], [350,450,500,270,0.78] ]) def nms(bboxes,thresh): x1 = bboxes[:, 0] y1 = bboxes[:, 1] x2 = bboxes[:, 2] y2 = bboxes[:, 3] areas = (x2 - x1 +1) * (y2 - y1 +1) #进行排序 scores = bboxes[:, 4] index = scores.argsort()[::-1] res = [] while index.size > 0: i = index[0] res.append(i) X1 = np.maximum(x1[i],x1[index[1:]]) Y1 = np.maximum(y1[i],y1[index[1:]]) X2 = np.minimum(x2[i],x2[index[1:]]) Y2 = np.minimum(y2[i],y2[index[1:]]) w = np.maximum(0,X2 -X1 +1) h = np.maximum(0,Y2 - Y1 +1) jiao = w*h iou = jiao / (areas[i] + areas[index[1:]] - jiao) idx = np.where(iou <= thresh)[0] index = index[idx +1] # idx =np.where(iou <= thresh)[0] # index = index[idx+1] # return res return res