NMS(非极大值抑制)算法 -- 理论、代码
目的
消除 anchor-based 算法中因为滑动窗口或者重叠采样等方式产生的众多冗余候选框，保留效果最佳的框。
定义
非极大值抑制（Non-Maximum Suppress，NMS）算法，就是抑制非极大值的目标（去冗余），从而搜索出局部极大值的目标（找最优）。
原理
输入
给定所有可能的预测边框 predictions = [ [x_max, x_min, y_max, y_min, score], [*], ..., [*]] 以及一个给定的IoU阈值 iou_threshold
有多少个类别，就执行多少次NMS算法
输出
经过NMS算法过滤后的一类物体的预测框 result = [x_max, x_min, y_max, y_min, score].
算法步骤
（1）将所有可能的预测框按类别划分为num_class+1个集合，其中1为背景类，背景类无需NMS处理；
（2）对于每一个集合（类别），按类别分数从高到低进行排序，得到num_class个降序列表 list；
（3）从一个 list 中取得第一个元素（分数最高），逐个计算该元素与列表中剩余元素的IoU，若IoU大于给定阈值则将该元素从列表中删除，同时将第一个元素保留；
（4）对处理过后的降序列表 list 重复执行步骤（3），直至 list 为空；这样返回的keep列表中就是图中该类别所有的物体的唯一框。
（5）对每一个类别都执行步骤（3）~（4），直接遍历完所有的类别；
代码（添加注释、测试）
import cv2
import numpy as np
import matplotlib.pyplot as plt
# iou = np.array([0.57184716, 0.76505679])
# order = np.array([2, 1, 4, 5, 6, 8])
# inds = np.where(iou <= 0.8)[0] # 将重叠度大于给定阈值的边框剔除掉，仅保留剩下的边框，返回相应的下标 [1,2]
# print('inds:', inds) # 得到分数小于阈值的框索引 [0 1]
# order = order[inds + 1] # 从剩余的候选框中继续筛选[1 2]
# print('order:', order) # [1 4 ]
class NMS:
def __init__(self, center=False, scale=1.0):
"""
:param center: the format of coordinate -> diagonal [x1, y1, x2, y2] or center [x, y, w, h]
"""
self.center = center
self.scale = scale
def compute_iou(self, bbox1, bbox2, eps=1e-8):
"""
compute the IoU of two bounding boxes.
:param eps: avoid none
:param bbox1: bounding box No.1.
:param bbox2: bounding box No.2.
:return: IoU of bbox1 and bbox2.
"""
if self.center:
x1, y1, w1, h1 = bbox1
xmin1, ymin1 = int(x1-w1/2.0), int(y1-h1/2.0)
xmax1, ymax1 = int(x1+w1/2.0), int(y1+h1/2.0)
x2, y2, w2, h2 = bbox2
xmin2, ymin2 = int(x2-w2/2.0), int(y2-h2/2.0)
xmax2, ymax2 = int(x2 + w2 / 2.0), int(y2 + h2 / 2.0)
else:
xmin1, ymin1, xmax1, ymax1 = bbox1
xmin2, ymin2, xmax2, ymax2 = bbox2
xx1 = np.max([xmin1, xmin2]) # 计算交集的对角坐标点
yy1 = np.max([ymin1, ymin2])
xx2 = np.min([xmax1, xmax2])
yy2 = np.min([ymax1, ymax2])
w = np.max([0.0, xx2 - xx1 + 1]) # 与 0.0 比较是为了防止两个不相交的情况
h = np.max([0.0, yy2 - yy1 + 1]) # 这里边长+1 是因为x1y1x2y2是像素点，相减+1得到距离
area_intersection = w * h # 计算交集面积
area1 = (xmax1 - xmin1 + 1) * (ymax1 - ymin1 + 1)
area2 = (xmax2 - xmin2 + 1) * (ymax2 - ymin2 + 1)
area_union = area1 + area2 - area_intersection # 计算并集面积（这里要记得去掉重叠的面积，避免重复计算）
iou = area_intersection / (area_union + eps) # 计算两个边框的交并比
return iou # eps很小，防止除数为0
@classmethod # 不用实例化就能调用的方法 可以NMS.py_cpu_nms(**)
def py_cpu_nms(cls, dets, iou_thresh=0.5, score_thresh=0.5):
"""
Pure Python NMS baseline. -> take reference from the Fast R-CNN.
:param iou_thresh: iou thresh -> default 0.5.
:param score_thresh: cls score thresh -> default 0.5.
:param dets: detection results -> [[xmin1, ymin1, xmax1, ymax1, score1], [xmin2, ymin2, xmax2, ymax2, score2], ...].
:return: optimal bounding boxes.
"""
# list的[]中有三个参数，用冒号分割 list[param1:param2:param3]
# param1，相当于start_index，可以为空，默认是0 一般[0,list.size]为左闭右开
# param2，相当于end_index，可以为空，默认是list.size
# param3，步长，默认为1。步长为-1时，返回倒序原序列
# numpy的切片操作，一般结构如num[a:b,c:d]，分析时以逗号为分隔符，
# 逗号之前为要取的num行的下标范围(a到b-1)，逗号之后为要取的num列的下标范围(c到d-1)，前面是行索引，后面是列索引。
# np.where(condition, x, y) 满足条件(condition)，输出x，不满足输出y。
# np.where(condition) 只有条件 (condition)，没有x和y，则输出满足条件元素的索引坐标
# a = np.array([2,4,6,8,10])
# np.where(a > 5) 返回a的索引 array([2,3,4])
# a[ np.where(a > 5) ] 等价于取出a[a>5] array([6,8,10])
# np.where()返回的元组为([],dtype=*) [0]取出所需要的列表 就是符合条件的dets的索引
dets = dets[ np.where(dets[:, -1] >= score_thresh) [0] ] # 只取出高值，过滤掉低于分数阈值的预测框
# print('dets:', dets) # [[ 30. 10. 200. 200. 0.95]
# [ 25. 15. 180. 220. 0.98]
# [ 35. 40. 190. 170. 0.96]]
xmin = dets[:, 0] # xmin -> [xmin1, xmin2, ...]
# print('xmin:', xmin) # [30. 25. 35.]
ymin = dets[:, 1] # ymin -> [ymin1, ymin2, ...]
xmax = dets[:, 2] # xmax -> [xmax1, xmax2, ...]
ymax = dets[:, 3] # ymax -> [ymax1, ymax2, ...]
# predict bbox class score -> [score1, score2, score3]
scores = dets[:, 4]
# print('scores:', scores) # [0.95 0.98 0.96]
# 按score降序排序（升序排序后逆序），argsort返回降序后的索引214356
order = scores.argsort()[::-1]
# print('order:', order) # [1 2 0]
areas = (xmax - xmin + 1) * (ymax - ymin + 1) # 计算面积list
# print('areas:', areas) # [32661. 32136. 20436.]
keep = [] # 保留最优的结果
# 搜索最佳边框
while order.size > 0:
top1_idx = order[0] # 选取得分最高的边框 索引
# print('top1_idx:', top1_idx) # 1
keep.append(top1_idx) # 添加到候选列表
# 将得分最高的边框与剩余边框进行比较
xx1 = np.maximum(xmin[top1_idx], xmin[order[1:]])
# print('xx1:', xx1) # [35. 30.]
# [] 与 [[],[],[],[]...[]] 比较 得到[[],[],[]...[]]
yy1 = np.maximum(ymin[top1_idx], ymin[order[1:]])
xx2 = np.minimum(xmax[top1_idx], xmax[order[1:]])
yy2 = np.minimum(ymax[top1_idx], ymax[order[1:]])
w = np.maximum(0.0, xx2 - xx1 + 1) # 计算交集 []
# print('w:', w) # [146. 151.]
h = np.maximum(0.0, yy2 - yy1 + 1)
intersection = w * h # 交集 []
# print('intersection:', intersection) # [19126. 28086.]
union = areas[top1_idx] + areas[order[1:]] - intersection # 计算并集 # []
# print('union:', union) # [33446. 36711.]
iou = intersection / union # 计算交并比 [0.5,0.6,0.7]
# print('iou:', iou) # [0.57184716 0.76505679]
# 将重叠度大于给定阈值的边框剔除掉，仅保留剩下的边框，返回相应的下标 [1,2]
inds = np.where(iou <= iou_thresh) [0]
# print('inds:', inds) # 得到分数小于阈值的框索引 []
order = order[inds + 1] # 从剩余的候选框中继续筛选
# print('order:', order) # []
return keep
if __name__ == '__main__':
img = cv2.imread("image.jpg")
img_cp = np.copy(img)
thickness = 2
info = np.array([
[30, 10, 200, 200, 0.95],
[25, 15, 180, 220, 0.98],
[35, 40, 190, 170, 0.96],
[60, 60, 90, 90, 0.3],
[20, 30, 40, 50, 0.1],
])
colors = [[0, 0, 255], [0, 255, 0], [255, 0, 0], [255, 255, 0], [0, 255, 255]]
plt.subplot(121)
plt.axis('off')
plt.title("Input image")
for i in range(len(colors)):
x1, y1, x2, y2, _, = info[i]
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)),
colors[i], thickness=thickness)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img)
plt.subplot(122)
plt.axis('off')
plt.title("After NMS")
indx = NMS.py_cpu_nms(dets=info, iou_thresh=0.5, score_thresh=0.5) # NMS
# print(indx)
for i in indx:
x1, y1, x2, y2, _ = info[i]
cv2.rectangle(img_cp, (int(x1), int(y1)), (int(x2), int(y2)), colors[i], thickness=thickness)
img_cp = cv2.cvtColor(img_cp, cv2.COLOR_BGR2RGB)
plt.imshow(img_cp)
# 保存并显示图片
plt.savefig('results.png')
plt.show()
参考：
非极大值抑制(NMS)算法讲解|理论+代码 - 知乎是