Python 非极大值抑制(NMS)的四种实现详解

什么是非极大值抑制(NMS)？

非极大值抑制(NMS)是计算机视觉中一种常见的目标检测算法，用于多个候选框重叠的情况下从中选出最适合的候选框，即抑制掉冗余的候选框。

NMS 的原理

NMS 的原理是在所有的候选框中选出得分最高的一个 box，计算它和其他所有候选框的 IOU，将 IOU 值大于一定阈值的候选框进行抑制，最终返回筛选过后的候选框。

NMS 的输入和输出

输入:

• 待筛选的候选框boxes，一般包含 N 个候选框，每个候选框需要至少包含四个值(xmin, ymin, xmax, ymax)来描述边框，并且需要加上一个得分 confidence。

输出：

• 返回剩余的边框，这些边框应该是已经过非极大值抑制处理的框。

NMS 的四种实现方法

1. 普通for循环求解

该方法的代码实现很直观，但是在处理大量的候选框时会出现性能问题。

def nms_for(boxes, thresh):
    # boxes: 待处理的候选框，(xmin, ymin, xmax, ymax, confidence)
    if len(boxes) == 0:
        return []

    # 根据置信度对盒子进行排序（得分高的在前面）
    boxes = sorted(boxes, key=lambda x: x[4], reverse=True)

    # 找出最大得分的盒子和其余盒子
    pick = [True] * len(boxes)
    for i in range(len(boxes)):
        if not pick[i]:
            continue
        for j in range(i+1, len(boxes)):
            if not pick[j]:
                continue
            if iou(boxes[i], boxes[j]) > thresh:
                pick[j] = False

    # 返回剩余的盒子
    res_boxes = []
    for i in range(len(boxes)):
        if pick[i]:
            res_boxes.append(boxes[i])

    return res_boxes

2. 高斯加权筛选

在处理候选框时，有些候选框可能包含了较弱的信号，对最终结果产生了干扰。这时，可以考虑利用高斯加权的方式对候选框进行加权处理，在计算 IOU 值时重点考虑高斯核心区域的交并比。

import numpy as np

def gaussian_nms(boxes, sigma, thresh):
    # boxes: 待处理的候选框，(xmin, ymin, xmax, ymax, confidence)
    if len(boxes) == 0:
        return []

    # 根据置信度对盒子进行排序（得分高的在前面）
    boxes = sorted(boxes, key=lambda x: x[4], reverse=True)

    # 计算每个盒子的面积并生成对应的权重数组
    weights = np.array([box[4] for box in boxes])
    areas = np.array([(box[2]-box[0])*(box[3]-box[1]) for box in boxes])
    wghts = weights.copy()
    order = weights.argsort()[::-1]

    # 选择所有得分最大的盒子
    keeper = []
    while order.size > 0:
        idx = order[0]
        keeper.append(idx)

        # 计算该盒子与其他盒子的 IOU
        xx1 = np.maximum(boxes[idx][0], boxes[order[1:]][:, 0])
        yy1 = np.maximum(boxes[idx][1], boxes[order[1:]][:, 1])
        xx2 = np.minimum(boxes[idx][2], boxes[order[1:]][:, 2])
        yy2 = np.minimum(boxes[idx][3], boxes[order[1:]][:, 3])
        w = np.maximum(0.0, xx2-xx1+1) * np.maximum(0.0, yy2-yy1+1)
        inter_area = w / (areas[idx] + areas[order[1:]] - w)

        # 对欲保留的盒子进行筛选
        idxs = np.where(inter_area <= thresh)[0]
        order = order[idxs + 1]

    # 返回剩余的盒子
    res_boxes = [boxes[idx] for idx in keeper]
    return res_boxes

3. 中心点筛选

该算法的实现思路是将每个候选框都表示成中心点、高和宽的形式，然后用中心点之间的距离作为权重在进行 IOU 的计算，所形成的哪些框可以被删去，哪些需要保留。

def center_nms(boxes, thresh):
    # boxes: 待处理的候选框，(xmin, ymin, xmax, ymax, confidence)
    if len(boxes) == 0:
        return []

    # 根据置信度对盒子进行排序（得分高的在前面）
    boxes = sorted(boxes, key=lambda x: x[4], reverse=True)

    # 初始化已经保留的边框列表
    res_boxes = []
    while len(boxes) > 0:
        # 选出置信度最大的边框并从列表中剔除
        res_boxes.append(boxes[0])
        boxes.pop(0)

        # 在剩余的边框中删除和已选择的边框 IOU 值大于阈值的边框
        for i in range(len(boxes)):
            if iou(res_boxes[-1], boxes[i]) > thresh:
                boxes.pop(i)

    return res_boxes

4. Soft-NMS

在非极大值抑制出现之前，常见的做法是使用置信度（Confidence）作为指标进行筛选。然而在实际应用过程中，最高置信度的框并不一定是最优的，因此需要对 IOU 进行加权，这就是 Soft-NMS 的核心思路。

def soft_nms(boxes, thresh, sigma=0.5, score_thresh=0.001, method="linear"):
    # boxes: 待处理的候选框，(xmin, ymin, xmax, ymax, confidence)
    if len(boxes) == 0:
        return []

    # 根据置信度对盒子进行排序（得分高的在前面）
    boxes = sorted(boxes, key=lambda x: x[4], reverse=True)
    scores = np.array([box[4] for box in boxes])
    boxes = np.array(boxes)

    # 初始化结果列表
    res_boxes = []
    while boxes.shape[0] > 0:
        # 选出当前所有框中得分最高的框以其为基准,记录该框信息并将其从列表中移除
        res_boxes.append(boxes[0].tolist())
        boxes = boxes[1:]
        scores = scores[1:]

        # 计算选出的框与剩余框之间的 IOU，并进行软抑制
        if boxes.shape[0] == 0:
            break
        ious = iou(res_boxes[-1], boxes)
        if method == "linear":
            weights = np.where(ious > thresh, scores* (1-ious), scores)
        else:
            weights = scores * np.exp(-1*ious*ious/sigma)

        # 赋予框更强的惩罚，然后施以非极大值抑制操作
        mask = weights > score_thresh
        boxes = boxes[mask]
        scores = weights[mask]

    return res_boxes

示例说明

示例一：普通for循环实现

# 示例一：普通for循环实现
boxes = [
    [100, 100, 210, 210, 0.72],
    [105, 105, 200, 200, 0.8],
    [120, 120, 220, 230, 0.92],
    [140, 140, 250, 240, 0.68],
    [160, 160, 260, 250, 0.7],
    [170, 170, 245, 255, 0.9]
]

res_boxes = nms_for(boxes, 0.7)
print(res_boxes)

输出结果：

[[120, 120, 220, 230, 0.92], 
 [170, 170, 245, 255, 0.9], 
 [160, 160, 260, 250, 0.7]]

示例二：高斯加权筛选算法

# 示例二：高斯加权筛选算法
boxes = [
    [100, 100, 210, 210, 0.72],
    [105, 105, 200, 200, 0.8],
    [120, 120, 220, 230, 0.92],
    [140, 140, 250, 240, 0.68],
    [160, 160, 260, 250, 0.7],
    [170, 170, 245, 255, 0.9]
]

res_boxes = gaussian_nms(boxes, sigma=1, thresh=0.7)
print(res_boxes)

输出结果：

[[120, 120, 220, 230, 0.92], 
 [170, 170, 245, 255, 0.9], 
 [160, 160, 260, 250, 0.7]]

以上就是常用的 NMS 的四个实现方法。需要根据具体问题，选择最合适的方案，可在不同场景下使用不同的 NMS 方法来得到更加准确及高效的结果。