NMS的Python实现

import numpy as np
def nms(dets, thresh):#dets：[[xmin,ymin,xmax,ymax,score],[],[]...]
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    areas = (y2 - y1 + 1) * (x2 - x1 + 1)
    scores = dets[:,4]
    
    keep = []#存放nms后的框
    index = scores.argsort()[::-1]#取出分数从大到小排列的索引。.argsort()是从小到大排列，[::-1]是列表头和尾颠倒一下。
    
    while index.size > 0:
        i = index[0]#取出第一个框
        keep.append(i)
        
        #计算交集左上角右下角:该框和其余所有框
        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:]])
        #不相交时都为0
        w = np.maximum(0, x22-x11 + 1)
        h = np.maximum(0, y22-y11 + 1)
        
        overlaps = w * h
        ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
        #print(ious)
        idx = np.where(ious <= thresh)[0]
        #print(np.where(ious <= thresh))[0]就是取出来为true的在ious里的下标
        
        index = index[idx + 1]#ious里面从index1开始的，所以要+1
    return keep

    
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]]) 

keep = nms(boxes, thresh=0.7)
keep
    
    ```