ssd.pytorch源码分析（三）— 非极大值抑制NMS

NMS源码

SSD论文链接

NMS介绍

吴恩达对于NMS（非极大值抑制）的介绍：

说白了，NMS的作用就是去掉目标检测任务重复的检测框。 例如，一个目标有多个选择框，现在要去掉多余的选择框。怎么做呢？循环执行步骤1和2， 直到只剩下一个框：

• 1、选出置信度p_c最高的框；
• 2、去掉和这个框IOU＞0.7的框。

相关函数

一、torch.clamp( ）

torch.clamp(input, min, max, out=None) → Tensor
           

将输入input张量每个元素夹紧到区间 [min,max]，并返回结果到一个新张量。

类似于numpy中的np.clip

操作定义如下：

| min, if x_i < min
y_i = | x_i, if min <= x_i <= max
      | max, if x_i > max
           

参数：

• input (Tensor) – 输入张量
• min (Number) – 限制范围下限
• max (Number) – 限制范围上限
• out (Tensor, optional) – 输出张量

例子：

>>> a = torch.randn(4)
>>> a
 1.3869
 0.3912
-0.8634
-0.5468
[torch.FloatTensor of size 4]
>>> torch.clamp(a, min=-0.5, max=0.5)
 0.5000
 0.3912
-0.5000
-0.5000
[torch.FloatTensor of size 4]
           

二、torch.index_select()

torch.index_select(input, dim, index, out=None) → Tensor
           

沿着指定维度对输入进行切片。

参数：

• input (Tensor) – 输入张量
• dim (int) – 索引的轴
• index (LongTensor) – 包含索引下标的一维张量
• out (Tensor, optional) – 目标张量

例子：

>>> x = torch.randn(3, 4)
>>> x

 1.2045  2.4084  0.4001  1.1372
 0.5596  1.5677  0.6219 -0.7954
 1.3635 -1.2313 -0.5414 -1.8478
[torch.FloatTensor of size 3x4]

>>> indices = torch.LongTensor([0, 2])
>>> torch.index_select(x, 0, indices)

 1.2045  2.4084  0.4001  1.1372
 1.3635 -1.2313 -0.5414 -1.8478
[torch.FloatTensor of size 2x4]

>>> torch.index_select(x, 1, indices)

 1.2045  0.4001
 0.5596  0.6219
 1.3635 -0.5414
[torch.FloatTensor of size 3x2]
           

注意，index_select函数中的参数index表示了有哪些索引值是需要保留的。

三、 torch.numel()

返回input 张量中的元素个数。

复现代码

以下为ssd.pytorch中NMS（实际上在任何anchor based的目标检测框架中都适用）。其中：

• 为了减少计算量，作者仅选取置信度前top_k=200个框；
• 代码中包含了IOU的计算。关于IOU计算推荐阅读这篇文章；

def nms(boxes, scores, overlap=0.7, top_k=200):
    """
    输入:
        boxes: 存储一个图片的所有预测框。[num_positive,4].
        scores:置信度。如果为多分类则需要将nms函数套在一个循环内。[num_positive].
        overlap: nms抑制时iou的阈值.
        top_k: 先选取置信度前top_k个框再进行nms.
    返回:
        nms后剩余预测框的索引.
    """
    
    keep = scores.new(scores.size(0)).zero_().long() 
    # 保存留下来的box的索引 [num_positive]
    # 函数new(): 构建一个有相同数据类型的tensor 
    
	#如果输入box为空则返回空Tensor
    if boxes.numel() == 0: 
        return keep
        
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    area = torch.mul(x2 - x1, y2 - y1) #并行化计算所有框的面积
    v, idx = scores.sort(0)  # 升序排序
    idx = idx[-top_k:]  # 前top-k的索引，从小到大
    xx1 = boxes.new()
    yy1 = boxes.new()
    xx2 = boxes.new()
    yy2 = boxes.new()
    w = boxes.new()
    h = boxes.new()

    count = 0
    while idx.numel() > 0:
        i = idx[-1]  # 目前最大score对应的索引
        keep[count] = i #存储在keep中
        count += 1
        if idx.size(0) == 1: #跳出循环条件：box被筛选完了
            break
        idx = idx[:-1]  # 去掉最后一个
        
        #剩下boxes的信息存储在xx，yy中
        torch.index_select(x1, 0, idx, out=xx1)
        torch.index_select(y1, 0, idx, out=yy1)
        torch.index_select(x2, 0, idx, out=xx2)
        torch.index_select(y2, 0, idx, out=yy2)
        
        # 计算当前最大置信框与其他剩余框的交集，不知道clamp的同学确实容易被误导
        xx1 = torch.clamp(xx1, min=x1[i])  #max(x1,xx1)
        yy1 = torch.clamp(yy1, min=y1[i])  #max(y1,yy1)
        xx2 = torch.clamp(xx2, max=x2[i])  #min(x2,xx2)
        yy2 = torch.clamp(yy2, max=y2[i])  #min(y2,yy2)
        w.resize_as_(xx2)
        h.resize_as_(yy2)
        w = xx2 - xx1 #w=min(x2,xx2)−max(x1,xx1)
        h = yy2 - yy1 #h=min(y2,yy2)−max(y1,yy1)
        w = torch.clamp(w, min=0.0) #max(w,0)
        h = torch.clamp(h, min=0.0) #max(h,0)
        inter = w*h
        
		#计算当前最大置信框与其他剩余框的IOU
        # IoU = i / (area(a) + area(b) - i)
        rem_areas = torch.index_select(area, 0, idx)  # 剩余的框的面积
        union = rem_areas + area[i]- inter #并集
        IoU = inter/union  # 计算iou
        
        # 选出IoU <= overlap的boxes(注意le函数的使用)
        idx = idx[IoU.le(overlap)]
    return keep,          count
    	   #[num_remain], num_remain