非极大值抑制NMS的python实现

什么是非极大值抑制

非极大值抑制的主要目的是为了消除多余的框，找到最佳的物体检测的位置。

比如我们想要检测手的时候, RCNN网络在训练之后会给出许多个预测框(比图上的更多), 我们先通过他们的置信度筛选出一批不符合的框, 剩下如图的预测框, 还是很多, 所以通过非极大值抑制保留出个别几个最准确的框(置信度得分最高的)
在这里插入图片描述
非极大值抑制后留下最佳框

NMS也存在一定的问题, 它会将两个重叠在一起的物体框只保留一个, 所以柔性非极大抑制Soft-NMS在筛选框的过程中也考虑了分数的影响, 可以保留一些重叠在一起的物体

伪代码

B 表示输入的预测框列表, 其中的每个bi包含了x, y, w, h的信息
S 表示每个预测框的预测结果得分
D 存储NMS之后剩下的框
伪代码运行流程
argmaxS 找到最大得分的下标m
将bm 也就是最大得分对应的预测框从B中删除, 加入到D中
遍历B(预测框列表) 将其中的bi与之前选出的得分最好的预测框之间计算IOU, 如果结果大于给定的门限值, 就删除当前的bi与si 意思就是说如果两个框的IOU比较高, 就认为他们预测的是同一个目标, 而得分最高的框已经取出, 所以就删除当前的框
下面的部分是Soft-NMS 它以一个权重的形式，将获得的IOU取高斯指数后乘上原得分，之后重新排序, 也就是说不仅仅根据IOU来筛选框, 也会保留一些分数很高的框

Python实现

用python实现伪代码, 具体注释已经写在下面了

# Author    : JokerTong
# Datetime  : 2023-02-03 16:40
# File      : NMS实现.py
import math
import numpy as np


def iou(box1, box2):
    # 计算框与框之间的IOU
    # x1, y1, x2, y2
    # 0   1   2   3
    area1 = (box1[3] - box1[1]) * (box1[2] - box1[0])  # 计算第一个框的面积
    area2 = (box2[3] - box2[1]) * (box2[2] - box2[0])  # 计算第二个框的面积
    inter_area = (min(box1[2], box2[2]) - max(box1[0], box2[0])) * \
                 (min(box1[3], box2[3]) - max(box1[1], box2[1]))  # 计算两个框重叠的面积
    return inter_area / area1 + area2 - inter_area  # 返回IOU


def spm(iou, mode='linear', sigma=0.3):
    # score penalty mechanism (soft-nms)
    if mode == 'linear':
        return 1 - iou
    elif mode == 'gaussian':
        return math.e ** (- (iou ** 2) / sigma)
    else:
        raise NotImplementedError


def NMS(lists, conf_thre, iou_thre, soft=True, soft_thre=0.001):
    # 根据阈值先过滤一批不符合要求的框
    lists = filter(lambda x: x[4] >= conf_thre, lists)
    # 将所有的预测框按分数从高大到低排序
    lists = sorted(lists, key=lambda x: x[4], reverse=True)
    keep = []

    while lists:
        m = lists.pop(0)  # 取出分数最高的框
        keep.append(m)
        for i, pred in enumerate(lists):  # 遍历剩下的框, 根据iou的值来筛选
            _iou = iou(m, pred)
            if _iou >= iou_thre:
                if soft:  # 如果采用soft-NMS, 则更新score
                    pred[4] *= spm(_iou, mode='gaussian', sigma=0.3)
                    keep.append(lists.pop(i))
                else:
                    lists.pop(i)

    if soft:
        keep = list(filter(lambda x: x[4] >= soft_thre, keep))
        keep = sorted(keep, key=lambda x: x[4], reverse=True)

    return keep


if __name__ == '__main__':
    np.random.seed(0)
    x1y1 = np.random.randint(0, 300, (300, 2)) / 600  # 初始化x1,y1坐标, 假设图像为(600*600)大小
    x2y2 = np.random.randint(300, 600, (300, 2)) / 600  # 初始化x2,y2坐标
    boxes = np.concatenate((x1y1, x2y2), 1)  # 将 x1,y1,x2,y2 拼接
    scores = np.random.rand(300, 1)  # 初始化分数

    lists = list(np.concatenate((boxes, scores), 1))
    detections = NMS(lists, conf_thre=0.1, iou_thre=0.7, soft=False, soft_thre=0.1)  # 非极大值抑制
    print(len(detections), detections)  # 输出留下的框, 以及具体的信息