# https://blog.csdn.net/a1103688841/article/details/89711120
import numpy as np
boxes=np.array([[100,100,210,210,0.72],
[250,250,420,420,0.8],
[220,220,320,330,0.92],
[100,100,210,210,0.72],
[230,240,325,330,0.81],
[220,230,315,340,0.9]])
def py_cpu_nms(dets, thresh):
#dets的数据格式是dets[[xmin,ymin,xmax,ymax,scores]....]
x1 = dets[:,0] #取所有行的第0个数,即所有的xmin组成一个列表
y1 = dets[:,1]
x2 = dets[:,2]
y2 = dets[:,3]
areas = (y2-y1+1) * (x2-x1+1) #每个框的面积,组成列表
scores = dets[:,4] # 所有的分数
# 这边的keep用于存放,NMS后剩余的方框
keep = []
# 取出分数从大到小排列的索引。.argsort()是从小到大排列,[::-1]是列表头和尾颠倒一下
index = scores.argsort()[::-1]
# 上面这两句比如分数[0.72 0.8 0.92 0.72 0.81 0.9 ]
# 对应的索引index[ 2 5 4 1 3 0 ]记住是取出索引,scores列表没变。
# index会剔除遍历过的方框,和合并过的方框。
while index.size >0:
i = index[0] # every time the first is the biggst, and add it directly
# keep保留的是索引值,不是具体的分数。
keep.append(i)
# 计算交集的左上角和右下角
x11 = np.maximum(x1[i], x1[index[1:]]) # calculate the points of overlap
y11 = np.maximum(y1[i], y1[index[1:]])
# print(y11)
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
print(x11, y11, x22, y22)
w = np.maximum(0, x22-x11+1) # the weights of overlap
h = np.maximum(0, y22-y11+1) # the height of overlap
overlaps = w*h
ious = overlaps / (areas[i]+areas[index[1:]] - overlaps)
idx = np.where(ious<=thresh)[0]
index = index[idx+1] # because index start from 1
return keep
import matplotlib.pyplot as plt
def plot_bbox(dets, c='k'):
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")
plt.figure(1)
ax1 = plt.subplot(1,2,1)
ax2 = plt.subplot(1,2,2)
plt.sca(ax1)
plot_bbox(boxes,'k') # before nms
keep = py_cpu_nms(boxes, thresh=0.7)
plt.sca(ax2)
plot_bbox(boxes[keep], 'r')# after nms
plt.show()
Reference:
NMS的python实现