Python实现NMS、DIOU_NMS

Python实现NMS，DIOU_NMS

NMS

import numpy as np

def NMS(dects, threshhold):
    x1 = dects[:, 0]
    y1 = dects[:, 1]
    x2 = dects[:, 2]
    y2 = dects[:, 3]
    scores = dects[:, 4]

    areas = (x2 - x1) * (y2 - y1) # 获取面积
    index = scores.argsort()[::-1]  # scores 降序排列
    keep = []

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

        xx1 = np.maximum(x1[i], x1[index[1:]])
        yy1 = np.maximum(y1[i], y1[index[1:]])
        xx2 = np.minimum(x2[i], x2[index[1:]])
        yy2 = np.minimum(y2[i], y2[index[1:]])

        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)

        inter = w * h
        ious = inter/(areas[i] + areas[index[1:]] - inter)
        ins = np.where(ious <= threshhold)[0]   # 过滤iou大于阈值的候选框
        index = index[ins + 1]  # 更新index

    print(boxes[keep])

if __name__ == "__main__":
    threshhold = 0.7
    boxes = np.array([[100, 100, 210, 210, 0.71],
                      [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]])
    NMS(boxes, threshhold)

DIOU_NMS

import numpy as np

def DIoU_NMS(dects, threshhold):
    x1 = dects[:, 0]
    y1 = dects[:, 1]
    x2 = dects[:, 2]
    y2 = dects[:, 3]
    scores = dects[:, 4]

    areas = (x2 - x1) * (y2 - y1) # 获取面积
    index = scores.argsort()[::-1]  # scores 降序排列
    keep = []

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

        # IoU
        xx1 = np.maximum(x1[i], x1[index[1:]])
        yy1 = np.maximum(y1[i], y1[index[1:]])
        xx2 = np.minimum(x2[i], x2[index[1:]])
        yy2 = np.minimum(y2[i], y2[index[1:]])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h

        # 中心点距离
        center_x1 = (x1[i]+x2[i])/2
        center_y1 = (y1[i]+y2[i])/2
        center_x2 = (x1[index[1:]]+x2[index[1:]])/2
        center_y2 = (y1[index[1:]]+y2[index[1:]])/2
        center_distance = (center_x1 - center_x2)**2 + (center_y1 - center_y2)**2

        # 对角线距离
        c_xx1 = np.minimum(x1[i], x1[index[1:]])
        c_yy1 = np.minimum(y1[i], y1[index[1:]])
        c_xx2 = np.maximum(x2[i], x2[index[1:]])
        c_yy2 = np.maximum(y2[i], y2[index[1:]])
        c_w = np.maximum(0.0, c_xx2 - c_xx1)
        c_h = np.maximum(0.0, c_yy2 - c_yy1)
        c_distance = c_w**2 + c_h**2

        ious = inter/(areas[i] + areas[index[1:]] - inter) - center_distance/c_distance
        ins = np.where(ious <= threshhold)[0]   # 过滤iou大于阈值的候选框
        index = index[ins + 1]  # 更新index

    print(boxes[keep])

if __name__ == "__main__":
    threshhold = 0.7
    boxes = np.array([[100, 100, 210, 210, 0.71],
                      [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]])
    DIoU_NMS(boxes, threshhold)