PyTorch实现非极大值抑制(NMS)

NMS即non maximum suppression即非极大抑制，顾名思义就是抑制不是极大值的元素，搜索局部的极大值。在最近几年常见的物体检测算法（包括rcnn、sppnet、fast-rcnn、faster-rcnn等）中，最终都会从一张图片中找出很多个可能是物体的矩形框，然后为每个矩形框为做类别分类概率。本文来通过Pytorch实现NMS算法。

如果你在做计算机视觉(特别是目标检测)，你肯定会听说过非极大值抑制(nms)。网上有很多不错的文章给出了适当的概述。简而言之，非最大抑制使用一些启发式方法减少了输出边界框的数量，例如交叉除以并集(iou)。

在PyTorch的文档中说：NMS 迭代地删除与另一个（得分较高）框的 IoU 大于 iou_threshold 的得分较低的框。

为了研究其如何工作，让我们加载一个图像并创建边界框

 from PIL import Image
 import torch
 import matplotlib.pyplot as plt
 import numpy as np
 
 # credit https://i0.wp.com/craffic.co.in/wp-content/uploads/2021/02/ai-remastered-rick-astley-never-gonna-give-you-up.jpg?w=1600&ssl=1
 img = Image.open("./samples/never-gonna-give-you-up.webp")
 img
1
2
3
4
5
6
7
8

我们手动创建 两个框，一个人脸，一个话筒

 original_bboxes = torch.tensor([
     # head
     [ 565, 73, 862, 373],
     # mic
     [807, 309, 865, 434]
 ]).float()
 
 w, h = img.size
 # we need them in range [0, 1]
 original_bboxes[...,0] /= h
 original_bboxes[...,1] /= w
 original_bboxes[...,2] /= h
 original_bboxes[...,3] /= w
1
2
3
4
5
6
7
8
9
10
11
12
13

这些bboxes 都是在[0,1]范围内的，虽然这不是必需的，但当有多个类时，这是非常有用的(我们稍后将看到为什么)。

 from torchvision.utils import draw_bounding_boxes
 from torchvision.transforms.functional import to_tensor
 from typing import List
 
 def plot_bboxes(img : Image.Image, bboxes: torch.Tensor, *args, **kwargs) -> plt.Figure:
     w, h = img.size
     # from [0, 1] to image size
     bboxes = bboxes.clone()
     bboxes[...,0] *= h
     bboxes[...,1] *= w
     bboxes[...,2] *= h
     bboxes[...,3] *= w
     fig = plt.figure()
     img_with_bboxes = draw_bounding_boxes((to_tensor(img) * 255).to(torch.uint8), bboxes, *args, **kwargs, width=4)
     return plt.imshow(img_with_bboxes.permute(1,2,0).numpy())
 
 plot_bboxes(img, original_bboxes, labels=["head", "mic"])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

为了说明，我们添加一些重叠的框

 max_bboxes = 3
 scaling = torch.tensor([1, .96, .97, 1.02])
 shifting = torch.tensor([0, 0.001, 0.002, -0.002])
 
 # broadcasting magic (2, 1, 4) * (1, 3, 1)
 bboxes = (original_bboxes[:,None,:] * scaling[..., None] + shifting[..., None]).view(-1, 4)
 
 plot_bboxes(img, bboxes, colors=[*["yellow"] * 4, *["blue"] * 4], labels=[*["head"] * 4, *["mic"] * 4])

1
2
3
4
5
6
7
8
9

现在可以看到，有6个bboxes ，这里我们还需要定义一个分数，这通常由模型输出。

 scores = torch.tensor([
     0.98, 0.85, 0.5, 0.2, # for head
     1, 0.92, 0.3, 0.1 # for mic
 ])
1
2
3
4

我们标签的分类，0代表人脸，1代表麦克风

 labels = torch.tensor([0,0,0,0,1,1,1,1])
1

最后，让我们排列一下这些数据

 perm = torch.randperm(scores.shape[0])
 bboxes = bboxes[perm]
 scores = scores[perm]
 labels = labels[perm]
1
2
3
4

让我们看看结果

 plot_bboxes(img, bboxes, 
             colors=["yellow" if el.item() == 0 else "blue" for el in labels], 
             labels=["head" if el.item()  == 0 else "mic" for el in labels]
            )
1
2
3
4

好了，这样我们模拟了模型的输出了，下面进入正题。

NMS是通过迭代删除低分数重叠的边界框来工作的。步骤如下。

bboxes are sorted by score in decreasing order
init a vector keep with ones
for i in len(bboxes):
    # was suppressed
    if keep[i] == 0:
        continue
    # compare with all the others
    for j in len(bbox):
        if keep[j]:
            if (iou(bboxes[i], bboxes[j]) > iou_threshold):
                keep[j] = 0

return keep
1
2
3
4
5
6
7
8
9
10
11
12
13

我们的Pytorch实现，采用三个参数（这实际上是从pytorch的文档中复制和粘贴的）：

• box (Tensor[N, 4])) – 用于执行 NMS 的框。它们应该是 (x1, y1, x2, y2) 格式，0 <= x1 < x2 和 0 <= y1 < y2。
• score (Tensor[N]) – 每个box 的得分
• iou_threshold (float) – 丢弃所有 IoU > iou_threshold 的框
• 返回值是非抑制边界框的索引

from torchvision.ops.boxes import box_iou

def nms(bboxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
    order = torch.argsort(-scores)
    indices = torch.arange(bboxes.shape[0])
    keep = torch.ones_like(indices, dtype=torch.bool)
    for i in indices:
        if keep[i]:
            bbox = bboxes[order[i]]
            iou = box_iou(bbox[None,...],(bboxes[order[i + 1:]]) * keep[i + 1:][...,None])
            overlapped = torch.nonzero(iou > iou_threshold)
            keep[overlapped + i + 1] = 0
    return order[keep]
1
2
3
4
5
6
7
8
9
10
11
12
13

让我们详细说明下这个参数：

order = scores.argsort()
1

根据分数得到排序的指标

indices = torch.arange(bboxes.shape[0])
1

创建用于迭代bboxes的索引 indices

keep = torch.ones_like(indices, dtype=torch.bool)
1

keep是用于判断一个bbox是否应该保留的向量，如果Keep [i] == 1，则bboxes[order[i]]不被抑制

for i in indices:
    ...
1
2

for循环遍历所有的box,如果当前box未被抑制，则keep[i] = 1

bbox = bboxes[order[i]]]
1

来通过已排序的位置获取bbox

iou = box_iou(bbox[None,...], (bboxes[order[i + 1:]]) * keep[i + 1:][...,None])
1

计算当前bbox和所有其他候选bbox之间的iou。这将把所有抑制框设置为零(因为keep将等于0)

(bboxes ...)[order[i + 1:]]
1

在排序的顺序中与后面所有的框进行比较，因为需要跳过当前的框，所以这里是i+ 1,

overlapped = torch.nonzero(iou > iou_threshold)
keep[overlapped + i + 1] = 0
1
2

计算和选择iou大于iou_threshold的索引。

我们之前对bboxes进行了切片，(bboxes…)[i + 1:])，所以我们需要添加这些索引的偏移量，这就是后面+ i + 1的原因。

最后返回order[keep]，这样映射回原始的box索引(未排序)，这样一个简单的函数就执行完成了。

让我们看看结果

nms_indices = nms(bboxes, scores, .45)
plot_bboxes(img, 
            bboxes[nms_indices],
            colors=["yellow" if el.item() == 0 else "blue" for el in labels[nms_indices]], 
            labels=["head" if el.item()  == 0 else "mic" for el in labels[nms_indices]]
           )
1
2
3
4
5
6

因为有多个类，所以需要让nms在同一个类中计算iou。还记得上面我们提到的在[0,1]之间吗?可以给它们添加标签，把不同类的框区分开。

nms_indices = nms(bboxes + labels[..., None], scores, .45)
plot_bboxes(img, 
            bboxes[nms_indices],
            colors=["yellow" if el.item() == 0 else "blue" for el in labels[nms_indices]], 
            labels=["head" if el.item()  == 0 else "mic" for el in labels[nms_indices]]
           )
1
2
3
4
5
6

如果我们将阈值更改为0.1，就得到了下图

让我们对比下pytorch官方的实现：

from torchvision.ops.boxes import nms as torch_nms
nms_indices = torch_nms(bboxes + labels[..., None], scores, .45)
plot_bboxes(img, 
            bboxes[nms_indices],
            colors=["yellow" if el.item() == 0 else "blue" for el in labels[nms_indices]], 
            labels=["head" if el.item()  == 0 else "mic" for el in labels[nms_indices]]
           )
1
2
3
4
5
6
7

结果是一样的。然我们看看时间：

%%timeit
nms(bboxes + labels[..., None], scores, .45)
#534 µs ± 22.1 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%%timeit
torch_nms(bboxes + labels[..., None], scores, .45)
#54.4 µs ± 3.29 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
1
2
3
4
5
6
7

我们的实现慢了大约10倍，哈，这个结果很正常，因为我们我们没有使用自定义的cpp内核!但是这并不代表我们的实现没有用，因为手写代码我们完全了解了NMS的工作原理，这是本文的真正意义，总之在这篇文章中我们看到了如何在PyTorch中实现非最大抑制，这对你了解目标检测的相关知识是非常有帮助的。

https://avoid.overfit.cn/post/1ffeb08f8ea4494cb992b0ad05db174b

作者：Francesco Zuppichini