原始的NMS参考这篇博客。
def cpu_nms(dets, thresh):
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1) # 计算每个框的面积
scores = dets[:, 4] # 取置信度
#print('areas: ', areas)
#print('scores: ', scores)
# keep用于存放NMS后剩余的方框
keep = []
# 取出置信度从大到小排列的索引。 argsort() 是从小到大的排列， [::-1]是列表头和尾颠倒一下
index = scores.argsort()[::-1]
#print(index)
# index 会删除遍历过的方框，和合并过的方框
while index.size > 0:
#print("index size: ", index.size)
# 取出第一个方框进行和其他方框比对，看有没有可以合并的
i = index[0]
#print(i)
# 因为我们这边分数已经从大到小排序了
# 所以如果有合并存在，也是保留分数最高的这个。
# keep 保留的是索引值，不是具体的分数。
keep.append(i)
#print(keep)
#print('x1', x1[i])
#print(x1[index[1:]])
# 计算交集的左上角和右下角
# 这里要注意，比如x1[i]这个方框的左上角x和所有其他的方框的左上角x
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
#print(x11, y11, x22, y22)
# 这边要注意，如果两个方框相交，x22-x11和y22-y11是正的。
# 如果两个方框不相交，x22-x11 和 y22-y11是负的，我们把不相交的w和h设为0
w = np.maximum(0, x22-x11+1)
h = np.maximum(0, y22-y11+1)
# 计算重叠面积，不相交因为w和h都是0，所以不相交面积为0
overlaps = w*h
#print('overlap is: ', overlaps)
# 这个就是IoU公式
# 得出来的IoU是一个列表，里面拥有当前方框和其他所有方框的IoU结果。
ious = overlaps / (areas[i]+areas[index[1:]] - overlaps)
#print('ious is: ', ious)
# 接下来是合并重叠度最大的方框，也就是合并ious中值大于thresh的方框，
# 我们合并的操作就是把它们提出，因为我们合并这些方框只保留下分数最高的。
# 我们经过排序当前我们操作的方框就是分数最高的，所以我们剔除其他和当前重叠度最高的方框
# 这里np.where(ious<=thresh)[0]是一个固定的写法
idx = np.where(ious<=thresh)[0]
#print("idx: ", idx)
# 把留下来的框进行NMS操作
# 这边留下来的框是去除当前操作的框，和当前操作的框重叠度大于thresh的框
# 每一次都会现去除当前操作框，所以索引的列表就会向前移动移位，要还原就+1，向后移动一位
index = index[idx+1]
#print("index: ", index)
return keep
boxes = np.array([[100, 100, 210, 210, 0.72],
[250, 250, 420, 420, 0.80],
[220, 220, 320, 330, 0.92],
[120, 120, 210, 210, 0.72],
[230, 240, 325, 330, 0.81],
[220, 230, 315, 330, 0.90],
[50, 50, 150, 150, 0.60]
])
plt.figure(1)
ax1 = plt.subplot(1, 2, 1)
ax2 = plt.subplot(1, 2, 2)
plt.sca(ax1)
plot_bbox(boxes, 'k', title='original')
plt.sca(ax2)
keep = cpu_nms(boxes, 0.7)
plot_bbox(boxes[keep], 'r', title='NMS')
plt.show()
Soft-NMS 论文阅读笔记可以参考这篇博客。
Soft-NMS代码出自于:https://github.com/DocF/Soft-NMS/blob/master/soft_nms.py
先来看一下Soft-NMS的整体流程:
代码如下:
import numpy as np
import matplotlib.pyplot as plt
def plot_bbox(dets, c='k', title='None'):
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
plt.plot([x1, x2], [y1, y1], c)
plt.plot([x1, x1], [y1, y2], c)
plt.plot([x1, x2], [y2, y2], c)
plt.plot([x2, x2], [y1, y2], c)
plt.title("nms")
def cpu_softnms(dets, iou_thresh, gamma=0.5, thresh=0.001, method=2):
N = dets.shape[0] # the size of bboxes
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
scores = dets[:, 4]
for i in range(N):
temp_box = dets[i, :4]
temp_score = scores[i]
temp_area = areas[i]
pos = i + 1
# m <---- argmax S
if i != N - 1:
maxscore = np.max(scores[pos:])
maxpos = np.argmax(scores[pos:])
else:
maxscore = scores[-1]
maxpos = -1
#
if temp_score < maxscore:
dets[i, :4] = dets[maxpos + i + 1, :4] # M <--- b_m
dets[maxpos + i + 1, :4] = temp_box # swap position
#temp_box = dets[i, :4]
scores[i] = scores[maxpos + i + 1]
scores[maxpos + i + 1] = temp_score # swap position
#temp_score = scores[i]
areas[i] = areas[maxpos + i + 1]
areas[maxpos + i + 1] = temp_area # swap position
#temp_area = areas[i]
# calculate IoU iou(M, b_i)
xx1 = np.maximum(x1[i], x1[pos:])
xx2 = np.minimum(x2[i], x2[pos:])
yy1 = np.maximum(y1[i], y1[pos:])
yy2 = np.minimum(y2[i], y2[pos:])
w = np.maximum(xx2 - xx1 + 1.0, 0.)
h = np.maximum(yy2 - yy1 + 1.0, 0.)
inters = w * h
ious = inters / (areas[i] + areas[pos:] - inters)
# Three methods, linear, gaussian, original
# f(iou(M, b_i))
if method == 1:
weight = np.ones(ious.shape)
weight[ious > iou_thresh] = weight[ious > iou_thresh] - ious[ious > iou_thresh]
elif method == 2:
weight = np.exp(-ious * ious / sigma)
else:
weight = np.ones(ious.shape)
weight[ious > iou_thresh] = 0
scores[pos:] = scores[pos:] * weight # s_i <---- s_i * f(iou(M, b_i))
inds = np.argwhere(scores >= thresh)
keep = inds.reshape(1,inds.shape[0])[0]
return keep
boxes = np.array([[100, 100, 210, 210, 0.72],
[250, 250, 420, 420, 0.80],
[220, 220, 320, 330, 0.92],
[120, 120, 210, 210, 0.72],
[230, 240, 325, 330, 0.81],
[220, 230, 315, 330, 0.90],
[50, 50, 150, 150, 0.60]
])
plt.figure(1)
ax1 = plt.subplot(1, 2, 1)
ax2 = plt.subplot(1, 2, 2)
plt.sca(ax1)
plot_bbox(boxes, 'k', title='original')
plt.sca(ax2)
keep = cpu_softnms(boxes, iou_thresh=0.9, thresh=0.001, method=2)
plot_bbox(boxes[keep], 'r', title='soft-NMS')
plt.show()
值得注意的是，虽然没有直接用到NMS的threshold，但是在最后引入了一个置信度，置信度的阈值也决定了soft-NMS性能。
Reference- https://blog.csdn.net/a1103688841/article/details/89711120
