非极大值抑制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()
标签：index,NMS,Python,横坐标,极大值,Bbox,np,左上角,bboxs 来源： https://blog.csdn.net/weixin_41761357/article/details/112665846