NMS算法实现

NMS算法(非极大值抑制)是目标检测算法中经典的后处理步骤,其本质是搜索局部最大值,抑制非极大值元素。主要利用目标检测框以及对应的置信度分数,设置一定的阈值来删除重叠较大的边界框。
其算法流程如下:

根据置信度得分进行排序
选择置信度最高的目标检测框添加到输出列表中,将其从检测框列表中删除
计算该检测框与剩余候选检测框的IOU
删除IOU大于阈值的检测框
重复上述4步,直至检测框列表为空

import numpy as np
#假设生成9个候选框
boxes = np.array([[100,100,210,210,0.72],
                 [250,25,420,420,0.87],
                 [220,220,320,330,0.92],
                 [120,130,210,210,0.73],
                 [230,240,325,330,0.81],
                 [220,230,315,340,0.93],
                 [300,400,450,220,0.95],
                 [200,150,500,100,0.85],
                 [350,450,500,270,0.78]
                 ])




def nms(bboxes,thresh):

    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]

    areas = (x2 - x1 +1) * (y2 - y1 +1)
       #进行排序



    scores = bboxes[:, 4]
    index = scores.argsort()[::-1]

    res = []
    while index.size > 0:
        i = index[0]
        res.append(i)

        X1 = np.maximum(x1[i],x1[index[1:]])
        Y1 = np.maximum(y1[i],y1[index[1:]])
        X2 = np.minimum(x2[i],x2[index[1:]])
        Y2 = np.minimum(y2[i],y2[index[1:]])

        w = np.maximum(0,X2 -X1 +1)
        h = np.maximum(0,Y2 - Y1 +1)

        jiao = w*h
        iou = jiao / (areas[i] + areas[index[1:]] - jiao)

        idx = np.where(iou <= thresh)[0]
        index = index[idx +1]



    #     idx =np.where(iou <= thresh)[0]
    #     index = index[idx+1]
    # return res


    return res


 

猜你喜欢

转载自blog.csdn.net/weixin_64043217/article/details/132508229
NMS
今日推荐