如有错误，恳请指出。
用这篇博客记录一下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
-
NMS主要步骤
-
下面是yolov3spp中，nms算法的主要步骤：
1）剔除置信度较低的背景目标
x = x[x[:, 4] > conf_thres]
2）剔除宽高较小或者较大的目标
min_wh, max_wh = 2, 4096
x = x[((x[:, 2:4] > min_wh) & (x[:, 2:4] < max_wh)).all(1)]
3）剔除类别概率较低的目标
这里可以选择对每一个类别都进行检测其概率，挑选出那些大于阈值的正样本，返回预测样本索引与每个样本的满足条件的阈值索引。也可以对概率最大的类别进行筛选，这种比较直观。就是当概率最大的类别满足阈值条件，说明该样本可以符合筛选条件，也就不需要在乎该正样本其他概率较小的类别
if multi_label:
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()
x = torch.cat((box[i], x[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)
else:
conf, j = x[:, 5:].max(1)
x = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]
4）检测数据是否为有限数
if not torch.isfinite(x).all():
x = x[torch.isfinite(x).all(1)]
5）剩余的数据进行nms，保留前max_num个数据
这个部分验证可能有用，测试基本没用。因为测试阶段一般不会有超过max_num个数据
c = x[:, 5] * 0 if agnostic else x[:, 5]
boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4]
i = torchvision.ops.nms(boxes, scores, iou_thres)
i = i[: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 according
to their intersection-over-union (IoU).
NMS iteratively removes lower scoring boxes which have an
IoU greater than iou_threshold with another (higher scoring)
box.
If multiple boxes have the exact same score and satisfy the IoU
criterion with respect to a reference box, the selected box is
not guaranteed to be the same between CPU and GPU. This is similar
to the behavior of argsort in PyTorch when repeated values are present.
Args:
boxes (Tensor[N, 4])): boxes to perform NMS on. They
are expected to be in 看了函数介绍，感觉正常来说就是普通传入边界框信息与置信度信息就可以了，所以这里我尝试对其进行了更改，更改为也更加直观的版本：
boxes, scores = x[:, :4], x[:, 4]
i = torchvision.ops.nms(boxes, scores, iou_thres)
i = i[:max_num]
下面分别查看更改前后，对图片的测试结果：
- 更改前：
- 更改后：
感觉更改前后效果好像没有什么不同，如果有知道为什么需要加上这么一个大的数值的朋友可以告诉我一下，谢谢。
- 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 results
param:
prediction: [batch, num_anchors(3个yolo预测层), (x+y+w+h+1+num_classes)] 3个anchor的预测结果总和
conf_thres: 先进行一轮筛选，将分数过低的预测框（iou_thres, 就将那个预测框置0
multi_label: 是否是多标签
max_num：筛选的最大数目
Returns detections with shape:
(x1, y1, x2, y2, object_conf, class)
"""
merge = True
min_wh, max_wh = 2, 4096
time_limit = 10.0
t = time.time()
nc = prediction[0].shape[1] - 5
multi_label &= nc > 1
output = [None] * prediction.shape[0]
for xi, x in enumerate(prediction):
x = x[x[:, 4] > conf_thres]
x = x[((x[:, 2:4] > min_wh) & (x[:, 2:4] < max_wh)).all(1)]
if not x.shape[0]:
continue
x[..., 5:] *= x[..., 4:5]
box = xywh2xyxy(x[:, :4])
if multi_label:
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()
x = torch.cat((box[i], x[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)
else:
conf, j = x[:, 5:].max(1)
x = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]
if classes:
x = x[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)]
if not torch.isfinite(x).all():
x = x[torch.isfinite(x).all(1)]
n = x.shape[0]
if not n:
continue
x = x[x[:, 4].argsort(descending=True)]
boxes, scores = x[:, :4], x[:, 4]
i = torchvision.ops.nms(boxes, scores, iou_thres)
i = i[:max_num]
if merge and (1 < n < 3E3):
try:
iou = box_iou(boxes[i], boxes) > iou_thres
weights = iou * scores[None]
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)
except:
print(x, i, x.shape, i.shape)
pass
output[xi] = x[i]
if (time.time() - t) > time_limit:
break
return output
补充说明：函数一般经过上诉的5个步骤进行筛选，一般就已经可以得到最后的预测结果的。 尤其是在对类别进行阈值筛选哪里，就已经把上前的样本筛选剩下几十几百个正样本，在经过nms稍稍处理，就可以挑选出比较合适的预测结果的。但是，在函数的最后还有一个设定，就是是否使用Merge nms进行进一步的处理。
所以，其实还有很多变种的nms操作，下面再介绍一下其他的nms处理方法，来对nms进行一个完整的认识与了解。
- 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 NMS
will not be applied between elements of different categories.
Args:
boxes (Tensor[N, 4]): boxes where NMS will be performed. They
are expected to be in 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 according
to their intersection-over-union (IoU).
NMS iteratively removes lower scoring boxes which have an
IoU greater than iou_threshold with another (higher scoring)
box.
If multiple boxes have the exact same score and satisfy the IoU
criterion with respect to a reference box, the selected box is
not guaranteed to be the same between CPU and GPU. This is similar
to the behavior of argsort in PyTorch when repeated values are present.
Args:
boxes (Tensor[N, 4])): boxes to perform NMS on. They
are expected to be in 3.3 and-nms
- *主要的思路：
其实和上面的自定义的nms方法是一致的，只不过这里进行了简化。一般来说，如果置信度最高的预测框与剩余的所以预测框的iou，只有有大于一定数值的，那么就可以判断为有重复的，那么同样的方法，剔除最好框与与最好框重复率较高的框，不断的重复来处理当前的类别的预测框。当当前类别处理好之后，再进行下一类别的预测框处理，唯一不同的就是判断的精简。但是，本质上来说都是类似的，处理的思路也是类似的；或者说，只是实现的方式不一样而已。
- *参考代码：
pred = pred[pred[:, 4].argsort(descending=True)]
det_max = []
cls = pred[:, -1]
for c in cls.unique():
dc = pred[cls == c]
n = len(dc)
if method == 'and':
while len(dc) > 1:
iou = bbox_iou(dc[0], dc[1:])
if iou.max() > 0.5:
det_max.append(dc[:1])
dc = dc[1:][iou < nms_thres]
if len(det_max):
det_max = torch.cat(det_max)
output[image_i] = det_max[(-det_max[:, 4]).argsort()]
3.4 merge-nms
- *主要的思路：
对于一批重复率较高的框不是简单的去置信度最高的预测框，而且根据置信度赋予每个预测框一个权重值，置信度较高权重也较高，因为置信度高有理由更加的看重这个预测框。所以，对于所以的预测框乘上一个置信度权重，简单来说就是对预测框信息做一个权重和取平均，思想上是可以重合每个边界框的信息。然后其他操作类似，处理好这批相识度较高的预测框之后，继续处理当前类别下一批相识度较高的预测框。然后处理完当前类别，再进行下一类别预测框的处理。
- *参考代码：
主要代码如下，我注释非常详细，可以直接观看：
pred = pred[pred[:, 4].argsort(descending=True)]
det_max = []
cls = pred[:, -1]
for c in cls.unique():
dc = pred[cls == c]
n = len(dc)
if method == 'merge':
while len(dc):
if len(dc) == 1:
det_max.append(dc)
break
i = bbox_iou(dc[0], dc) > nms_thres
weights = dc[i, 4:5]
dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()
det_max.append(dc[:1])
dc = dc[i == 0]
if len(det_max):
det_max = torch.cat(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带来的提升应该是不会太高的。
- *参考代码：
pred = pred[pred[:, 4].argsort(descending=True)]
det_max = []
cls = pred[:, -1]
for c in cls.unique():
dc = pred[cls == c]
n = len(dc)
if method == 'soft_nms':
sigma = 0.5
while len(dc):
det_max.append(dc[:1])
if len(dc) == 1:
break
iou = bbox_iou(dc[0], dc[1:])
dc = dc[1:]
dc[:, 4] *= torch.exp(-iou ** 2 / sigma)
dc = dc[dc[:, 4] > conf_thres]
if len(det_max):
det_max = torch.cat(det_max)
output[image_i] = det_max[(-det_max[:, 4]).argsort()]
3.6 iou-nms
- *主要的思路：
对预测信息的类别进行不断遍历，每次只处理一个类别。对于当前的类别，当前的置信度最高的预测框，肯定可以作为之后的输出结果，那么还有一些预测框可能与当前置信度最高的预测框有重叠，那么就需要计算当前挑选框与剩余全部框的一个iou，当iou大于某个阈值说明重复率过高需要剔除；那么现在就可以更新当前需要处理的预测框了，不断的剔除最好框与与最好框重复率较高的框，处理完所有框后，就可以进行下一个类别预测框的处理，不断的循环。
- *参考代码：
pred = pred[pred[:, 4].argsort(descending=True)]
det_max = []
cls = pred[:, -1]
for c in cls.unique():
dc = pred[cls == c]
n = len(dc)
if method == 'iou_nms':
while dc.shape[0]:
det_max.append(dc[:1])
if len(dc) == 1:
break
iou = bbox_iou(dc[0], dc[1:])
dc = dc[1:][iou < nms_thres]
if len(det_max):
det_max = torch.cat(det_max)
output[image_i] = det_max[(-det_max[:, 4]).argsort()]
3.7 diou_nms
- *主要的思路：
把普通的iou计算换成diou计算仅此而已，其他的与上述的iou一样，计算当前挑选框与剩余全部框的一个iou，当iou大于某个阈值说明重复率过高需要剔除。不断的剔除最好框与与最好框重复率较高的框，处理完所有框后，就可以进行下一个类别预测框的处理，不断的循环。
- *参考代码：
pred = pred[pred[:, 4].argsort(descending=True)]
det_max = []
cls = pred[:, -1]
for c in cls.unique():
dc = pred[cls == c]
n = len(dc)
if method == 'diou_nms':
while dc.shape[0]:
det_max.append(dc[:1])
if len(dc) == 1:
break
diou = bbox_iou(dc[0], dc[1:], DIoU=True)
dc = dc[1:][diou < nms_thres]
if len(det_max):
det_max = torch.cat(det_max)
output[image_i] = det_max[(-det_max[:, 4]).argsort()]
主要不同就是将iou换成了diou ，使用的还是同一个函数，只是改变了一下参数，所以其实还可以使用giou_nms与ciou_nms，本质上没有变化。
- 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: 先进行一轮筛选，将分数过低的预测框（iou_thres, 就将那个预测框置0
multi_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_thres
multi_cls = multi_label
min_wh, max_wh = 2, 4096
output = [None] * len(prediction)
for image_i, pred in enumerate(prediction):
pred = pred[pred[:, 4] > conf_thres]
pred = pred[(pred[:, 2:4] > min_wh).all(1) & (pred[:, 2:4] < max_wh).all(1)]
if len(pred) == 0:
continue
pred[..., 5:] *= pred[..., 4:5]
box = xywh2xyxy(pred[:, :4])
if multi_cls or conf_thres < 0.01:
i, j = (pred[:, 5:] > conf_thres).nonzero(as_tuple=False).t()
pred = torch.cat((box[i], pred[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)
else:
conf, j = pred[:, 5:].max(1)
pred = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]
if len(pred) == 0:
continue
pred = pred[torch.isfinite(pred).all(1)]
pred = pred[pred[:, 4].argsort(descending=True)]
if method == 'hard_nms_batch':
output[image_i] = pred[torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres)]
continue
det_max = []
cls = pred[:, -1]
for c in cls.unique():
dc = pred[cls == c]
n = len(dc)
if n == 1:
det_max.append(dc)
continue
elif n > 500:
dc = dc[:500]
if method == 'hard_nms':
det_max.append(dc[torchvision.ops.boxes.nms(dc[:, :4], dc[:, 4], nms_thres)])
elif method == 'and_nms':
while len(dc) > 1:
iou = bbox_iou(dc[0], dc[1:])
if iou.max() > 0.5:
det_max.append(dc[:1])
dc = dc[1:][iou < nms_thres]
elif method == 'merge_nms':
while len(dc):
if len(dc) == 1:
det_max.append(dc)
break
i = bbox_iou(dc[0], dc) > nms_thres
weights = dc[i, 4:5]
dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()
det_max.append(dc[:1])
dc = dc[i == 0]
elif method == 'soft_nms':
sigma = 0.5
while len(dc):
det_max.append(dc[:1])
if len(dc) == 1:
break
iou = bbox_iou(dc[0], dc[1:])
dc = dc[1:]
dc[:, 4] *= torch.exp(-iou ** 2 / sigma)
dc = dc[dc[:, 4] > conf_thres]
elif method == 'iou_nms':
while dc.shape[0]:
det_max.append(dc[:1])
if len(dc) == 1:
break
iou = bbox_iou(dc[0], dc[1:])
dc = dc[1:][iou < nms_thres]
elif method == 'diou_nms':
while dc.shape[0]:
det_max.append(dc[:1])
if len(dc) == 1:
break
diou = bbox_iou(dc[0], dc[1:], DIoU=True)
dc = dc[1:][diou < nms_thres]
if len(det_max):
det_max = torch.cat(det_max)
output[image_i] = det_max[(-det_max[:, 4]).argsort()]
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
Original: https://blog.csdn.net/weixin_44751294/article/details/124364145
Author: Clichong
Title: 目标检测的Tricks | 【Trick9】nms非极大值抑制处理（包括变体merge-nms、and-nms、soft-nms、diou-nms等介绍）
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/682342/
转载文章受原作者版权保护。转载请注明原作者出处！