nms

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