pytorch实现yolov3中使用的nms(非最大抑制)理解

前期在研究pytorch实现yolov3时，虽然大体搞明白了yolov3的原理，但是由于项目紧急，没来得及做笔记，趁着这段时间不是那么忙了，整理一遍过程。首先说说nms(非极大抑制)。

对于尺寸为416×416的输入图像，YOLO预测（（52×52）+（26×26）+ 13×13））×3 = 10647个边界框。 但是，在我的图像中只有三个物体(本人，水杯，手机)。 我们如何将检测结果从10647减少到3呢？

分两步：

2、基于对象置信度的阈值。
首先，我们根据对象分数(objectness score)过滤框。 通常，具有低于阈值分数的框被忽略。

3、非最大抑制（NMS）
NMS打算解决同一图像的多重检测问题。通过抑制不是极大值的元素，可以理解为局部最大搜索。这个局部代表的是一个邻域，邻域有两个参数可变，一是邻域的维数，二是邻域的大小。这里讨论用于目标检测中提取分数最高的窗口的。例如在行人检测中，滑动窗口经提取特征，经分类器分类识别后，每个窗口都会得到一个分数。但是滑动窗口会导致很多窗口与其他窗口存在包含或者大部分交叉的情况。这时就需要用到NMS来选取那些邻域里分数最高（是行人的概率最大），并且抑制那些分数低的窗口。

3.1、算法过程

Step1：按置信概率排列相应的备选框

Step2：取最大的框作为保留框，与其IOU大于阈值的框删除掉

Step3：剩下的框执行Step2

图例解释

首先根据候选框的类别分类概率做排序：A<B<C<D<E<F

1）先标记最大概率矩形框F是我们要保留下来的；
2）从最大概率矩形框F开始，分别判断A~E与F的重叠度IOU（两框的交并比）是否大于某个设定的阈值，假设B、D与F的重叠度超过阈值，那么就扔掉B、D；

3）从剩下的矩形框A、C、E中，选择概率最大的E，标记为要保留下来的，然后判读E与A、C的重叠度，扔掉重叠度超过设定阈值的矩形框

就这样一直重复下去，直到剩下的矩形框没有了，标记完所有要保留下来的矩形框

    batch_size = prediction.size(0)
    write = False
    for ind in range(batch_size):
        max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
        max_conf = max_conf.float().unsqueeze(1)
        max_conf_score = max_conf_score.float().unsqueeze(1)
        seq = (image_pred[:,:5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)
        non_zero_ind =  (torch.nonzero(image_pred[:,4])) 
        try:
            image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
        except:
            continue
        if image_pred_.shape[0] == 0:
            continue
        img_classes = unique(image_pred_[:,-1])
        for cls in img_classes:
            cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
            class_mask_ind   = torch.nonzero(cls_mask[:,-2]).squeeze()
            image_pred_class = image_pred_[class_mask_ind].view(-1,7)
            conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)   #Number of detections
            
            for i in range(idx):
                try:
                    ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
                except ValueError:
                    break
            
                except IndexError:
                    break
                iou_mask = (ious < nms_conf).float().unsqueeze(1)
                image_pred_class[i+1:] *= iou_mask       
                non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                image_pred_class = image_pred_class[non_zero_ind].view(-1,7)
                
            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
            seq = batch_ind, image_pred_class
            
            if not write:
                output = torch.cat(seq,1)
                write = True
            else:
                out = torch.cat(seq,1)
                output = torch.cat((output,out))

L1：batch_size = prediction.size(0)获取一个batch的图片数量

对每一张图片得分的预测值进行NMS操作，因为每张图片的目标数量不一样，所以有效得分的方框的数量不一样，没法将几张图片同时处理，因此一次只能完成一张图的置信度阈值的设置和NMS,不能将所涉及的操作向量化.所以必须在预测的第一个维度上（batch数量）上遍历每张图片，将得分低于一定分数的去掉，对剩下的方框进行进行NMS

L4：max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)

image_pred.shape → 10647 * 85，其中的85包含以下信息

(x1 , x2 , y1 , y2 , objectness score , cls0_score , cls1_score , cls2_score ,…, cls78_score , cls79_score)

因此image_pred[:,5:5+ num_classes] → (cls0_score , cls1_score , cls2_score ,…, cls78_score , cls79_score)

torch.max(image_pred[:,5:5+ num_classes], 1)表示在cls0_score 到cls79_score 中找的最大得分返回给max_conf，最大得分的index返回给max_conf_score

L8：image_pred = torch.cat(seq, 1)

image_pred .shape → 10647 * 7，其中的7代表

(x1 , x2 , y1 , y2 , objectness score , max_conf , max_conf_score )

L9：non_zero_ind = (torch.nonzero(image_pred[:,4]))

在预测的10647个predict bbox中，objectness score 大于 confidence的抽出，返回index

以上实现了上述两步走的第一步基于对象置信度的阈值

L11:image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)

根据L9中获取的index，从原image_pred抽出阈值大于confidence的predict bbox

print(image_pred_ )结果为

L14:if image_pred_.shape[0] == 0:表示没有预测到物体

L16：img_classes = unique(image_pred_[:,-1])

image_pred_的维度1的信息(x1 , x2 , y1 , y2 , objectness score , max_conf , max_conf_score )

所以image_pred_[:,-1]表示max_conf_score 既物体的分类，取0~79

对于我的场景(人，水杯，手机)，img_classes 对应(0,41,67)，其中0代表人，41代表水杯，67代表手机

print(img_classes)结果为 tensor([ 0., 41., 67.])

L18：cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)抽出其中的一个分类(例如抽出41 水杯)

print(cls_mask)结果如下

L19：class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()

cls_mask[:,-2]表示

L20:image_pred_class = image_pred_[class_mask_ind].view(-1,7)

print(image_pred_class )结果

L21：conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]

对image_pred_class[:,4] objectness score 排序

L27：ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])

计算第i个方框和i+1到最终的所有方框的IOU，这里实现了两步走的第二步非最大抑制（NMS）

L33：iou_mask = (ious < nms_conf).float().unsqueeze(1)和L34：image_pred_class[i+1:] *= iou_mask

将IOU大于阈值的框置位0

L35和L36

non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()

image_pred_class = image_pred_class[non_zero_ind].view(-1,7)

print(image_pred_class )结果

最终返回的output是一个batch中所有图片中剩下的方框的属性，一行对应一个方框，属性为(x1,y1,x2,y2,s,s_cls,index_cls)，
ind 是这个方框所属图片在这个batch中的序号，x1,y1是在网络输入图片(416x416)坐标系中，方框左上角的坐标；x2,y2是在网络输入
图片(416x416)坐标系中，方框右下角的坐标。s是这个方框含有目标的得分s_cls是这个方框中所含目标最有可能的类别的概率得分，index_cls是s_cls对应的这个类别所对应的序号