【算法手撕代码】手撕 IOU 与 NMS（Pytorch实现）

NMS

在目标检测后处理的时候，因为网络产生的候选框太多了，通常我们都需要使用 NMS (Non-maximum suppression) 非极大值抑制来保留最大值的候选框，也就是最合适的候选框。

NMS 算法的输入与返回值

输入值有 3 项：

• box
• score
• threshold

返回值：没有被抑制的 index 矩阵

其实仔细观察可以看到，目标检测的话，出了候选框的位置信息，网络还会输出候选框内部物体的分类信息，但是在 NMS 计算里却没有，这是为什么？

因为对于每个类别，都会进行一次 NMS，也就是有多少个类别就会进行多少次 NMS。

NMS 的流程

1. 根据 score 的大小进行排序（降序）
2. 从 score 的最大的 box 开始遍历，求其与别的 box 的 IOU 值
3. 如果 IOU 值大于 阈值 threshold，则去除，否则则保留。
4. 最后返回没有被抑制的 index 矩阵

def nms(self, bboxes, scores, thresh=0.5):
	# bboxe: [N, 4], scores: [N, ]
    x1, y1 = bboxes[:, 0], bboxes[:, 1]
    x2, y2 = bboxes[:, 1], bboxes[:, 1]
    areas = (x2 - x1) * (y2 - y1)  # bboxes area: [N, ]
    
    _, order = scores.sort(0, descending=True)  # 根据 score 进行降序排序，返回索引 order: [N, ]

    keep = []   # keep保留了NMS后留下的边框box
    
    while order.numel() > 0:
    
       	# 结束判断 ----------------------------------------------
        if order.numel() == 1: # 保留框只剩1个
            i = order.item()
            keep.append(i)
            break
        else: # 还有保留框没有NMS
            i = order[0].item() # 保留scores最大的那个框box[i]
            keep.append(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)  # inter: [N-1, ]
        union = areas[i] + areas[order[1: ]] - inter
        # 计算每一个框和当前框的IoU
        IoU = inter / union

        # 保留IoU小于threshold的边框索引 ---------------------------------------
        idx = (IoU <= thresh).nonzero().squeeze()
        if idx.numel() == 0:
            break
        order = order[idx+1]  # 这里+1是为了补充idx和order之间的索引差

    return torch.LongTensor(keep)  # 返回保留下的所有边框的索引

while 循环外面的计算 bbox 面积和根据 score 大小进行降序排序的部分都很好理解。

主要是在 while 循环内部，其实主要分成 3 个部分。

第一部分

        if order.numel() == 1: # 保留框只剩1个
            i = order.item()
            keep.append(i)
            break
        else: # 还有保留框没有NMS
            i = order[0].item() # 保留scores最大的那个框box[i]
            keep.append(i)

这段代码怎么理解呢？

因为是根据 score 的大小，将 bbox 从大到小遍历

• 如果最后遍历到最后，只剩下一个候选框的话，就直接将这个候选框的 index 保留，并且结束。
• 如果还有很多候选框没有遍历的话，就保留最大的候选框，然后遍历剩下的 order，去抑制和这个最大候选框接近的结果。

第二部分

这个部分是 IOU 的计算：

        # 计算 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)  # inter: [N-1, ]
        union = areas[i] + areas[order[1: ]] - inter  # areas[i]: [1, ], union: [N-1, ]
        # 计算每一个框和当前框的IoU
        IoU = inter / union # iou: [N-1, ]

torch.clamp 是怎么个用法，就是一个截断函数（虽然 clamp 直接翻译是“夹紧”）

• tensor_a.clamp(min=a1)：tensor_a 中小于 a1 的值都会被设成 a1。
• tensor_a.clamp(max=a2)：tensor_a 中大于 a2 的值都会被设成 a2。
• tensor_a.clamp(min=a1, max=a2)：tensor_a 中小于 a1 的值都会被设成 a1，大于 a2 的值都会被设成 a2。

当然也可以写成这样的形式： torch.clamp(tensor_a, min=a1, max=a2)

• xx1：就是所有大于 x1[i] 的 x1 值都会被保留，也就是上图中红点，小于 x1[i] 的都会被设成 x1[i]。
• xx2：就是说所有大于 x2[i] 的都会被设成 x2[i]，也就是上图的蓝点。

yy1 与 yy2 也是同理，这样就相当于保留了中间交集部分矩形的左上角和右上角的坐标，就可以用来计算交集矩形的面积 inter。这个交集肯定是不包含与自己的交集的，所有是 $[N-1,]$。

union 也是 $[N-1,]$， areas[i] 是 $[1,]$，但是在进行矩阵运算的时候，自动变成了 $[N-1,]$ 来进行计算。最后相除得到的 iou 也是 $[N-1,]$。

第三部分

        # 保留IoU小于threshold的边框索引 ---------------------------------------
        idx = (IoU <= thresh).nonzero().squeeze()
        if idx.numel() == 0:
            break
        order = order[idx+1]  # 这里+1是为了补充idx和order之间的索引差， order: [len(idx)+1, ]

tensor.nonzero()：会返回一个 [M, ] 的矩阵，每一行都是一个非零值的索引，说明一共有 M 个非零值。
tensor.squeeze()：会将维度为1的给去掉。

• 如果 idx 为 0，说明剩下的候选框全部抑制掉了，可以结束 NMS 了。
• idx 不为 0，则说明还有别的候选框没有被抑制掉，继续进行 NMS。

为什么 order[idx + 1] 要 idx + 1 呢?

因为计算 iou 的时候，会计算候选框和除了自己（第一个）别的候选框的 iou，所以得到的 iou 结果应该是 N-1 个，而 order 是 N 个，所以为了对齐，要加 1。

最后返回 keep

最后返回的 keep 理论上最大可能保留了所有的候选框，就是如果你的 threshold 设置得很大的时候，例如 0.99 之类的。

最少应该也能保留一个 score 值最大的候选框吧。

IOU

IOU 其实上面计算 NMS 其实已经涉及到了，原理是一样的，也是求交集和并集的比例。但是可能对于不同形式的输入，写法可能有些区别。

在求 NMS 中，是利用矩阵运算同时计算一个候选框和别的候选框的 IOU。

这里，我们也可以简单求两个候选框之间的 IOU， 也可以矩阵运算一次性求解多个候选框之间的 IOU。

矩阵运算计算 IOU 矩阵

def iou(self, box1, box2):
        N, M = box1.size(0), box2.size(0)

        lt = torch.max(  # 左上角的点
            box1[:, :2].unsqueeze(1).expand(N, M, 2),   # [N,2]->[N,1,2]->[N,M,2]
            box2[:, :2].unsqueeze(0).expand(N, M, 2),   # [M,2]->[1,M,2]->[N,M,2]
        )

        rb = torch.min(  # 右下角的点
            box1[:, 2:].unsqueeze(1).expand(N, M, 2),
            box2[:, 2:].unsqueeze(0).expand(N, M, 2),
        )

        wh = rb - lt  # [N,M,2]
        wh[wh < 0] = 0   # 两个box没有重叠区域
        inter = wh[:,:,0] * wh[:,:,1]   # [N,M]

        area1 = (box1[:, 2]-box1[:, 0]) * (box1[:, 3]-box1[:, 1])  # [N,]
        area2 = (box2[:, 2]-box2[:, 0]) * (box2[:, 3]-box2[:, 1])  # [M,]
        area1 = area1.unsqueeze(1).expand(N, M)  # [N, M]
        area2 = area2.unsqueeze(0).expand(N, M)  # [N, M]

        return inter / (area1 + area2 - inter)  # [N, M]

box1 和 box2 分别是 [N, 4] 和 [M, 4] 的维度，最后获得的 IOU 也是 [N, M] 的维度。

两个 box 的 IOU

def iou(self, box1, box2):
    lt_x, lt_y = max(box1[0], box2[0]), max(box1[1], box2[1])  
    rb_x, rb_y = min(box1[2], box2[2]), min(box1[3], box2[3])  

    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])

    w, h = rb_x - lt_x, rb_y - lt_y

    if w <= 0 or h <= 0:
        return 0  

    inter = w * h
    return inter / (area1 + area2 - inter)

box1 和 box2 只有一个维度，也就是 list 的长度，都是 4，所以最后获得的 IOU 是一个数值。