NMS和soft-NMS实现

前言

Soft-NMS – Improving Object Detection With One Line of Code.通过阅读这篇文章，从中了解到相比较于NMS还有改进的Soft-NMS，文章说的是可以提升在很多数据集上的检测精度。

2021-07-05 更新了fast nms的程序

无意中看到Matrix nms算法，本想实现这个算法，但是无奈没找到详细的原理，如果大家有合适的算法解析，可以私聊我分享呢，感谢。

正文

下面的计算IOU函数

def iou(predicted_bound, ground_truth_bound):       #计算IOU
    """
    computing the IoU of two boxes.
    Args:
        box: (xmin, ymin, xmax, ymax),通过左下和右上两个顶点坐标来确定矩形位置
    Return:
        IoU: IoU of box1 and box2.
    """
    pxmin, pymin, pxmax, pymax = predicted_bound
    # print("预测框P的坐标是：({}, {}, {}, {})".format(pxmin, pymin, pxmax, pymax))
    gxmin, gymin, gxmax, gymax = ground_truth_bound
    # print("原标记框G的坐标是：({}, {}, {}, {})".format(gxmin, gymin, gxmax, gymax))

    parea = (pxmax - pxmin) * (pymax - pymin)  # 计算P的面积
    garea = (gxmax - gxmin) * (gymax - gymin)  # 计算G的面积
    # print("预测框P的面积是：{}；原标记框G的面积是：{}".format(parea, garea))

    # 求相交矩形的左下和右上顶点坐标(xmin, ymin, xmax, ymax)
    xmin = max(pxmin, gxmin)  # 得到左下顶点的横坐标
    ymin = max(pymin, gymin)  # 得到左下顶点的纵坐标
    xmax = min(pxmax, gxmax)  # 得到右上顶点的横坐标
    ymax = min(pymax, gymax)  # 得到右上顶点的纵坐标

    # 计算相交矩形的面积
    w = xmax - xmin
    h = ymax - ymin
    if w <=0 or h <= 0:
        return 0

    area = w * h  # G∩P的面积
    # area = max(0, xmax - xmin) * max(0, ymax - ymin)  # 可以用一行代码算出来相交矩形的面积
    # print("G∩P的面积是：{}".format(area))

    # 并集的面积 = 两个矩形面积 - 交集面积
    IoU = area / (parea + garea - area)

    return IoU

使用NMS筛选boxes

def nms(boxes,threshold=0.5):
    boxes_copy = boxes.copy()
    result = []
    while len(boxes) > 0:
        index = [i[4] for i in boxes]
        max_ = np.argmax(index)

        max_cor = boxes[max_]
        result.append(boxes[max_])

        boxes = np.delete(boxes, max_, axis=0)
        res = []
        for i in range(len(boxes)):
            if iou(max_cor[:-1], boxes[i][:-1]) < threshold:
                res.append(boxes[i])
        boxes = res

    plt.figure()
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 2)
    plt.imshow(img)
    plt.title('nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes_copy:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()
    return result

使用soft-nms函数筛选目标框

def soft_nms(boxes,threshold,score,λ):
    boxes_copy = boxes.copy()
    result = []
    while len(boxes)>0:
        boxes = [i for i in boxes if i[-1]>score]
        if len(boxes)==0:
            break
        index = [i[-1] for i in boxes]
        max_ = np.argmax(index)
        max_cor = boxes[max_]
        boxes = np.delete(boxes, max_, axis=0)
        result.append(max_cor)
        for i in range(len(boxes)):
            if iou(max_cor[:-1],boxes[i][:-1])>threshold:
                boxes[i][-1] = boxes[i][-1]*np.exp(-1*iou(max_cor[:-1],boxes[i][:-1])/λ)

    plt.figure()
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 2)
    plt.imshow(img)
    plt.title('soft-nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes_copy:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()
    result = [[int(i[0]),int(i[1]),int(i[2]),int(i[3]),float(i[4])] for i in result]
    return result

这是fast nms的程序

def fast_nms(boxes,threshold=0.5):
    score = [i[-1] for i in boxes]
    matrix = np.zeros((len(boxes),len(boxes)))
    for i in range(len(boxes)-1):
        for j in range(i+1,len(boxes)):
            matrix[i,j] = iou(boxes[i][:-1],boxes[j][:-1])
    res1 = [max(matrix[:,i]) for i in range(len(matrix))]
    # print(res1)
    result = []
    for i in range(len(res1)):
        if res1[i]<threshold:
            result.append(boxes[i])
    print(result)
    plt.figure()
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 2)
    plt.imshow(img)
    plt.title('fast-nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0), 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()

这些是调试的程序

boxes = np.array([[100, 110, 210, 210, 0.71],
                  [250, 250, 420, 420, 0.8],
                  [220, 200, 320, 330, 0.92],
                  [120, 100, 210, 210, 0.72],
                  [230, 240, 325, 330, 0.81],
                  [220, 230, 315, 340, 0.91]])

nms(boxes,threshold=0.5)
soft_nms(boxes,0.5,0.4,0.1)    #boxes,threshold,score,λ(gauss函数参数）

得到的NMS算法的结果：

得到的soft-NMS算法的结果：

得到 Fast NMS的效果

有问题可以直接提问，欢迎一起交流。

参考文章
[1]:《Improving Object Detection With One Line of Code》