PyTorch實現非極大值抑制(NMS)

NMS即non maximum suppression即非極大抑制，顧名思義就是抑制不是極大值的元素，搜索局部的極大值。在最近幾年常見的物體檢測算法（包括rcnn、sppnet、fast-rcnn、faster-rcnn等）中，最終都會從一張圖片中找出很多個可能是物體的矩形框，然后為每個矩形框為做類別分類概率。

如果你在做計算機視覺(特別是目標檢測)，你肯定會聽說過非極大值抑制(nms)。網上有很多不錯的文章給出了適當的概述。簡而言之，非最大抑制使用一些啟發(fā)式方法減少了輸出邊界框的數量，例如交叉除以并集(iou)。

在PyTorch的文檔中說：NMS 迭代地刪除與另一個（得分較高）框的 IoU 大于 iou_threshold 的得分較低的框。

為了研究其如何工作，讓我們加載一個圖像并創(chuàng)建邊界框

 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

我們手動創(chuàng)建 兩個框，一個人臉，一個話筒

 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

這些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"])

為了說明，我們添加一些重疊的框

 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])

現在可以看到，有6個bboxes ，這里我們還需要定義一個分數，這通常由模型輸出。

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

我們標簽的分類，0代表人臉，1代表麥克風。

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

最后，讓我們排列一下這些數據：

 perm = torch.randperm(scores.shape[0]) bboxes = bboxes[perm] scores = scores[perm]
 labels = labels[perm]

讓我們看看結果：

 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]
            )

好了，這樣我們模擬了模型的輸出了，下面進入正題。

NMS是通過迭代刪除低分數重疊的邊界框來工作的。步驟如下。

bboxes are sorted by score in decreasing orderinit a vector keep with onesfor 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

我們的Pytorch實現，采用三個參數（這實際上是從pytorch的文檔中復制和粘貼的）：

• box (Tensor[N, 4])) – 用于執(zhí)行 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_ioudef 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]

讓我們詳細說明下這個參數：

order = scores.argsort()

根據分數得到排序的指標

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

創(chuàng)建用于迭代bboxes的索引 indices

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

keep是用于判斷一個bbox是否應該保留的向量，如果Keep [i] == 1，則bboxes[order[i]]不被抑制

for i in indices:    ...

for循環(huán)遍歷所有的box,如果當前box未被抑制，則keep[i] = 1

bbox = bboxes[order[i]]]

來通過已排序的位置獲取bbox

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

計算當前bbox和所有其他候選bbox之間的iou。這將把所有抑制框設置為零(因為keep將等于0)

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

在排序的順序中與后面所有的框進行比較，因為需要跳過當前的框，所以這里是i+ 1,

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

計算和選擇iou大于iou_threshold的索引。

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

最后返回order[keep]，這樣映射回原始的box索引(未排序)，這樣一個簡單的函數就執(zhí)行完成了。

讓我們看看結果。

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]]
           )

因為有多個類，所以需要讓nms在同一個類中計算iou。還記得上面我們提到的在[0,1]之間嗎?可以給它們添加標簽，把不同類的框區(qū)分開。

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]]
           )

如果我們將閾值更改為0.1，就得到了下圖

讓我們對比下pytorch官方的實現：

from torchvision.ops.boxes import nms as torch_nmsnms_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]]
           )

結果是一樣的。然我們看看時間：

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

%%timeittorch_nms(bboxes + labels[..., None], scores, .45)
#54.4 μs ± 3.29 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

我們的實現慢了大約10倍，哈，這個結果很正常，因為我們我們沒有使用自定義的cpp內核!但是這并不代表我們的實現沒有用，因為手寫代碼我們完全了解了NMS的工作原理，這是本文的真正意義，總之在這篇文章中我們看到了如何在PyTorch中實現非最大抑制，這對你了解目標檢測的相關知識是非常有幫助的。