NMS algorithm implementation

The NMS algorithm (non-maximum suppression) is a classic post-processing step in the target detection algorithm. Its essence is to search for local maximum values ​​and suppress non-maximum elements. It mainly uses the target detection frame and the corresponding confidence score to set a certain threshold to delete bounding boxes with large overlaps.
The algorithm flow is as follows:

Sort according to the confidence score.
Select the target detection frame with the highest confidence and add it to the output list and delete it from the detection frame list. Calculate
the IOU between the detection frame and the remaining candidate detection frames.
Delete the detection frames with IOU greater than the threshold
and repeat the above 4 steps. , until the detection box list is empty

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