话不多说,直接上代码,代码是根据soft max论文提供的代码,修改的,原来的代码是cpu版本,且返回的keep参数并不是原来bbox的序列参数,所以做了部分修改,本人亲测pytorch版本的maskrcnn可用,提高了我的部分结果。
def mine_soft_nms(boxes, sigma = 0.05, Nt = 0.3, threshold = 0.001, method = 2):
“ori list is an order to log the bbox order!”
N = boxes.shape[0]
ori_list = list(np.arange(N))
for i in range(N):
maxscore = boxes[i, 4]
maxpos = i
tx1 = boxes[i, 0]
ty1 = boxes[i, 1]
tx2 = boxes[i, 2]
ty2 = boxes[i, 3]
ts = boxes[i, 4]
pos = i + 1
while pos < N:
if maxscore < boxes[pos, 4]:
maxscore = boxes[pos, 4]
maxpos = pos
pos = pos + 1
boxes[i, 0] = boxes[maxpos, 0]
boxes[i, 1] = boxes[maxpos, 1]
boxes[i, 2] = boxes[maxpos, 2]
boxes[i, 3] = boxes[maxpos, 3]
boxes[i, 4] = boxes[maxpos, 4]
boxes[maxpos, 0] = tx1
boxes[maxpos, 1] = ty1
boxes[maxpos, 2] = tx2
boxes[maxpos, 3] = ty2
boxes[maxpos, 4] = ts
tx1 = boxes[i, 0]
ty1 = boxes[i, 1]
tx2 = boxes[i, 2]
ty2 = boxes[i, 3]
ts = boxes[i, 4]
ori_list[i], ori_list[maxpos] = ori_list[maxpos], ori_list[i]
pos = i + 1
while pos < N:
x1 = boxes[pos, 0]
y1 = boxes[pos, 1]
x2 = boxes[pos, 2]
y2 = boxes[pos, 3]
s = boxes[pos, 4]
area = (x2 - x1 + 1) * (y2 - y1 + 1)
iw = (min(tx2, x2) - max(tx1, x1) + 1)
if iw > 0:
ih = (min(ty2, y2) - max(ty1, y1) + 1)
if ih > 0:
ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih)
ov = iw * ih / ua
if method == 1:
if ov > Nt:
weight = 1 - ov
else:
weight = 1
elif method == 2:
weight = torch.exp(-(ov * ov) / sigma)
else:
if ov > Nt:
weight = 0
else:
weight = 1
boxes[pos, 4] = weight * boxes[pos, 4]
if boxes[pos, 4] < threshold:
boxes[pos, 0] = boxes[N - 1, 0]
boxes[pos, 1] = boxes[N - 1, 1]
boxes[pos, 2] = boxes[N - 1, 2]
boxes[pos, 3] = boxes[N - 1, 3]
boxes[pos, 4] = boxes[N - 1, 4]
ori_list[pos], ori_list[N - 1] = ori_list[N - 1], ori_list[pos]
N = N - 1
pos = pos - 1
pos = pos + 1
keep = [i for i in range(N)]
mine_keep = ori_list[:N]
print("soft nms N: ", N)
return torch.LongTensor(mine_keep).cuda()
输入的参数都能看懂,和你的原来的NMS代码的输入输出对其就可以了!
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)