【目标检测】NMS、Soft-NMS、Fast-NMS

【目标检测】NMS、Soft-NMS、fast-NMS

• NMS
• Soft-NMS
• fast-NMS

NMS

NMS是目标检测中常用的算法，最终网络的输出可能包括了很多预测框，其中一个目标周围会有许多候选框，如下图所示，而该算法目的是找出最佳的目标位置。

主要有两种实现方式，一种是将一张图片中的所有的预测框聚在一起进行NMS，一种是将一张图片中所有预测框按所预测的类别分别进行NMS（也就是以类别为外循环），但主要的思想还是不变，假设给定预测框bboxes和scores，维度分别为[N,4]、[N,]，其中N为预测框的个数，该算法步骤如下：

1. 先将预测框按scores大小降序排列，这样保证了留下的框置信度最大。
2. 从第一个预测框开始，分别计算该预测框（代号为i）与剩下的预测框之间的交并比IoU，排除掉IoU大于阈值（一般是0.5）的其他框，因为这些大于阈值的框就说明与当前预测框i太接近了，同时置信度也没有i大，所以需要排除掉，依次循环直至处理完所有框。

代码如下所示：

def nms(self, bboxes, scores, threshold=0.5):
        x1 = bboxes[:,0]
        y1 = bboxes[:,1]
        x2 = bboxes[:,2]
        y2 = bboxes[:,3]
        areas = (x2-x1)*(y2-y1)   # [N,]
        _, order = scores.sort(0, descending=True)    # 降序

        keep = []
        while order.numel() > 0:       # torch.numel()返回张量元素个数
            if order.numel() == 1:     # 保留框只剩一个
                i = order.item()
                keep.append(i)
                break
            else:
                i = order[0].item()    # 保留scores最大的那个框box[i]
                keep.append(i)

            # 计算box[i]与其余各框的IOU
            xx1 = x1[order[1:]].clamp(min=x1[i])   # [N-1,]
            yy1 = y1[order[1:]].clamp(min=y1[i])
            xx2 = x2[order[1:]].clamp(max=x2[i])
            yy2 = y2[order[1:]].clamp(max=y2[i])
            inter = (xx2-xx1).clamp(min=0) * (yy2-yy1).clamp(min=0)   # [N-1,]

            iou = inter / (areas[i]+areas[order[1:]]-inter)  # [N-1,]
            idx = (iou <= threshold).nonzero().squeeze() # 注意此时idx为[N-1,] 而order为[N,]
            if idx.numel() == 0:
                break
            order = order[idx+1]  # 修补索引之间的差值
        return torch.LongTensor(keep) 

Soft-NMS

该算法在论文Improving Object Detection With One Line of Code（CVPR2017）中被提出来，此处为github链接。

例如上图，如果按照上述传统的NMS算法，红色框比绿色框置信度高，所以排序后会先处理红色框，而此时在计算与其他框的IoU时，绿色框和红色框的IoU大于阈值，从而会排除掉绿色框，同时NMS的阈值也难以确定，设置高了会误检，设置低了会漏检，如上图所示的情况。而Soft-NMS的思路就是将所有IoU大于阈值的框降低其置信度，而不是删除。

算法如下图所示，$MMM$si​为相应框的得分。

传统的NMS可以描述为将IoU大于阈值的框的得分置为0：

而Soft-NMS提出了两种方式，一种是线性加权，一种是高斯加权：

线性加权

高斯加权

实验结果如下图所示，可以看出来在不增加额外的计算量下可以平均提升1%的精度。

Fast-NMS

该算法是在实例分割论文YOLACT中所提出，此处为github链接。该论文主打实时(fps>30)，说传统的NMS可以利用矩阵简化从而降低时间，但不得不牺牲一些精度，实验结果显示虽然降低了0.1mAP，但时间比基于Cython实现的NMS快11.8ms。

算法代码如下：

def fast_nms(self, boxes, masks, scores, iou_threshold:float=0.5, top_k:int=200, second_threshold:bool=False):
        '''
        boxes:  torch.Size([num_dets, 4])
        masks:  torch.Size([num_dets, 32])
        scores: torch.Size([num_classes, num_dets])
        '''
        # step1: 每一类的框按照scores降序排序后取前top_k个
        scores, idx = scores.sort(1, descending=True) 
        # scores为降序排列 
        # idx为原顺序的索引 
        idx = idx[:, :top_k].contiguous() # 取前top_k个框 
        scores = scores[:, :top_k] 
        num_classes, num_dets = idx.size()

        boxes = boxes[idx.view(-1), :].view(num_classes, num_dets, 4) # torch.Size([num_classes, num_dets, 4])
        masks = masks[idx.view(-1), :].view(num_classes, num_dets, -1) # torch.Size([num_classes, num_dets, 32]) 其中32为生成的系数个数
        # step2: 计算每一类中，box与box之间的IoU
        iou = jaccard(boxes, boxes) # torch.Size([num_classes, num_dets, num_dets])
        iou.triu_(diagonal=1) # triu_()取上三角 tril_()取下三角 此处将矩阵的下三角和对角线元素删去
        iou_max, _ = iou.max(dim=1) # 按列取大值 torch.Size([num_classes, num_dets])

        # 过滤掉iou大于阈值的框 
        keep = (iou_max <= iou_threshold) # torch.Size([num_classes, num_dets])

        if second_threshold: # 保证保留的框满足一定的置信度
            keep *= (scores > self.conf_thresh)

        # Assign each kept detection to its corresponding class
        classes = torch.arange(num_classes, device=boxes.device)[:, None].expand_as(keep)
        '''
        tensor([[ 0,  0,  0,  ...,  0,  0,  0],
        [ 1,  1,  1,  ...,  1,  1,  1],
        [ 2,  2,  2,  ...,  2,  2,  2],
        ...,
        [77, 77, 77,  ..., 77, 77, 77],
        [78, 78, 78,  ..., 78, 78, 78],
        [79, 79, 79,  ..., 79, 79, 79]])
        '''
        classes = classes[keep]
        boxes = boxes[keep]
        masks = masks[keep]
        scores = scores[keep]
        # Only keep the top cfg.max_num_detections highest scores across all classes
        scores, idx = scores.sort(0, descending=True)
        idx = idx[:cfg.max_num_detections]
        scores = scores[:cfg.max_num_detections]
        classes = classes[idx]e
        boxes = boxes[idx]
        masks = masks[idx]
        return boxes, masks, classes, scores # torch.Size([max_num_detections])

下面通过一个例子🌰来模拟该过程，假设属于狗这一类目前有3个预测框，并且已经按得分降序排列，即该类别下$score(bbox1)>score(bbox2)score(bbox1)>score(bbox2)scor{e}_{(bbox1)}>scor{e}_{(bbox2)}$score(bbox1)​>score(bbox2)​，下列表格为预测框之间的IoU值。

	bbox1	bbox2	bbox3
bbox1	1	0.7	0.2
bbox2	0.7	1	0.8
bbox3	0.2	0.8	1

之后通过torch中triu这个函数取出上三角并将对角线置0。

	bbox1	bbox2	bbox3
bbox1	0	0.7	0.2
bbox2	0	0	0.8
bbox3	0	0	0

按列取最大值。

bbox1	bbox2	bbox3
0	0.7	0.8

舍弃超过阈值的框，假设阈值为0.5，那么该类中bbox2和bbox3都要被舍弃，只留下了bbox1。因为首先预测框已经按得分降序排列好了，并且每一个元素都是行号小于列号，元素大于阈值就代表着这一列对应的框与比它置信度高的框过于重叠，比如0.7，就代表着bbox2和bbox1过于重叠，并且bbox1的置信度较高，应该排除bbox2。但是如果按照传统的NMS方法，bbox2会被排除，bbox3会被保留。所以fast-NMS会去掉更多的框，但是作者认为速度至上，并且实验也证明精度下降一点点速度却大大提高，如下图所示。

参考博客：
[1] NMS
[2] Soft-NMS
[3] fast-NMS