import numpy as np
boxes = np.array([[200, 200, 400, 400], [220, 220, 420, 420], [200, 240, 400, 440], [240, 200, 440, 400], [1, 1, 2, 2]], dtype=np.float32)
boxscores = np.array([0.9, 0.8, 0.7, 0.6, 0.5], dtype=np.float32) # 不放在里面是因为scores会发生变化
def softnms(boxes_x1y1x2y2, score_boxes, threshold_iou = 0.3, sigma=0.5, threshold_end=0.001, method=2):
"""
softnms
:param boxes_x1y1x2y2: boexs 坐标矩阵 format [y1, x1, y2, x2]
:param score_boxes: 每个 boxes 对应的分数
:param threshold_iou: 其余boxes与一个最好score的box分别求解iou，大于score_boxes值表示与该box重叠很高的重复框去掉，小于该值的可能是其它object的框， 交叠门限
:param sigma: 使用 gaussian 函数的方差
:param threshold_end: 最后的分数门限,将一次一次求解最好的box对应的iou值保存在scores中，而后保留大于此阈值的scores对应boxes
:param method: 使用的方法
:return: 留下的 boxes 的 index
"""
# indexes concatenate boxes with the last column
N = boxes_x1y1x2y2.shape[0] # 表示boxes的个数
# the order of boxes coordinate is [y1,x1,y2,x2]
y1 = boxes_x1y1x2y2[:, 0]
x1 = boxes_x1y1x2y2[:, 1]
y2 = boxes_x1y1x2y2[:, 2]
x2 = boxes_x1y1x2y2[:, 3]
scores = score_boxes
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
for i in range(N):
# intermediate parameters for later parameters exchange
tBD = boxes_x1y1x2y2[i, :].copy() # 对应的第i个坐标
tscore = scores[i].copy() # 对应的第i个scores
tarea = areas[i].copy() # 对应的第i个面积
pos = i + 1 # pos为跳过当前第i个的box的索引
if i != N - 1: # 判断i是否取最大了
maxscore = np.max(scores[pos:], axis=0) # 跳过当前第i个后的scores分数取最大的
maxpos = np.argmax(scores[pos:], axis=0) # 跳过当前第i个后的scores分数取最大的分数对应的索引
else:
maxscore = scores[-1] # 如果没有（即只有一个），就取这个
maxpos = 0 # 因为只有一个，即索引为0
if tscore < maxscore: # 第一个与
# 每次把最高score的 往上拿
boxes_x1y1x2y2[i, :] = boxes_x1y1x2y2[maxpos + i + 1, :] # 行置换，score，area也一样
boxes_x1y1x2y2[maxpos + i + 1, :] = tBD
tBD = boxes_x1y1x2y2[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]
# IoU calculate
xx1 = np.maximum(boxes_x1y1x2y2[i, 1], boxes_x1y1x2y2[pos:, 1])
yy1 = np.maximum(boxes_x1y1x2y2[i, 0], boxes_x1y1x2y2[pos:, 0])
xx2 = np.minimum(boxes_x1y1x2y2[i, 3], boxes_x1y1x2y2[pos:, 3])
yy2 = np.minimum(boxes_x1y1x2y2[i, 2], boxes_x1y1x2y2[pos:, 2])
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)
# Three methods: 1.linear 2.gaussian 3.original NMS
if method == 1: # linear
weight = np.ones(ovr.shape)
weight[ovr > threshold_iou] = weight[ovr > threshold_iou] - ovr[ovr > threshold_iou]
elif method == 2: # gaussian
weight = np.exp(-(ovr * ovr) / sigma)
else: # original NMS
weight = np.ones(ovr.shape)
weight[ovr > threshold_iou] = 0 # 表示大于Nt的iou值就和比较的box重合很多，因此大于这个阈值的iou为0，则相当于该score为0
scores[pos:] = weight * scores[pos:] # 给定scores的值，若不好就删除
# select the boxes and keep the corresponding indexes
boxes_new = boxes_x1y1x2y2[scores > threshold_end] # 判断求解的scores是否满足阈值，若满足就保留
return boxes_new
# boxes and scores
boees_slect = softnms(boxes, boxscores, method=3)
print(boees_slect)
原始NMS结果如下：
线性soft-nms:
高斯nms：
