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

---

前言

NMS的原理传送门
伪代码：

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）获取位置坐标，本代码用对角坐标表示位置

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

上述代码获取的坐标结果如下所示：(这里我后期修改了一下数值，图没改，不过只看图中大概对应关系是什么就行，数值无影响)

（2）计算每个Bbox的面积

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

（3）对Bbox的置信度得分排序

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

scores获取置信度一栏，然后按置信度用argsort()函数进行排序。但是argsort函数默认是从小到大进行排序，所以在后面用[::-1]实现倒序。（红字是倒序前的排序结果，蓝字是索引）

index中是排序后的各数组的索引。（粗体蓝字）

不过由于倒序处理，所以index中实际是[6,5,2,1,7,4,3,2,0]的形式。

(4)初始化结果集，对应伪代码中D

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

(5)选取最大值，遍历并计算iou

因为要跟所有的框计算一下iou所以需要写在一个循环里，直至所有的都计算过一遍，所以根据index的长度进行遍历：

    while index.size > 0:

5.1 读取置信度最高的框，并直接送入结果集中

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

5.2 计算其余框和最高分框的重叠面积

首先，重叠面积怎么计算？

可见第一个Bbox的右下角和第二个Bbox的左上角构成的黑色区域为重叠部分，但是怎么表示出来呢？

我们在x1中存储了[x11,x12,x13……]，分别表示第一个框的左上角横坐标，第二个框的左上角横坐标……等。Y1[y11,y12,y13……],X2[x21,x22,x23……],Y2[y21,y22,y23……]同理。

显然，重叠区域的
左上角横坐标为：左上角横坐标中大的那个
左上角纵坐标为：左上角纵坐标中小的那个
右下角横坐标为：右下角横坐标中小的那个
右下角纵坐标为：右下角纵坐标中大的那个

而Python中可以用np.maximum(X, Y) 用于逐元素比较两个array相对位置的元素大小。

代码表示重叠区域的左上角和右下角：

        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)  # 因为我们不确定框的方位，没有重叠的情况下，两个坐标是不构成区域的，所以需要用maximuw进行判断
        h = np.maximum(0, y22 - y11 + 1)  # the height of overlap
        overlaps = w * h

5.3 计算iou

前面只是得到了交集的面积，现在需要求IOU

        ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)

iou是一个列表，最高分框和其余框的iou值。其中index[1:]表示从index中下标为1的开始到最后。记住要减去一个重叠区域的面积，否组重叠区域就会重复计算两次。

5.4 按照IOU阈值删选Bbox

        idx = np.where(ious <= thresh)[0]
        index = index[idx + 1]  # because index start from 1

np.where()函数，按函数中的条件对数组进行返回。

（6）返回结果集

    return res

4.测试

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.title(" nms")


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

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

keep = nms(boxes, thresh=0.3)
plt.sca(ax2)
plot_bbox(boxes[keep], 'r')  # after nms
plt.show()

调整阈值为0.1：效果更明显了一点

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

总结

贴完整的代码

# 实现非极大值抑制
# 以单类为例
# 1.导入必要的库
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.计算每个框的面积
    areas = (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 / (areas[i] + areas[index[1:]] - overlaps) # index[1:]从下标1开始取到列表结束 最高分的面积加其余的面积

        idx = np.where(ious <= thresh)[0]
        index = index[idx + 1]  # because index start from 1

    return res





def plot_bbox(bboxs, c='k'):
    x1 = bboxs[:, 0]
    y1 = bboxs[:, 1]
    x2 = bboxs[:, 2]
    y2 = bboxs[:, 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.title(" nms")


plt.figure(1)
ax1 = plt.subplot(1, 2, 1)  # 定义子图，把NMS前和NMS后的绘制在一起
ax2 = plt.subplot(1, 2, 2)

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

res = nms(boxes, thresh=0.2)
plt.sca(ax2)
plot_bbox(boxes[res], 'r')  # after nms
plt.show()