NMS算法(非极大值抑制)是目标检测算法中经典的后处理步骤,其本质是搜索局部最大值,抑制非极大值元素。主要利用目标检测框以及对应的置信度分数,设置一定的阈值来删除重叠较大的边界框。
其算法流程如下:
- 根据置信度得分进行排序
- 选择置信度最高的目标检测框添加到输出列表中,将其从检测框列表中删除
- 计算该检测框与剩余候选检测框的IOU
- 删除IOU大于阈值的检测框
- 重复上述4步,直至检测框列表为空
其算法实现如下
import numpy as np
def nms(dets, thresh): # x1, y1, x2, y2, score
x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 各个方框的面积
order = scores.argsort()[::-1] # 按置信度排序后的index, 作为候选集
keep = [] # 保存筛选出来的方框的index
while order.size > 0:
i = order[0] # 当前置信度最大的方框
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy1 = np.maximum(y1[i], y1[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 / (areas[i] + areas[order[1:]] - inter)
inds = np.where(overlap <= thresh)[0] # 交并比小于thresh的仍然保留在候选集里, 大的过滤掉
order = order[inds + 1] # inds + 1对应原来order中overlap小于thresh的项
return keep
if __name__ == '__main__':
detections = [
[10, 20, 100, 100, 0.9],
[20, 10, 110, 100, 0.88],
[20, 20, 110, 110, 0.86],
[40, 50, 200, 200, 0.95],
[45, 52, 198, 202, 0.87]
]
detections = np.array(detections)
keeps = nms(detections, 0.5)
print(detections[keeps])