非极大值抑制NMS的python实现
什么是非极大值抑制
非极大值抑制的主要目的是为了消除多余的框，找到最佳的物体检测的位置。
比如我们想要检测手的时候, RCNN网络在训练之后会给出许多个预测框(比图上的更多), 我们先通过他们的置信度筛选出一批不符合的框, 剩下如图的预测框, 还是很多, 所以通过非极大值抑制保留出个别几个最准确的框(置信度得分最高的)
非极大值抑制后留下最佳框
NMS也存在一定的问题, 它会将两个重叠在一起的物体框只保留一个, 所以柔性非极大抑制Soft-NMS在筛选框的过程中也考虑了分数的影响, 可以保留一些重叠在一起的物体
伪代码
B表示输入的预测框列表, 其中的每个bi包含了x, y, w, h的信息S表示每个预测框的预测结果得分D存储NMS之后剩下的框
伪代码运行流程argmaxS找到最大得分的下标m- 将
bm也就是最大得分对应的预测框从B中删除, 加入到D中 - 遍历
B(预测框列表) 将其中的bi与之前选出的得分最好的预测框之间计算IOU, 如果结果大于给定的门限值, 就删除当前的bi与si意思就是说如果两个框的IOU比较高, 就认为他们预测的是同一个目标, 而得分最高的框已经取出, 所以就删除当前的框 - 下面的部分是
Soft-NMS它以一个权重的形式，将获得的IOU取高斯指数后乘上原得分，之后重新排序, 也就是说不仅仅根据IOU来筛选框, 也会保留一些分数很高的框
Python实现
用python实现伪代码, 具体注释已经写在下面了
# Author : JokerTong
# Datetime : 2023-02-03 16:40
# File : NMS实现.py
import math
import numpy as np
def iou(box1, box2):
# 计算框与框之间的IOU
# x1, y1, x2, y2
# 0 1 2 3
area1 = (box1[3] - box1[1]) * (box1[2] - box1[0]) # 计算第一个框的面积
area2 = (box2[3] - box2[1]) * (box2[2] - box2[0]) # 计算第二个框的面积
inter_area = (min(box1[2], box2[2]) - max(box1[0], box2[0])) * \
(min(box1[3], box2[3]) - max(box1[1], box2[1])) # 计算两个框重叠的面积
return inter_area / area1 + area2 - inter_area # 返回IOU
def spm(iou, mode='linear', sigma=0.3):
# score penalty mechanism (soft-nms)
if mode == 'linear':
return 1 - iou
elif mode == 'gaussian':
return math.e ** (- (iou ** 2) / sigma)
else:
raise NotImplementedError
def NMS(lists, conf_thre, iou_thre, soft=True, soft_thre=0.001):
# 根据阈值先过滤一批不符合要求的框
lists = filter(lambda x: x[4] >= conf_thre, lists)
# 将所有的预测框按分数从高大到低排序
lists = sorted(lists, key=lambda x: x[4], reverse=True)
keep = []
while lists:
m = lists.pop(0) # 取出分数最高的框
keep.append(m)
for i, pred in enumerate(lists): # 遍历剩下的框, 根据iou的值来筛选
_iou = iou(m, pred)
if _iou >= iou_thre:
if soft: # 如果采用soft-NMS, 则更新score
pred[4] *= spm(_iou, mode='gaussian', sigma=0.3)
keep.append(lists.pop(i))
else:
lists.pop(i)
if soft:
keep = list(filter(lambda x: x[4] >= soft_thre, keep))
keep = sorted(keep, key=lambda x: x[4], reverse=True)
return keep
if __name__ == '__main__':
np.random.seed(0)
x1y1 = np.random.randint(0, 300, (300, 2)) / 600 # 初始化x1,y1坐标, 假设图像为(600*600)大小
x2y2 = np.random.randint(300, 600, (300, 2)) / 600 # 初始化x2,y2坐标
boxes = np.concatenate((x1y1, x2y2), 1) # 将 x1,y1,x2,y2 拼接
scores = np.random.rand(300, 1) # 初始化分数
lists = list(np.concatenate((boxes, scores), 1))
detections = NMS(lists, conf_thre=0.1, iou_thre=0.7, soft=False, soft_thre=0.1) # 非极大值抑制
print(len(detections), detections) # 输出留下的框, 以及具体的信息