NMS算法步骤:
对于某一类的所有预测框dets,都有其相对应的置信度score,设最终要输出的结果列表为keep=[ ]
- 保存当前集合中置信度最高的边框至keep中,并从预测框集合中去掉它
- 依次计算该预测框与其他所有边框的IOU
- 去除所有IOU大于阈值的边框
- 如果预测框集合不为空,返回第一步重复上述三个步骤。
def NMS(dets,thresh):
#dets为某个类的框,形式为[[x1,y1,x2,y2,socre],[x1,y1,x2,y2,socre]]
x1 = dets[:,0]
y1 = dets[:,1]
x2 = dets[:,2]
y2 = dets[:,3]
scores = dets[:4]
# 每一个框的面积
areas = (x2-x1+1)*(y2-y1+1)
# 按照score降序排列,其中orders是序号
orders = scores.argsort()[::-1]
keep = []# 最终结果保存在keep中
while orders.size()>0:
i = orders[0]
keep.append(i) # 保留该类剩余box中得分最高的一个
#得到相交区域的左上和右下坐标
xx1 = np.maximum(x1[i],x1[orders[1:]])
yy1 = np.maximum(y1[i],y1[orders[1:]])
xx2 = np.maximum(x2[i],x2[orders[1:]])
yy2 = np.maximum(y2[i],y2[orders[1:]])
w = np.maximum(0.0,xx2-xx1+1)
h = np.maximum(0.0,yy2-yy1+1)
# 计算出相交区域的面积,不重叠时面积为0
inter = w*h
over = inter/(areas[i]+areas[orders[1:]]-inter)
# 保留IOU小于阈值的box
inds = np.where(over<=thresh)[0]
orders = orders[inds+1]
return keep