python实现NMS

python实现NMS

什么叫NMS

NMS即non maximum suppression即非极大抑制。专业的自己看链接，我就只用一句话说我的理解。
就是在多个目标框里，每个目标选择唯一最好的框（红色）。
注意：最好的这个红框可能不是最完美的，但是是可选择的里面最好的。

如何选择最好的，这里就会产生一个新的概念，叫做IoU：交并比(Intersection-over-Union,IoU)

参考博文

目标检测中NMS（非极大抑制）的概念理解

以下是我自己写的代码，实现nms功能

def func_IoU(a, b):
    '''
    lx =a[0], ly = a[1] ← 左上角
                            右下角 →  rx =a[2],ry = a[3]
    '''
    # 如果范围出错，退出计算，返回0
    if a[0] > a[2] or a[1] > a[3] or b[0] > b[2] or b[1] > b[3]:
        return 0.0

    # 计算相交面积
    w = min(a[2], b[2]) - max(a[0], b[0])
    h = min(a[3], b[3]) - max(a[1], b[1])
    if w<0 or h <0:
        return 0.0
    in_areas = w * h

    # 计算并集面积
    a_areas = (a[2] - a[0]) * (a[3] - a[1])
    b_areas = (b[2] - b[0]) * (b[3] - b[1])
    uni_areas =a_areas+b_areas-in_areas
    return in_areas/(uni_areas+1e-6)


# 1.选出得分最高的
# 2.和最高得分的面积进行iou计算，大于threshold的去除

def nms(bboxes, confidence, threshold):
    class_list=[]
    conf_list=[]
    if len(bboxes)>0:
        keep_box_list = []
        keep_conf_list= []
        max_num = confidence.index(max(confidence))
        max_box = bboxes.pop(max_num)
        max_conf = confidence.pop(max_num)
        for i in range(len(bboxes)):
            if func_IoU(max_box,bboxes[i])<=threshold:
                keep_box_list.append(bboxes[i])
                keep_conf_list.append(confidence[i])
                print(f'max_box:{max_box},max_conf:{max_conf},\ '
                      f'conf:{confidence[i]}的IoU：{func_IoU(max_box,bboxes[i])}')
        #递归几次，说明有几个prediction窗口
        class_list,conf_list =nms(keep_box_list,keep_conf_list,threshold)
        # 将返回的窗口加入到列表
        class_list.append(max_box)
        conf_list.append(max_conf)
        return class_list,conf_list
    return class_list,conf_list

if __name__ == '__main__':
    bounding=[(187,82,337,317),(150,67,305,282),(246,121,368,304)]
    confidence_score = [0.9,0.75,0.8]
    threshold = 0.6
    picked_boxes,picked_score= nms(bounding,confidence_score,threshold)
    print('最终bbox列表：',picked_boxes)
    print('最终conf分数列表：',picked_score)
    print('threshold:',threshold)

大神简单的实现方式：

import numpy as np
def nms(bboxes, confidence_score, threshold):
    """非极大抑制过程
    :param bboxes: 同类别候选框坐标
    :param confidence: 同类别候选框分数
    :param threshold: iou阈值
    :return:
    """
    # 1、传入无候选框返回空
    if len(bboxes) == 0:
        return [], []
    # 强转数组
    bboxes = np.array(bboxes)
    score = np.array(confidence_score)

    # 取出n个的极坐标点
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]

    # 2、对候选框进行NMS筛选
    # 返回的框坐标和分数
    picked_boxes = []
    picked_score = []
    # 对置信度进行排序, 获取排序后的下标序号, argsort默认从小到大排序
    order = np.argsort(score)
    areas = (x2 - x1) * (y2 - y1)
    while order.size > 0:
        # 将当前置信度最大的框加入返回值列表中
        index = order[-1]
        picked_boxes.append(bboxes[index])
        picked_score.append(confidence_score[index])

        # 获取当前置信度最大的候选框与其他任意候选框的相交面积
        x11 = np.maximum(x1[index], x1[order[:-1]])
        y11 = np.maximum(y1[index], y1[order[:-1]])
        x22 = np.minimum(x2[index], x2[order[:-1]])
        y22 = np.minimum(y2[index], y2[order[:-1]])
        w = np.maximum(0.0, x22 - x11)
        h = np.maximum(0.0, y22 - y11)
        intersection = w * h

        # 利用相交的面积和两个框自身的面积计算框的交并比, 将交并比大于阈值的框删除
        ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)
        keep_boxes_indics = np.where(ratio < threshold)
        order = order[keep_boxes_indics]
    return picked_boxes, picked_score


if __name__ == '__main__':
    bounding = [(187, 82, 337, 317), (150, 67, 305, 282), (246, 121, 368, 304)]
    confidence_score = [0.9, 0.75, 0.8]
    threshold = 0.4
    picked_boxes, picked_score = nms(bounding, confidence_score, threshold)
    print('阈值threshold为:', threshold)
    print('最终bbox列表：', picked_boxes)
    print('最终confidence分数列表：', picked_score)