ICode9

精准搜索请尝试: 精确搜索
首页 > 编程语言> 文章详细

非极大值抑制NMS代码实现(Python)

2021-01-15 15:57:16  阅读:1043  来源: 互联网

标签:index NMS Python 横坐标 极大值 Bbox np 左上角 bboxs


文章目录


前言

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()





标签:index,NMS,Python,横坐标,极大值,Bbox,np,左上角,bboxs
来源: https://blog.csdn.net/weixin_41761357/article/details/112665846

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有