如有错误，恳请指出。

用这篇博客记录一下nms，也就是非极大值抑制处理，算是目标检测后处理的一个难点。

在训练阶段是不需要nms处理的，只有在验证或者是测试阶段才需要将预测结果进行非极大值抑制处理，来挑选最佳的正样本。下面就详细查看一下非极大值抑制处理算法的一个大致流程。

文章目录

1. NMS主要步骤
2. NMS代码实现
3. NMS的变体与实现
- 3.1 hard_nms_batch
- 3.2 hard_nms
- 3.3 and-nms
- 3.4 merge-nms
- 3.5 soft-nms
- 3.6 iou-nms
- 3.7 diou_nms
4. NMS变体代码完整展示

1. NMS主要步骤

下面是yolov3spp中，nms算法的主要步骤：

1）剔除置信度较低的背景目标

x = x[x[:, 4] > conf_thres]

2）剔除宽高较小或者较大的目标

min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
x = x[((x[:, 2:4] > min_wh) & (x[:, 2:4] < max_wh)).all(1)]

3）剔除类别概率较低的目标

这里可以选择对每一个类别都进行检测其概率，挑选出那些大于阈值的正样本，返回预测样本索引与每个样本的满足条件的阈值索引。也可以对概率最大的类别进行筛选，这种比较直观。就是当概率最大的类别满足阈值条件，说明该样本可以符合筛选条件，也就不需要在乎该正样本其他概率较小的类别

if multi_label:  # 针对每个类别执行非极大值抑制# (x[:, 5:] > conf_thres).nonzero(as_tuple=False): torch.Size([nums, 2])# 返回是满足条件的位置，每个nums表示第几行第几列非0# i: 满足条件的x# j: 满足条件的x所在位置索引i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()# 由于这里的j是直接根据20类来进行索引，但是在x中是维度是25，前5个维度是坐标与置信度信息，所以需要加5来进行偏移# 然后将边界框信息(xyxy), 置信度，索引，重新整合成新xx = torch.cat((box[i], x[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)
else:  # best class only  直接针对每个类别中概率最大的类别进行非极大值抑制处理conf, j = x[:, 5:].max(1)x = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]

4）检测数据是否为有限数

# Apply finite constraint: 检测数据是否为有限数
if not torch.isfinite(x).all():x = x[torch.isfinite(x).all(1)]

5）剩余的数据进行nms，保留前max_num个数据

这个部分验证可能有用，测试基本没用。因为测试阶段一般不会有超过max_num个数据

# Batched NMS: x[box(xyxy), conf, max_index]
c = x[:, 5] * 0 if agnostic else x[:, 5]  # classes
boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4]  # boxes (offset by class), scores
# 非极大值抑制处理，返回筛选后的索引
i = torchvision.ops.nms(boxes, scores, iou_thres)
i = i[:max_num]  # 最多只保留前max_num个目标信息
# 获取最后的筛选结果
output[xi] = x[i]

这里需要注意的点是，对于torchvision.ops.nms函数来说，传入的是预测的边界框，置信度以及iou阈值，这里是和真实边界框ground true是没有任何关系的，而返回是就是筛选出的索引。

而且，这里的边界框信息并不是普通的预测出来的边界框x1y1x2y2，其还需要加上一个预测类别*最大宽高的这么一个数值，所以最后得到的boxes是比较大的。这一点我不是很理解，为什么需要再乘上一个比边界框大好几倍的数值？不过这里函数封装成了库文件，没有办法查看代码了，这里贴上nms函数的介绍：

def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:"""Performs non-maximum suppression (NMS) on the boxes accordingto their intersection-over-union (IoU).NMS iteratively removes lower scoring boxes which have anIoU greater than iou_threshold with another (higher scoring)box.If multiple boxes have the exact same score and satisfy the IoUcriterion with respect to a reference box, the selected box isnot guaranteed to be the same between CPU and GPU. This is similarto the behavior of argsort in PyTorch when repeated values are present.Args:boxes (Tensor[N, 4])): boxes to perform NMS on. Theyare expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and``0 <= y1 < y2``.scores (Tensor[N]): scores for each one of the boxesiou_threshold (float): discards all overlapping boxes with IoU > iou_thresholdReturns:Tensor: int64 tensor with the indices of the elements that have been keptby NMS, sorted in decreasing order of scores"""_assert_has_ops()return torch.ops.torchvision.nms(boxes, scores, iou_threshold)

看了函数介绍，感觉正常来说就是普通传入边界框信息与置信度信息就可以了，所以这里我尝试对其进行了更改，更改为也更加直观的版本：

# 源码的例子：
# Batched NMS: x[box(xyxy), conf, max_index]
# c = x[:, 5] * 0 if agnostic else x[:, 5]  # classes
# boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4]  # boxes (offset by class), scores# 我更改的例子：
# 非极大值抑制处理，根据边界框与置信度来进行筛选，返回筛选后的索引
boxes, scores = x[:, :4], x[:, 4]
i = torchvision.ops.nms(boxes, scores, iou_thres)
i = i[:max_num]  # 最多只保留前max_num个目标信息

下面分别查看更改前后，对图片的测试结果：

更改前：
更改后：

感觉更改前后效果好像没有什么不同，如果有知道为什么需要加上这么一个大的数值的朋友可以告诉我一下，谢谢。

2. NMS代码实现

如上所说，这里我作了一点小小的更改以更加的直观且容易理解，注释得比较详细，可以直接观看。

YOLOv3-SPP代码：

# 筛选以及非极大值抑制处理
def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6,multi_label=True, classes=None, agnostic=False, max_num=100):"""Performs  Non-Maximum Suppression on inference resultsparam: prediction: [batch, num_anchors(3个yolo预测层), (x+y+w+h+1+num_classes)]  3个anchor的预测结果总和conf_thres: 先进行一轮筛选，将分数过低的预测框（<conf_thres）删除（分数置0）nms_thres: iou阈值, 如果其余预测框与target的iou>iou_thres, 就将那个预测框置0multi_label: 是否是多标签max_num：筛选的最大数目Returns detections with shape:(x1, y1, x2, y2, object_conf, class)"""# Settingsmerge = True  # merge for best mAPmin_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and heighttime_limit = 10.0  # seconds to quit aftert = time.time()nc = prediction[0].shape[1] - 5  # number of classesmulti_label &= nc > 1  # multiple labels per boxoutput = [None] * prediction.shape[0]for xi, x in enumerate(prediction):  # image index, image inference 遍历每张图片# Apply constraints# torch.Size([5040, 25]) -> torch.Size([3051, 25])x = x[x[:, 4] > conf_thres]  # confidence 根据obj confidence虑除背景目标# torch.Size([3051, 25]) -> torch.Size([3051, 25])x = x[((x[:, 2:4] > min_wh) & (x[:, 2:4] < max_wh)).all(1)]  # width-height 虑除小目标# If none remain process next imageif not x.shape[0]:continue# Compute conf: x[:, xywh + conf + cls_score]x[..., 5:] *= x[..., 4:5]  # conf = obj_conf * cls_conf# Box (center x, center y, width, height) to (x1, y1, x2, y2)box = xywh2xyxy(x[:, :4])# Detections matrix nx6 (xyxy, conf, cls)if multi_label:  # 针对每个类别执行非极大值抑制# (x[:, 5:] > conf_thres).nonzero(as_tuple=False): torch.Size([nums, 2])# 返回是满足条件的位置，每个nums表示第几行第几列非0# i: 满足条件的x# j: 满足条件的x所在位置索引i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()# 由于这里的j是直接根据20类来进行索引，但是在x中是维度是25，前5个维度是坐标与置信度信息，所以需要加5来进行偏移# 然后将边界框信息(xyxy), 置信度，索引，重新整合成新xx = torch.cat((box[i], x[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)else:  # best class only  直接针对每个类别中概率最大的类别进行非极大值抑制处理conf, j = x[:, 5:].max(1)x = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]# Filter by classif classes:x = x[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)]# Apply finite constraint: 检测数据是否为有限数if not torch.isfinite(x).all():x = x[torch.isfinite(x).all(1)]# If none remain process next imagen = x.shape[0]  # number of boxesif not n:continue# Sort by confidencex = x[x[:, 4].argsort(descending=True)]# Batched NMS: x[box(xyxy), conf, max_index]# c = x[:, 5] * 0 if agnostic else x[:, 5]  # classes# boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4]  # boxes (offset by class), scores# 非极大值抑制处理，根据边界框与置信度来进行筛选，返回筛选后的索引boxes, scores = x[:, :4], x[:, 4]i = torchvision.ops.nms(boxes, scores, iou_thres)i = i[:max_num]  # 最多只保留前max_num个目标信息# iou筛选处理: 使用加权平均值合并框if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)try:  # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix# [11, 78] * [1, 78] = [11, 78]weights = iou * scores[None]  # box weights# torch.mm: 矩阵a和b矩阵相乘# torch.mul: 矩阵a和b数乘，维度不变# torch.mm(weights, x[:, :4]).float(): [11, 78] * [78, 4] = [11, 4]# weights.sum(1, keepdim=True): torch.Size([11, 1])x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes# i = i[iou.sum(1) > 1]  # require redundancyexcept:  # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139print(x, i, x.shape, i.shape)passoutput[xi] = x[i]if (time.time() - t) > time_limit:break  # time limit exceededreturn output

补充说明：函数一般经过上诉的5个步骤进行筛选，一般就已经可以得到最后的预测结果的。尤其是在对类别进行阈值筛选哪里，就已经把上前的样本筛选剩下几十几百个正样本，在经过nms稍稍处理，就可以挑选出比较合适的预测结果的。但是，在函数的最后还有一个设定，就是是否使用Merge nms进行进一步的处理。

所以，其实还有很多变种的nms操作，下面再介绍一下其他的nms处理方法，来对nms进行一个完整的认识与了解。

3. NMS的变体与实现

3.1 hard_nms_batch

官方实现的函数：torchvision.ops.boxes.batched_nms

需要传入三个参数：预测边界框（xyxy），置信度，预测类别信息

简单使用：

i = torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres)
output[image_i] = pred[i]

详细介绍：

def batched_nms(boxes: Tensor,scores: Tensor,idxs: Tensor,iou_threshold: float,
) -> Tensor:"""Performs non-maximum suppression in a batched fashion.Each index value correspond to a category, and NMSwill not be applied between elements of different categories.Args:boxes (Tensor[N, 4]): boxes where NMS will be performed. Theyare expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and``0 <= y1 < y2``.scores (Tensor[N]): scores for each one of the boxesidxs (Tensor[N]): indices of the categories for each one of the boxes.iou_threshold (float): discards all overlapping boxes with IoU > iou_thresholdReturns:Tensor: int64 tensor with the indices of the elements that have been kept by NMS, sortedin decreasing order of scores"""# Benchmarks that drove the following thresholds are at# https://github.com/pytorch/vision/issues/1311#issuecomment-781329339# Ideally for GPU we'd use a higher thresholdif boxes.numel() > 4_000 and not torchvision._is_tracing():return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)else:return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)

3.2 hard_nms

官方实现函数：torchvision.ops.boxes.nms或torchvision.ops.nms

需要传入两个参数：预测边界框（xyxy），置信度，不需要预测类别信息，比batch_nms少需要一个参数，更加的方便

简单使用：

i = torchvision.ops.boxes.nms(pred[:, :4], pred[:, 4], nms_thres)
output[image_i] = pred[i]

详细介绍：

def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:"""Performs non-maximum suppression (NMS) on the boxes accordingto their intersection-over-union (IoU).NMS iteratively removes lower scoring boxes which have anIoU greater than iou_threshold with another (higher scoring)box.If multiple boxes have the exact same score and satisfy the IoUcriterion with respect to a reference box, the selected box isnot guaranteed to be the same between CPU and GPU. This is similarto the behavior of argsort in PyTorch when repeated values are present.Args:boxes (Tensor[N, 4])): boxes to perform NMS on. Theyare expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and``0 <= y1 < y2``.scores (Tensor[N]): scores for each one of the boxesiou_threshold (float): discards all overlapping boxes with IoU > iou_thresholdReturns:Tensor: int64 tensor with the indices of the elements that have been keptby NMS, sorted in decreasing order of scores"""_assert_has_ops()return torch.ops.torchvision.nms(boxes, scores, iou_threshold)

3.3 and-nms

主要的思路：

其实和上面的自定义的nms方法是一致的，只不过这里进行了简化。一般来说，如果置信度最高的预测框与剩余的所以预测框的iou，只有有大于一定数值的，那么就可以判断为有重复的，那么同样的方法，剔除最好框与与最好框重复率较高的框，不断的重复来处理当前的类别的预测框。当当前类别处理好之后，再进行下一类别的预测框处理，唯一不同的就是判断的精简。但是，本质上来说都是类似的，处理的思路也是类似的；或者说，只是实现的方式不一样而已。

参考代码：

# 降序排列 为NMS做准备  [43, 6]
pred = pred[pred[:, 4].argsort(descending=True)]det_max = []  # 存放分数最高的框 即target
cls = pred[:, -1]
for c in cls.unique():  # 对所有的种类(不重复)dc = pred[cls == c]  # dc: 选出pred中所有类别是c的结果n = len(dc)  # 有多少个类别是c的预测框# 在hard-nms的逻辑基础上，增加是否为单独框的限制，删除没有重叠框的框（减少误检）。if method == 'and':  # requires overlap, single boxes erasedwhile len(dc) > 1:iou = bbox_iou(dc[0], dc[1:])  # iou with other boxesif iou.max() > 0.5:  # 删除没有重叠框的框/iou小于0.5的框（减少误检）det_max.append(dc[:1])# 首先需要去除分数最高的，以及其他重复的预测框，再进行下一轮的筛选dc = dc[1:][iou < nms_thres]  # remove ious > threshold# 对于同一张图的nms处理结果，进行拼接处理，并且按置信度进行降序排序(也就是从大到小排序)
if len(det_max):det_max = torch.cat(det_max)  # concatenate  因为之前是append进det_max的output[image_i] = det_max[(-det_max[:, 4]).argsort()]  # 排序

3.4 merge-nms

主要的思路：

对于一批重复率较高的框不是简单的去置信度最高的预测框，而且根据置信度赋予每个预测框一个权重值，置信度较高权重也较高，因为置信度高有理由更加的看重这个预测框。所以，对于所以的预测框乘上一个置信度权重，简单来说就是对预测框信息做一个权重和取平均，思想上是可以重合每个边界框的信息。然后其他操作类似，处理好这批相识度较高的预测框之后，继续处理当前类别下一批相识度较高的预测框。然后处理完当前类别，再进行下一类别预测框的处理。

参考代码：

主要代码如下，我注释非常详细，可以直接观看：

# 降序排列 为NMS做准备  [43, 6]
pred = pred[pred[:, 4].argsort(descending=True)]det_max = []  # 存放分数最高的框 即target
cls = pred[:, -1]
for c in cls.unique():  # 对所有的种类(不重复)dc = pred[cls == c]  # dc: 选出pred中所有类别是c的结果n = len(dc)  # 有多少个类别是c的预测框# 在hard-nms的基础上，增加保留框位置平滑策略（重叠框位置信息求解平均值），使框的位置更加精确。if method == 'merge':  # weighted mixture boxwhile len(dc):if len(dc) == 1:det_max.append(dc)break# 筛选出iou大于阈值的索引，这部分的预测框可以看出是重复的，对这部分筛选出来的预测框的置信度作为权重weightsi = bbox_iou(dc[0], dc) > nms_thres  # i = True/False的集合weights = dc[i, 4:5]     # 根据i，保留所有True# weights: [8, 1], dc[i, :4]: [8, 4], (weights * dc[i, :4]): [8, 4]# .sum(0): {Tensor: 4}, .sum(1): {Tensor: 8}, weights.sum(): {Tensor}# 就是对当前重复的预测框进行一个平均边界预测，用挑选出来的置信度与边界框进行相乘，将边界框的平均偏移量作为需要挑选的预测框，所以需要处于权重和# 权重较高的预测框，更需要重视，但是也不忽视其他的预测框就是merge的核心dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()  # 重叠框位置信息求解平均值# 将根据权重和处理好的预测框作为这批重复率高预测框的最好框，从而继续处理其他重复率高的预测框det_max.append(dc[:1])# i=0表示重合度比较低，提出重复率高的预测框，接下来继续循环这些剩下的重合度较低的预测框dc = dc[i == 0]# 对于同一张图的nms处理结果，进行拼接处理，并且按置信度进行降序排序(也就是从大到小排序)
if len(det_max):det_max = torch.cat(det_max)  # concatenate  因为之前是append进det_max的output[image_i] = det_max[(-det_max[:, 4]).argsort()]  # 排序

3.5 soft-nms

主要的思路：

对于当前置信度最高框与剩余框做一个iou，对于iou较高的预测框重复率应该是比较高的，对于这些预测框用宇哥衰减公式，其作用是对于iou较高的预测框衰减得最快，因为是呈指数形式衰减的。衰减后保留置信度仍然对于置信度阈值（这里设置为0.1），对于剩下的预测框的含义是重复率不算太高的。这里我觉得这是换了另外的一种方式来筛选，所以称之为soft筛选，之前的都是硬筛选（直接选择最好的，剩下的全部过滤）。

个人感觉，这样操作之后对密集预测会有帮助。因为有些密集的目标会被nms直接过滤掉，但是如果是用soft来衰减，对于这批可能重复率高，但是对于下一批重复率就不高，也就是当前样本可能是几个贴近的正样本，避免了被过滤。那么说，对于非密集预测，soft-nms带来的提升应该是不会太高的。

参考代码：

# 降序排列 为NMS做准备  [43, 6]
pred = pred[pred[:, 4].argsort(descending=True)]det_max = []  # 存放分数最高的框 即target
cls = pred[:, -1]
for c in cls.unique():  # 对所有的种类(不重复)dc = pred[cls == c]  # dc: 选出pred中所有类别是c的结果n = len(dc)  # 有多少个类别是c的预测框# soft-NMS: https://arxiv.org/abs/1704.04503# 推理时间：0.0030sif method == 'soft_nms':  sigma = 0.5  # soft-nms sigma parameterwhile len(dc):# if len(dc) == 1:  这是U版的源码 我做了个小改动#     det_max.append(dc)#     break# det_max.append(dc[:1])det_max.append(dc[:1])   # append dc的第一行  即targetif len(dc) == 1:breakiou = bbox_iou(dc[0], dc[1:])  # 计算target与其他框的iou# 这里和上面的直接置0不同，置0不需要管维度dc = dc[1:]  # dc=target往后的所有预测框# dc必须不包括target及其前的预测框，因为还要和值相乘, 维度上必须相同# dc[:, 4]: {Tensor: 36}# torch.exp(-iou ** 2 / sigma): {Tensor: 36}dc[:, 4] *= torch.exp(-iou ** 2 / sigma)  # 得分衰减{Tensor: 36}# 另外一种方式来挑选预测框, 衰减完后置信度还比较高的预测框，物理含义是重复率不算太高的# 因为重复率高，衰减的系数也会大，置信度会变得很小dc = dc[dc[:, 4] > conf_thres]# 对于同一张图的nms处理结果，进行拼接处理，并且按置信度进行降序排序(也就是从大到小排序)
if len(det_max):det_max = torch.cat(det_max)  # concatenate  因为之前是append进det_max的output[image_i] = det_max[(-det_max[:, 4]).argsort()]  # 排序

3.6 iou-nms

主要的思路：

对预测信息的类别进行不断遍历，每次只处理一个类别。对于当前的类别，当前的置信度最高的预测框，肯定可以作为之后的输出结果，那么还有一些预测框可能与当前置信度最高的预测框有重叠，那么就需要计算当前挑选框与剩余全部框的一个iou，当iou大于某个阈值说明重复率过高需要剔除；那么现在就可以更新当前需要处理的预测框了，不断的剔除最好框与与最好框重复率较高的框，处理完所有框后，就可以进行下一个类别预测框的处理，不断的循环。

参考代码：

# 降序排列 为NMS做准备  [43, 6]
pred = pred[pred[:, 4].argsort(descending=True)]det_max = []  # 存放分数最高的框 即target
cls = pred[:, -1]
for c in cls.unique():  # 对所有的种类(不重复)dc = pred[cls == c]  # dc: 选出pred中所有类别是c的结果n = len(dc)  # 有多少个类别是c的预测框# 推理时间：0.00299 是官方写的3倍if method == 'iou_nms':  # Hard NMS 自己写的 只支持单类别输入while dc.shape[0]:  # dc.shape[0]: 当前class的预测框数量det_max.append(dc[:1])  # 让score最大的一个预测框(排序后的第一个)为targetif len(dc) == 1:  # 出口 dc中只剩下一个框时，breakbreak# dc[0] ：target     dc[1:] ：其他预测框# 做的内容就是将当前所挑选的最好的预测框与其他剩余的预测框计算iou，当iou比较高说明重复率大，可以删除# 这里因为要剔除重复框，所以只保留小于阈值的预测框，因为大于阈值的预测框说明是重复的iou = bbox_iou(dc[0], dc[1:])  # 计算 普通iou# remove target and iou > threshold# 首先需要去除分数最高的，以及其他重复的预测框，再进行下一轮的筛选dc = dc[1:][iou < nms_thres]# 对于同一张图的nms处理结果，进行拼接处理，并且按置信度进行降序排序(也就是从大到小排序)
if len(det_max):det_max = torch.cat(det_max)  # concatenate  因为之前是append进det_max的output[image_i] = det_max[(-det_max[:, 4]).argsort()]  # 排序

3.7 diou_nms

主要的思路：

把普通的iou计算换成diou计算仅此而已，其他的与上述的iou一样，计算当前挑选框与剩余全部框的一个iou，当iou大于某个阈值说明重复率过高需要剔除。不断的剔除最好框与与最好框重复率较高的框，处理完所有框后，就可以进行下一个类别预测框的处理，不断的循环。

参考代码：

# 降序排列 为NMS做准备  [43, 6]
pred = pred[pred[:, 4].argsort(descending=True)]det_max = []  # 存放分数最高的框 即target
cls = pred[:, -1]
for c in cls.unique():  # 对所有的种类(不重复)dc = pred[cls == c]  # dc: 选出pred中所有类别是c的结果n = len(dc)  # 有多少个类别是c的预测框# 与iou_nms只是计算iou方面不一样而已if method == 'diou_nms':  # DIoU NMS  https://arxiv.org/pdf/1911.08287.pdfwhile dc.shape[0]:  # dc.shape[0]: 当前class的预测框数量det_max.append(dc[:1])  # 让score最大的一个预测框(排序后的第一个)为targetif len(dc) == 1:  # 出口 dc中只剩下一个框时，breakbreak# dc[0] ：target     dc[1:] ：其他预测框diou = bbox_iou(dc[0], dc[1:], DIoU=True)  # 计算 dioudc = dc[1:][diou < nms_thres]  # remove dious > threshold  保留True 删去False# 对于同一张图的nms处理结果，进行拼接处理，并且按置信度进行降序排序(也就是从大到小排序)
if len(det_max):det_max = torch.cat(det_max)  # concatenate  因为之前是append进det_max的output[image_i] = det_max[(-det_max[:, 4]).argsort()]  # 排序

主要不同就是将iou换成了diou ，使用的还是同一个函数，只是改变了一下参数，所以其实还可以使用giou_nms与ciou_nms，本质上没有变化。

4. NMS变体代码完整展示

需要注意，以下代码和yolov3spp代码是不一样的，不过可以直接替换使用。yolov3spp中使用的方法只是hard_nms处理，并且设置了一个可控参数选择是否使用merge_nms，这些nms的处理方法在以下代码中均可以选择是使用。

基于yolov3spp的代码更改：

def non_max_suppression(prediction, conf_thres=0.1,iou_thres=0.6, multi_label=True, method='iou_nms'):"""Removes detections with lower object confidence score than 'conf_thres'Non-Maximum Suppression to further filter detections.param:prediction: [batch, num_anchors(3个yolo预测层), (x+y+w+h+1+num_classes)]  3个anchor的预测结果总和conf_thres: 先进行一轮筛选，将分数过低的预测框（<conf_thres）删除（分数置0）nms_thres: iou阈值, 如果其余预测框与target的iou>iou_thres, 就将那个预测框置0multi_label: 是否是多标签method: nms方法  (https://github.com/ultralytics/yolov3/issues/679)(https://github.com/ultralytics/yolov3/pull/795)-hard_nms: 普通的 (hard) nms 官方实现(c函数库),可支持gpu,只支持单类别输入-hard_nms_batch: 普通的 (hard) nms 官方实现(c函数库),可支持gpu,支持多类别输入                    -and_nms: 在hard-nms的逻辑基础上，增加是否为单独框的限制，删除没有重叠框的框（减少误检）。-merge_nms: 在hard-nms的基础上，增加保留框位置平滑策略（重叠框位置信息求解平均值），使框的位置更加精确。-soft_nms: soft nms 用一个衰减函数作用在score上来代替原来的置0-iou_nms: 普通的 (hard) nms 只支持单类别输入-diou_nms: 普通的 (hard) nms 的基础上引入DIoU(普通的nms用的是iou)Returns detections with shape:(x1, y1, x2, y2, object_conf, class)"""nms_thres = iou_thresmulti_cls = multi_label# Box constraintsmin_wh, max_wh = 2, 4096  # (pixels) 宽度和高度的大小范围 [min_wh, max_wh]output = [None] * len(prediction)  # batch_size个output  存放最终筛选后的预测框结果for image_i, pred in enumerate(prediction):# 开始  pred = [12096, 25]# 第一层过滤   根据conf_thres虑除背景目标(obj_conf<conf_thres 0.1的目标 置信度极低的目标)pred = pred[pred[:, 4] > conf_thres]  # pred = [45, 25]# 第二层过滤   虑除超小anchor标和超大anchor  x=[45, 25]pred = pred[(pred[:, 2:4] > min_wh).all(1) & (pred[:, 2:4] < max_wh).all(1)]# 经过前两层过滤后如果该feature map没有目标框了，就结束这轮直接进行下一张图if len(pred) == 0:continue# 计算 scorepred[..., 5:] *= pred[..., 4:5]  # score = cls_conf * obj_conf# Box (center x, center y, width, height) to (x1, y1, x2, y2)box = xywh2xyxy(pred[:, :4])# Detections matrix nx6 (xyxy, conf, cls)if multi_cls or conf_thres < 0.01:# 第三轮过滤:针对每个类别score(obj_conf * cls_conf) > conf_thres [43, 6]# 这里一个框是有可能有多个物体的，所以要筛选# nonzero: 获得矩阵中的非0(True)数据的下标  a.t(): 将a矩阵拆开# i: 下标 [43]   j: 类别index [43] 过滤了两个score太低的i, j = (pred[:, 5:] > conf_thres).nonzero(as_tuple=False).t()# pred = [43, xyxy+score+class] [43, 6]# unsqueeze(1): [43] => [43, 1] add batch dimension# box[i]: [43,4] xyxy# pred[i, j + 5].unsqueeze(1): [43,1] score  对每个i,取第（j+5）个位置的值（第j个class的值cla_conf）# j.float().unsqueeze(1): [43,1] classpred = torch.cat((box[i], pred[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)else:  # best class onlyconf, j = pred[:, 5:].max(1)  # 一个类别直接取分数最大类的即可pred = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]# 第三轮过滤后如果该feature map没有目标框了，就结束这轮直接进行下一个feature mapif len(pred) == 0:continue# 第四轮过滤  这轮可有可无，一般没什么用 [43, 6] 检测数据是否为有限数pred = pred[torch.isfinite(pred).all(1)]# 降序排列 为NMS做准备  [43, 6]pred = pred[pred[:, 4].argsort(descending=True)]# Batched NMS# Batched NMS推理时间：0.054if method == 'hard_nms_batch':  # 普通的(hard)nms: 官方实现(c函数库),可支持gpu,但支持多类别输入# batched_nms：参数1 [43, xyxy]  参数2 [43, score]  参数3 [43, class]  参数4 [43, nms_thres]output[image_i] = pred[torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres)]# print("hard_nms_batch")continue# All other NMS methods  都是单类别输入det_max = []  # 存放分数最高的框 即targetcls = pred[:, -1]for c in cls.unique():  # 对所有的种类(不重复)dc = pred[cls == c]  # dc: 选出pred中所有类别是c的结果n = len(dc)  # 有多少个类别是c的预测框if n == 1:# No NMS required if only 1 predictiondet_max.append(dc)continueelif n > 500:# limit to first 500 boxes: https://github.com/ultralytics/yolov3/issues/117# 密集性 主要考虑到NMS是一个速度慢的算法（O(n^2)）,预测框太多,算法的效率太慢 所以这里筛选一下（最多500个预测框）dc = dc[:500]# 推理时间：0.001if method == 'hard_nms':  # 普通的(hard)nms: 只支持单类别输入det_max.append(dc[torchvision.ops.boxes.nms(dc[:, :4], dc[:, 4], nms_thres)])# 在hard-nms的逻辑基础上，增加是否为单独框的限制，删除没有重叠框的框（减少误检）。elif method == 'and_nms':  # requires overlap, single boxes erasedwhile len(dc) > 1:iou = bbox_iou(dc[0], dc[1:])  # iou with other boxesif iou.max() > 0.5:  # 删除没有重叠框的框/iou小于0.5的框（减少误检）det_max.append(dc[:1])dc = dc[1:][iou < nms_thres]  # remove ious > threshold# 在hard-nms的基础上，增加保留框位置平滑策略（重叠框位置信息求解平均值），使框的位置更加精确。elif method == 'merge_nms':  # weighted mixture boxwhile len(dc):if len(dc) == 1:det_max.append(dc)break# 筛选出iou大于阈值的索引，这部分的预测框可以看出是重复的，对这部分筛选出来的预测框的置信度作为权重weightsi = bbox_iou(dc[0], dc) > nms_thres  # i = True/False的集合weights = dc[i, 4:5]     # 根据i，保留所有True# weights: [8, 1], dc[i, :4]: [8, 4], (weights * dc[i, :4]): [8, 4]# .sum(0): {Tensor: 4}, .sum(1): {Tensor: 8}, weights.sum(): {Tensor}# 就是对当前重复的预测框进行一个平均边界预测，用挑选出来的置信度与边界框进行相乘，将边界框的平均偏移量作为需要挑选的预测框，所以需要处于权重和# 权重较高的预测框，更需要重视，但是也不忽视其他的预测框就是merge的核心dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()  # 重叠框位置信息求解平均值# 将根据权重和处理好的预测框作为这批重复率高预测框的最好框，从而继续处理其他重复率高的预测框det_max.append(dc[:1])# i=0表示重合度比较低，提出重复率高的预测框，接下来继续循环这些剩下的重合度较低的预测框dc = dc[i == 0]# 推理时间：0.0030selif method == 'soft_nms':  # soft-NMS      https://arxiv.org/abs/1704.04503sigma = 0.5  # soft-nms sigma parameterwhile len(dc):# if len(dc) == 1:  这是U版的源码 我做了个小改动#     det_max.append(dc)#     break# det_max.append(dc[:1])det_max.append(dc[:1])   # append dc的第一行  即targetif len(dc) == 1:breakiou = bbox_iou(dc[0], dc[1:])  # 计算target与其他框的iou# 这里和上面的直接置0不同，置0不需要管维度dc = dc[1:]  # dc=target往后的所有预测框# dc必须不包括target及其前的预测框，因为还要和值相乘, 维度上必须相同# dc[:, 4]: {Tensor: 36}# torch.exp(-iou ** 2 / sigma): {Tensor: 36}dc[:, 4] *= torch.exp(-iou ** 2 / sigma)  # 得分衰减{Tensor: 36}# 另外一种方式来挑选预测框dc = dc[dc[:, 4] > conf_thres]# 推理时间：0.00299 是官方写的3倍elif method == 'iou_nms':  # Hard NMS 只支持单类别输入while dc.shape[0]:  # dc.shape[0]: 当前class的预测框数量det_max.append(dc[:1])  # 让score最大的一个预测框(排序后的第一个)为targetif len(dc) == 1:  # 出口 dc中只剩下一个框时，breakbreak# dc[0] ：target     dc[1:] ：其他预测框# 做的内容就是将当前所挑选的最好的预测框与其他剩余的预测框计算iou，当iou比较高说明重复率大，可以删除# 这里因为要剔除重复框，所以只保留小于阈值的预测框，因为大于阈值的预测框说明是重复的iou = bbox_iou(dc[0], dc[1:])  # 计算 普通iou# remove target and iou > threshold# 首先需要去除分数最高的，以及其他重复的预测框，再进行下一轮的筛选dc = dc[1:][iou < nms_thres]# 推理时间：0.00299elif method == 'diou_nms':  # DIoU NMS  https://arxiv.org/pdf/1911.08287.pdfwhile dc.shape[0]:  # dc.shape[0]: 当前class的预测框数量det_max.append(dc[:1])  # 让score最大的一个预测框(排序后的第一个)为targetif len(dc) == 1:  # 出口 dc中只剩下一个框时，breakbreak# dc[0] ：target     dc[1:] ：其他预测框# 只是将计算iou变成了计算dioudiou = bbox_iou(dc[0], dc[1:], DIoU=True)  # 计算 dioudc = dc[1:][diou < nms_thres]  # remove dious > threshold  保留True 删去False# 对于同一张图的nms处理结果，进行拼接处理，并且按置信度进行降序排序(也就是从大到小排序)if len(det_max):det_max = torch.cat(det_max)  # concatenate  因为之前是append进det_max的output[image_i] = det_max[(-det_max[:, 4]).argsort()]  # 排序# output tensor [7, 6]return output

各种nms特点一句话总结：

Hard-nms–直接删除相邻的同类别目标，密集目标的输出不友好。
Soft-nms–改变其相邻同类别目标置信度（有关iou的函数），后期通过置信度阈值进行过滤，适用于目标密集的场景。
or-nms–hard-nms的非官方实现形式，只支持cpu。
vision-nms–hard-nms的官方实现形式(c函数库)，可支持gpu（cuda），只支持单类别输入。
vision-batched-nms–hard-nms的官方实现形式(c函数库)，可支持gpu（cuda）,支持多类别输入。
and-nms–在hard-nms的逻辑基础上，增加是否为单独框的限制，删除没有重叠框的框（减少误检）。
merge-nms–在hard-nms的基础上，增加保留框位置平滑策略（重叠框位置信息求解平均值），使框的位置更加精确。
diou-nms–在hard-nms的基础上，用diou替换iou，里有参照diou的优势。

以上几种nms的性能表现：https://github.com/ultralytics/yolov3/issues/679

看了一下几种nms的效果，看来无脑用merge-nms的效果是最好的，不过在实际工程项目中可以自己逐个试一下，这与计算边界框偏移损失一样，iou/giou/diou/ciou各种计算方法均也试一试。

当然，可能还存在其他各种各样的变体，这里就不再细诉了。

参考资料：

https://blog.csdn.net/qq_38253797/article/details/117920079
https://blog.csdn.net/qq_33270279/article/details/103721790

目标检测的Tricks | 【Trick9】nms非极大值抑制处理（包括变体merge-nms、and-nms、soft-nms、diou-nms等介绍）相关推荐

sklearn逻辑回归极大似然损失_收藏！攻克目标检测难点秘籍二，非极大值抑制与回归损失优化之路...
点击上方"AI算法修炼营",选择加星标或"置顶" 标题以下,全是干货前面的话在前面的秘籍一中,我们主要关注了模型加速之轻量化网络,对目标检测模型的实时性难点 ...
目标窗口检测算法-NMS非极大值抑制
1. NMS(non maximum suppression)的定义和算法步骤 NMS(non maximum suppression),中文名非极大值抑制,在很多计算机视觉任务中都有广泛应用,如:边 ...
目标检测中的LOU(交并比)和NMS(非极大值抑制)代码实现
1.LOU, 两个box框的交集比上并集,示意图如下所示: 代码如下所示: #假设box1的维度为[N,4] box2的维度为[M,4] def Lou(box1, box2):N = box1.si ...
Non-Maximum Suppression，NMS非极大值抑制
Non-Maximum Suppression,NMS非极大值抑制概述非极大值抑制(Non-Maximum Suppression,NMS),顾名思义就是抑制不是极大值的元素,可以理解为局部最大搜 ...
深度学习自学（三）：NMS非极大值抑制总结
非极大值抑制(Non-Maximum Suppression,NMS) 顾名思义就是抑制不是极大值的元素,可以理解为局部最大搜索.这个局部代表的是一个邻域,邻域有两个参数可变,一是邻域的维数,二是邻域 ...
Susan角点检测python实现（边缘检测、角点检测、重心计算、非极大值抑制）
Susan角点检测(边缘检测.角点检测.重心计算.非极大值抑制) 写在前面黄宁然--看过你看过的算法,觉得好难. 参考文献镇楼 [1]https://blog.csdn.net/tostq/arti ...
FAST角点检测算法（二）- 非极大值抑制筛选fast特征点
FAST角点检测算法(二)- 非极大值抑制筛选fast特征点 author@jason_ql(lql0716) http://blog.csdn.net/lql0716 fast角点检测算法参考文章& ...
【概念梳理】NMS 非极大值抑制
写在最前本文对网上关于 NMS 的解释整理了一下一.原理 YOLO在最后的一个步骤就是对 SxSx(Bx5+C) 个向量进行非极大值抑制(Non-max suppression),一开始不是太明白 ...
NMS——非极大值抑制
NMS(non maximum suppression),中文名非极大值抑制,在很多计算机视觉任务中都有广泛应用,如:边缘检测.目标检测等. 这里主要以人脸检测中的应用为例,来说明NMS,并给出Mat ...

目标检测的Tricks | 【Trick9】nms非极大值抑制处理（包括变体merge-nms、and-nms、soft-nms、diou-nms等介绍）