目标检测中NMS(non maximum suppression)
一、原理
参考网址
NMS即non maximum suppression即非极大抑制，顾名思义就是抑制不是极大值的元素，搜索局部的极大值。在最近几年常见的物体检测算法（包括rcnn、sppnet、fast-rcnn、faster-rcnn等）中，最终都会从一张图片中找出很多个可能是物体的矩形框，然后为每个矩形框为做类别分类概率。
Soft NMS法：对于得分小于阈值的边框，不再直接舍弃，而是降低其得分。Soft NMS是对NMS的优化算法，它在不增加额外参数的情况下且只需要对NMS算法进行简单的改动就能提高AP。该Soft-NMS算法在标准数据集PASCAL VOC2007（较R-FCN和Faster-RCNN提升1.7%）和MS-COCO（较R-FCN提升1.3%，较Faster-RCNN提升1.1%）上均有提升。对于大多数数据集而言，作用比较小，提升效果非常不明显，它起作用的地方是大量密集的同类重叠场景，大量密集的不同类重叠场景其实也没什么作用。
二、实现步骤
1、NMS算法实现
参考网址
假设我们已经有了预测的框，每个预测框对应的类别，每个预测框对应的类别得分。本文中使用的案例是n个[x1,y1,x2,y2,confident, 0,0,0,1,0,0] --> n个[左上坐标，右下坐标，置信度，种类的one-hot编码]。
（2）对于每个类别而言，取该类中计算得分最大的框与其余预测框之间的IoU。
（3）根据设定的阈值，剔除掉该类中IOU大于阈值的预测框。
（4）对于每个类别而言，在(3)的基础上循环(1)-(3)步骤，在该类中已经排好序的情况下，直接下一个最大值作为该类的得分最大的框与其余预测框之间的IoU；直到剩余有用的预测框个数为0。
2、softNMS算法实现
论文地址：https://arxiv.org/pdf/1704.04503v2.pdf
（1）对于每个类别，按照类别得分的从大到小顺序排列
（2）对于每个类别而言，取该类中计算得分最大的框与其余预测框之间的IoU。
（3）根据计算的iou，使用高斯惩罚函数公式
Si 表示置信度，当iou越大时，该预测框的置信度就会施加更重的惩罚值，从而降低置信度。
（4）对于每个类别而言，在(3)对置信度施加惩罚后，再剔除掉该类中置信度小于阈值的预测框。
（5）重复(1) - (4)步骤，直至剩余的预测框个数为0。
三、代码实现
1、Pytorch版
import numpy as np
import cv2
import matplotlib.pyplot as plt
import torch
def bbox_iou(box1, box2, x1y1x2y2=True):
# 坐标分离
if not x1y1x2y2:
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2]/2, box1[:, 0] - box1[:, 2]/2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3]/2, box1[:, 1] - box1[:, 3]/2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2]/2, box2[:, 0] - box2[:, 2]/2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3]/2, box2[:, 1] - box2[:, 3]/2
else:
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
# 计算交集的两坐标（即交集的左上、右下坐标）
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
# 计算交集面积
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, min=0) *\
torch.clamp(inter_rect_y2 - inter_rect_y1, min=0)
# 分别计算两个目标框各自的面积
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
# 交集/并集
iou = inter_area / torch.clamp(b1_area + b2_area - inter_area, min= 1e-6)
return iou
def nms(boxes, nms_thres=0.5):
result = []
boxes = torch.Tensor(boxes)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = boxes[:, -1].cpu().unique()
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = boxes[boxes[:, -1] == c]
# 按照存在物体的置信度排序
_, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
detections_class = detections_class[conf_sort_index]
# 进行非极大抑制
max_detections = []
while detections_class.size(0):
# 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
max_detections.append(detections_class[0].unsqueeze(0))
if len(detections_class) == 1:
break
ious = bbox_iou(max_detections[-1], detections_class[1:])
detections_class = detections_class[1:][ious < nms_thres]
# 堆叠
max_detections = torch.cat(max_detections).data
result.append(max_detections)
result = torch.cat(result).data
# ===================================================================================== #
# 案例绘图 #
# ===================================================================================== #
plt.figure()
colors = [(255,0,0),(0,255,0),(0,0,255)]
img = np.zeros((616, 616, 3))
plt.subplot(121)
for i in result:
cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img)
plt.title('nms')
plt.subplot(122)
img_ = np.zeros((616, 616, 3))
for i in boxes:
cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img_)
plt.title('original')
plt.show()
# ===================================================================================== #
return result # 返回nms后的结果
def soft_nms(boxes,conf_thres=0.5,sigma=0.5):
result = []
boxes = torch.Tensor(boxes)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = boxes[:, -1].cpu().unique()
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = boxes[boxes[:, -1] == c]
# 按照存在物体的置信度排序
_, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
detections_class = detections_class[conf_sort_index]
# 进行非极大抑制
max_detections = []
while detections_class.size(0):
# 取出这一类置信度最高的，一步一步往下判断，根据iou设置一个对置信度的惩罚因子，去除置信度小于conf_thres的框
max_detections.append(detections_class[0].unsqueeze(0))
if len(detections_class) == 1:
break
ious = bbox_iou(max_detections[-1], detections_class[1:])
detections_class[1:, 4] = torch.exp(-(ious * ious) / sigma) * detections_class[1:, 4]
detections_class = detections_class[1:]
detections_class = detections_class[detections_class[:, 4] >= conf_thres]
arg_sort = torch.argsort(detections_class[:, 4], descending = True)
detections_class = detections_class[arg_sort]
# 堆叠
max_detections = torch.cat(max_detections).data
result.append(max_detections)
result = torch.cat(result).data
# ===================================================================================== #
# 案例绘图 #
# ===================================================================================== #
plt.figure()
colors = [(255,0,0),(0,255,0),(0,0,255)]
img = np.zeros((616, 616, 3))
plt.subplot(121)
for i in result:
cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img)
plt.title('nms')
plt.subplot(122)
img_ = np.zeros((616, 616, 3))
for i in boxes:
cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img_)
plt.title('original')
plt.show()
# ===================================================================================== #
return result # 返回nms后的结果
if __name__ == '__main__':
boxes = np.array([[100, 110, 210, 210, 0.71, 0.7, 0.3, 0.5],
[250, 250, 420, 420, 0.8, 0.1, 0.8, 0.6],
[220, 200, 320, 330, 0.92, 0.2, 0.5, 1.0],
[120, 100, 210, 210, 0.72, 0.8, 0.2, 0.3],
[230, 240, 325, 330, 0.81, 0.1, 0.9, 0.2],
[220, 230, 315, 340, 0.91, 0.2, 0.7, 0.6]])
# 转换成[x1,y1,x2,y2,confident,种类下标]
boxes = np.hstack((boxes[...,:5],np.expand_dims(np.argmax(boxes[:,5:],axis=-1),axis=-1)))
nms(boxes,nms_thres=0.5)
soft_nms(boxes,0.5,0.1) #boxes,threshold,λ(gauss函数参数）效果图：
2、Numpy版
import numpy as np
import cv2
import matplotlib.pyplot as plt
'''
numpy版
'''
import numpy as np
def bbox_iou(box1, box2, x1y1x2y2=True):
if not x1y1x2y2:
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2]/2, box1[:, 0] - box1[:, 2]/2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3]/2, box1[:, 1] - box1[:, 3]/2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2]/2, box2[:, 0] - box2[:, 2]/2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3]/2, box2[:, 1] - box2[:, 3]/2
else:
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
inter_rect_x1 = np.maximum(b1_x1, b2_x1)
inter_rect_y1 = np.maximum(b1_y1, b2_y1)
inter_rect_x2 = np.minimum(b1_x2, b2_x2)
inter_rect_y2 = np.minimum(b1_y2, b2_y2)
inter_area = np.maximum(inter_rect_x2 - inter_rect_x1, 0) *\
np.maximum(inter_rect_y2 - inter_rect_y1, 0)
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
iou = inter_area / np.maximum(b1_area + b2_area - inter_area, 1e-6)
return iou
def nms(boxes, nms_thres=0.5):
result = []
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = np.unique(boxes[:, -1])
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = boxes[boxes[:, -1] == c]
# 按照存在物体的置信度排序
conf_sort_index = np.argsort(detections_class[:, 4])[::-1]
detections_class = detections_class[conf_sort_index]
# 进行非极大抑制
max_detections = []
while len(detections_class) != 0:
# 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
max_detections.append(np.expand_dims(detections_class[0],axis=0))
if len(detections_class) == 1:
break
ious = bbox_iou(max_detections[-1], detections_class[1:])
detections_class = detections_class[1:][ious < nms_thres]
# 堆叠
max_detections = np.concatenate(max_detections)
result.append(max_detections)
result = np.concatenate(result)
# ===================================================================================== #
# 案例绘图 #
# ===================================================================================== #
plt.figure()
colors = [(255,0,0),(0,255,0),(0,0,255)]
img = np.zeros((616, 616, 3))
plt.subplot(121)
for i in result:
cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img)
plt.title('nms')
plt.subplot(122)
img_ = np.zeros((616, 616, 3))
for i in boxes:
cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img_)
plt.title('original')
plt.show()
# ========================================================================================== #
return result # 返回nms后的结果
def soft_nms(boxes,conf_thres=0.5,sigma=0.5):
result = []
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = np.unique(boxes[:, -1])
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = boxes[boxes[:, -1] == c]
# 按照存在物体的置信度排序
conf_sort_index = np.argsort(detections_class[:, 4])[::-1]
detections_class = detections_class[conf_sort_index]
# 进行非极大抑制
max_detections = []
while len(detections_class) != 0:
# 取出这一类置信度最高的，一步一步往下判断，根据iou设置一个对置信度的惩罚因子，去除置信度小于conf_thres的框
max_detections.append(np.expand_dims(detections_class[0],axis=0))
if len(detections_class) == 1:
break
ious = bbox_iou(max_detections[-1], detections_class[1:])
detections_class[1:, 4] = np.exp(-(ious * ious) / sigma) * detections_class[1:, 4]
detections_class = detections_class[1:]
detections_class = detections_class[detections_class[:, 4] >= conf_thres]
arg_sort = np.argsort(detections_class[:, 4])[::-1]
detections_class = detections_class[arg_sort]
# 堆叠
max_detections = np.concatenate(max_detections)
result.append(max_detections)
result = np.concatenate(result)
# =================================================================================== #
# 案例绘图 #
# =================================================================================== #
plt.figure()
colors = [(255,0,0),(0,255,0),(0,0,255)]
img = np.zeros((616, 616, 3))
plt.subplot(121)
for i in result:
cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img)
plt.title('nms')
plt.subplot(122)
img_ = np.zeros((616, 616, 3))
for i in boxes:
cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
plt.imshow(img_)
plt.title('original')
plt.show()
# ========================================================================================== #
return result # 返回nms后的结果
if __name__ == '__main__':
boxes = np.array([[100, 110, 210, 210, 0.71, 0.7, 0.3, 0.5],
[250, 250, 420, 420, 0.8, 0.1, 0.8, 0.6],
[220, 200, 320, 330, 0.92, 0.2, 0.5, 1.0],
[120, 100, 210, 210, 0.72, 0.8, 0.2, 0.3],
[230, 240, 325, 330, 0.81, 0.1, 0.9, 0.2],
[220, 230, 315, 340, 0.91, 0.2, 0.7, 0.6]])
# 转换成[x1,y1,x2,y2,confident,种类下标]
boxes = np.hstack((boxes[...,:5],np.expand_dims(np.argmax(boxes[:,5:],axis=-1),axis=-1)))
nms(boxes,nms_thres=0.5)
soft_nms(boxes,0.5,0.1) #boxes,threshold,λ(gauss函数参数）效果图：
三、总结
本文章分别由算法原理、算法实现步骤和代码展示三个部分组成。如果你觉得本章论文对你有帮助，请点个，谢谢。
四、参考
https://blog.csdn.net/weixin_44791964/article/details/106222846
https://blog.csdn.net/qq_25344301/article/details/120094530
https://www.bilibili.com/video/BV1aB4y1K7za/?p=1&vd_source=bb59b9d0f88f71826610fe604fbe264e
https://blog.csdn.net/lzzzzzzm/article/details/120151155?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522167272576716800222869958%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fall.%2522%257D&request_id=167272576716800222869958&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~first_rank_ecpm_v1~rank_v31_ecpm-4-120151155-null-null.142v68control,201v4add_ask,213v2t3_control2&utm_term=yolov5%E9%9D%9E%E6%9E%81%E5%A4%A7%E5%80%BC%E6%8A%91%E5%88%B6%E7%AE%97%E6%B3%95&spm=1018.2226.3001.4187