非极大值抑制NMS代码实现（基于Python）

前言

给出一张图片和上面许多物体检测的候选框（即每个框可能都代表某种物体），但是这些框很可能有互相重叠的部分，我们要做的就是只保留最优的框。假设有N个框，每个框被分类器计算得到的分数为Si, 1<=i<=N。

（1）建造一个存放待处理候选框的集合H，初始化为包含全部N个框；建造一个存放最优框的集合M，初始化为空集。
（2）将所有集合 H 中的框进行排序，选出分数最高的框 m，从集合 H 移到集合 M；
（3）遍历集合 H 中的框，分别与框m计算交并比（Interection-over-union，IoU），如果高于某个阈值（一般为0~0.5），则认为此框与 m重叠，将此框从集合 H 中去除。
（4）回到第2步进行迭代，直到集合 H 为空。集合 M 中的框为我们所需。

NMS代码实现

1.导入必要的库

# 1.导入必要的库
import numpy as np
from matplotlib import pyplot as plt

2.人为生成一组位置坐标，模拟候选框

# 2.假设生成如下Bbox
boxes = np.array([[100, 100, 210, 210, 0.72],
                  [250, 250, 420, 420, 0.87],
                  [220, 220, 320, 330, 0.92],
                  [120, 130, 210, 210, 0.73],
                  [230, 240, 325, 330, 0.81],
                  [220, 230, 315, 340, 0.93],
                  [300, 400, 450, 200, 0.95],
                  [200, 150, 500, 100, 0.85],
                  [350, 450, 500, 270, 0.78]])

3.定义NMS

def nms(bboxs, thresh):
    # bboxs:形似上面设置的boxes，是一组包含了诸多框坐标的数组
    # thresh： IOU阈值

    # 1.获取左上角右下角四个坐标
    x1 = bboxs[:, 0]                          # 获取所有框的左上角横坐标
    y1 = bboxs[:, 1]                          # 获取所有框的左上角纵坐标
    x2 = bboxs[:, 2]                          # 获取所有框的右下角横坐标
    y2 = bboxs[:, 3]                          # 获取所有框的右下角纵坐标

    # 2.计算每个框的面积
    area = (y2 - y1 + 1) * (x2 - x1 + 1)

    # 3.获取得分以排序
    scores = bboxs[:, 4];
    index = scores.argsort()[::-1]  # argsort默认从小到大排序，[::-1]实现翻转

    # 4.保留结果集，返回输出保留下来的Bbox最终结果
    res = []

    while index.size > 0:
        i = index[0]   # index中存储Bbox按分排序后的索引，所以第一个就是得分最高的Bbox索引，直接保留
        res.append(i)

        x11 = np.maximum(x1[i], x1[index[1:]])  # 用X11表示重叠区域的左上角横坐标
        y11 = np.maximum(y1[i], y1[index[1:]])  # 用y11表示重叠区域的左上角横坐标
        x22 = np.minimum(x2[i], x2[index[1:]])  # 用X22表示重叠区域的左上角横坐标
        y22 = np.minimum(y2[i], y2[index[1:]])  # 用y221表示重叠区域的左上角横坐标

        w = np.maximum(0, x22 - x11 + 1)  # the weights of overlap
        h = np.maximum(0, y22 - y11 + 1)  # the height of overlap

        overlaps = w * h
        ious = overlaps / (area[i] + area[index[1:]] - overlaps)
        idx = np.where(ious <= thresh)[0]
        index = index[idx + 1]

    return res

完整代码：

# -*- coding: utf-8 -*-
# @Time    : 2021/10/4 14:55
# @Author  : jhys
# @FileName: nms_test.py

import numpy as np
from matplotlib import pyplot as plt

# 2.假设生成如下Bbox
boxes = np.array([[100, 100, 210, 210, 0.72],
                  [250, 250, 420, 420, 0.87],
                  [220, 220, 320, 330, 0.92],
                  [120, 130, 210, 210, 0.73],
                  [230, 240, 325, 330, 0.81],
                  [220, 230, 315, 340, 0.93],
                  [300, 400, 450, 200, 0.95],
                  [200, 150, 500, 100, 0.85],
                  [350, 450, 500, 270, 0.78]])

def nms(bboxs, thresh):
    # bboxs:形似上面设置的boxes，是一组包含了诸多框坐标的数组
    # thresh： IOU阈值

    # 1.获取左上角右下角四个坐标
    x1 = bboxs[:, 0]                          # 获取所有框的左上角横坐标
    y1 = bboxs[:, 1]                          # 获取所有框的左上角纵坐标
    x2 = bboxs[:, 2]                          # 获取所有框的右下角横坐标
    y2 = bboxs[:, 3]                          # 获取所有框的右下角纵坐标

    # 2.计算每个框的面积
    area = (y2 - y1 + 1) * (x2 - x1 + 1)

    # 3.获取得分以排序
    scores = bboxs[:, 4];
    index = scores.argsort()[::-1]  # argsort默认从小到大排序，[::-1]实现翻转

    # 4.保留结果集，返回输出保留下来的Bbox最终结果
    res = []

    while index.size > 0:
        i = index[0]   # index中存储Bbox按分排序后的索引，所以第一个就是得分最高的Bbox索引，直接保留
        res.append(i)

        x11 = np.maximum(x1[i], x1[index[1:]])  # 用X11表示重叠区域的左上角横坐标
        y11 = np.maximum(y1[i], y1[index[1:]])  # 用y11表示重叠区域的左上角横坐标
        x22 = np.minimum(x2[i], x2[index[1:]])  # 用X22表示重叠区域的左上角横坐标
        y22 = np.minimum(y2[i], y2[index[1:]])  # 用y221表示重叠区域的左上角横坐标

        w = np.maximum(0, x22 - x11 + 1)  # the weights of overlap
        h = np.maximum(0, y22 - y11 + 1)  # the height of overlap

        overlaps = w * h
        ious = overlaps / (area[i] + area[index[1:]] - overlaps)
        idx = np.where(ious <= thresh)[0]
        index = index[idx + 1]

    return res

def plot_bbox(dets, c='k'):

    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]

    plt.plot([x1, x2], [y1, y1], c)
    plt.plot([x1, x1], [y1, y2], c)
    plt.plot([x1, x2], [y2, y2], c)
    plt.plot([x2, x2], [y1, y2], c)


plt.figure(1)
ax1 = plt.subplot(1, 2, 1)
ax2 = plt.subplot(1, 2, 2)

plt.sca(ax1)
plt.title("before nms")
plot_bbox(boxes, 'k')  # before nms

keep = nms(boxes, thresh=0.1)
print("keep: %s" %keep)
plt.sca(ax2)
plt.title("nms")
plot_bbox(boxes[keep], 'r')

plt.show()

 测试（thresh = 0.1）：

 调整阈值为0.9：相当于没进行NMS：