# 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实现