nms修改为soft nms
一,NMS与soft nms原理
先放一张论文中的图
在对上述图片中的目标进行检测过程中,使用非极大值抑制(NMS)进行筛选时只能获取到得分最高的框(得分0.95),对于得分较低的框会抑制掉。但由上图中可以明显看到两个框其实都框住了物体,这就会产生错误的筛选。
soft nms为了改善这种情况(并没有完全消除!)提出了如下的两个函数:
1.线性表达式:
2.高斯表达式(常用):
二,修改后的代码
import numpy as np
def non_max_suppression_fast(boxes,probs,Nt=0.3,sigma=0.5,overlap_thresh=0.001,method=2):
dets = boxes
N = dets.shape[0]
indexes = np.array([np.arange(N)])
dets = np.concatenate((dets,indexes.T),axis=1)
y1 = dets[:, 1]
x1 = dets[:, 0]
y2 = dets[:, 3]
x2 = dets[:, 2]
scores = probs
areas = (x2-x1+1)*(y2-y1+1)
for i in range(N):
tBD = dets[i, :].copy()
tscore = scores[i].copy()
tarea = areas[i].copy()
pos = i+1
if i !=N-1:
maxscore = np.max(scores[pos:],axis=0)
maxpos = np.argmax(scores[pos:],axis=0)
else:
maxscore = scores[-1]
maxpos = 0
if tscore<maxscore:
dets[i, :] = dets[maxpos+i+1, :]
dets[maxpos+i+1, :] = tBD
tBD = dets[i,:]
scores[i] = scores[maxpos+i+1]
scores[maxpos+i+1] = tscore
tscore = scores[i]
areas[i] = areas[maxpos+i+1]
areas[maxpos+i+1] = tarea
tarea = areas[i]
xx1 = np.maximum(dets[i,0],dets[pos:,0])
yy1 = np.maximum(dets[i,1],dets[pos:,1])
xx2 = np.minimum(dets[i,2],dets[pos:,2])
yy2 = np.minimum(dets[i,3],dets[pos:,3])
w = np.maximum(0.0,xx2-xx1+1)
h = np.maximum(0.0,yy2-yy1+1)
inter = w*h
ovr = inter/(areas[i]+areas[pos:]-inter)
if method ==1:
weight = np.ones(ovr.shape)
weight[ovr>Nt] = weight[ovr>Nt]-ovr[ovr>Nt]
elif method == 2:
weight = np.exp(-(ovr*ovr)/sigma)
else:
weight = np.ones(ovr.shape)
weight[ovr>Nt] = 0
scores[pos:] = weight*scores[pos:]
inds = dets[:, 4][scores>overlap_thresh]
keep = inds.astype(int)
keeps = keep.tolist()
boxes = boxes[keeps].astype("int")
probs = probs[keeps]
return boxes,probs