MTCNN中非极大值抑制NMS代码

def nms(boxes,thresh=0.3,isMin=False):#[x1,y1,x2,y2,c]
    #根据置信度排序
    _boxes = boxes[(-boxes[:,4]).argsort()]
    #保留剩余的框
    r_boxes = []
    while _boxes.shape[0]>1:
        #取出第一个框
        a_box = _boxes[0]
        #取出剩余的框
        b_boxes = _boxes[1:]
        #保留第一个框
        r_boxes.append(a_box)

        #比较IOU后保留阈值小的值
        index = np.where(iou(a_box,b_boxes,isMin)<thresh)
        _boxes = b_boxes[index]
    if _boxes.shape[0]>0:
        r_boxes.append(_boxes[0])
    #stack组装为矩阵
    return np.stack(r_boxes)

if __name__ == '__main__':
    # bs = np.array([[2,2,30,30,40],[3,3,25,25,60],[18,18,27,27,15]])
    # print(nms(bs))
    a=np.array([2,2,30,30,40])
    b=np.array([[3,3,25,25,60],[18,18,27,27,15]])
    print(iou(a,b,isMin=False))