python 实现nms
用keep来保存留下的索引,按score进行排序,依次计算当前box和余下box iou,留下iou小于阈值的,遍历到没有框了为止
import numpy as np
import cv2
def nms(scores, bboxes, thresh):
order= np.argsort(scores)[::-1]
keep = []
areas = (bboxes[:, 2] - bboxes[:, 0] + 1) * (bboxes[:,3]-bboxes[:,1]+1)
while order.size>0:
keep.append(order[0])
i = order[0]
xmin,ymin = bboxes[order[1:], 0],bboxes[order[1:], 1]
xmax,ymax = bboxes[order[1:], 2],bboxes[order[1:], 3]
xmin = np.maximum(xmin, bboxes[i][0])
ymin= np.maximum(ymin, bboxes[i][1])
xmax= np.minimum(xmax, bboxes[i][2])
ymax = np.minimum(ymax, bboxes[i][3])
union = (ymax-ymin+1) * (xmax-xmin+1)
iou = union/(areas[i]+ areas[order[1:]]-union)
print(iou)
idx = np.where(iou<=thresh)[0]
order = order[idx+1]
return keep
def plot_img(boxes,score):
img = np.ones((600, 600))
for i, bbboxes in enumerate(boxes):
cv2.rectangle(img, (bbboxes[0], bbboxes[1]),(bbboxes[2], bbboxes[3]),(0,0,255), 1)
cv2.putText(img, str(i)+': '+str(score[i]), (bbboxes[0], bbboxes[1]),fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=1, color=(0,0,244))
cv2.imshow('s', img)
cv2.waitKey(0)
bboxes = np.array([[100,100,260,260], [50,20,230,230], [110,110,360,360]])
scores = np.array([0.8,0.9,0.4])
plot_img(bboxes, scores)
keep= nms(bboxes=bboxes, scores=scores, thresh=0.3)
print(keep)
输出结果:[ 0.36550872 0.16916038]
[1, 2]
[]
框1score最大 从它开始计算 它与0框1框iou分别为[ 0.36550872 0.16916038] ,由于阈值为0.3 抑制掉了框0,剩下1,2