torchvision.ops.nms

torchvision中已經(jīng)有了nms

torchvision.ops.nms(boxes, scores, iou_threshold)

• boxes (Tensor[N, 4])) – bounding boxes坐標. 格式：(x1, y1, x2, y2)
• scores (Tensor[N]) – bounding boxes得分
• iou_threshold (float) – IoU過濾閾值

返回NMS過濾后的bouding boxes索引（降序排列）

import torch import torchvisionbox = torch.tensor([[2,3.1,7,5],[3,4,8,4.8],[4,4,5.6,7],[0.1,0,8,1]]) score = torch.tensor([0.5, 0.3, 0.2, 0.4])output = torchvision.ops.nms(boxes=box, scores=score, iou_threshold=0.3) print('IOU of bboxes:') iou = torchvision.ops.box_iou(box,box) print(iou) print(output)

計算多類別nms

torchvision.ops.boxes.batched_nms() 或 torchvision.ops.batched_nms()

torchvision.ops.boxes.batched_nms(boxes, scores, lvl, nms_thresh)

不同版本接口不太一樣

• batched_nms（）：根據(jù)每個類別進行過濾，只對同一種類別進行計算IOU和閾值過濾。
• nms（）：不區(qū)分類別對所有bbox進行過濾。如果有不同類別的bbox重疊的話會導致被過濾掉并不會分開計算。

import torchvision.ops as ops import torch from torchvision.ops import boxes as box_opsboxes = torch.Tensor([[2,2,4,4], [1,1,5,5], [3,3,3.5,3.9]]) # bbox classes = torch.Tensor([0,1,0]) # classes scores = torch.Tensor([0.8,0.8,0.8]) # scores# ops.batched_nms(b, s, c, 0.001) print(box_ops.batched_nms(boxes, scores, classes, 0.001)) #運行結果 tensor([0, 1]) #[2,2,4,4], [1,1,5,5] bbox實際上是有包含關系的，但是類別不一樣print(ops.nms(boxes, scores, 0.001)) # 運行結果 tensor([0]) # 可以看到 [1,1,5,5] 類別為1 但是被過濾掉了，只留下0號類別的[2,2,4,4]

?

?手動實現(xiàn)

from torch import Tensor import torch import torchvisiondef box_area(boxes: Tensor) -> Tensor:"""Computes the area of a set of bounding boxes, which are specified by its(x1, y1, x2, y2) coordinates.Arguments:boxes (Tensor[N, 4]): boxes for which the area will be computed. Theyare expected to be in (x1, y1, x2, y2) formatReturns:area (Tensor[N]): area for each box"""return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:"""Return intersection-over-union (Jaccard index) of boxes.Both sets of boxes are expected to be in (x1, y1, x2, y2) format.Arguments:boxes1 (Tensor[N, 4])boxes2 (Tensor[M, 4])Returns:iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2"""area1 = box_area(boxes1) # 每個框的面積 (N,)area2 = box_area(boxes2) # (M,)lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] # N中一個和M個比較； 所以由N，M 個rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]wh = (rb - lt).clamp(min=0) # [N,M,2] #小于0的為0 clamp 鉗；夾鉗；inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] iou = inter / (area1[:, None] + area2 - inter)return iou # NxM， boxes1中每個框和boxes2中每個框的IoU值；def nms(boxes: Tensor, scores: Tensor, iou_threshold: float):""":param boxes: [N, 4]， 此處傳進來的框，是經(jīng)過篩選（NMS之前選取過得分TopK）之后， 在傳入之前處理好的；:param scores: [N]:param iou_threshold: 0.7:return:"""keep = [] # 最終保留的結果， 在boxes中對應的索引；idxs = scores.argsort() # 值從小到大的 索引while idxs.numel() > 0: # 循環(huán)直到null； numel()： 數(shù)組元素個數(shù)# 得分最大框?qū)乃饕? 以及對應的坐標max_score_index = idxs[-1]max_score_box = boxes[max_score_index][None, :] # [1, 4]keep.append(max_score_index)if idxs.size(0) == 1: # 就剩余一個框了；breakidxs = idxs[:-1] # 將得分最大框 從索引中刪除； 剩余索引對應的框 和 得分最大框 計算IoU；other_boxes = boxes[idxs] # [?, 4]ious = box_iou(max_score_box, other_boxes) # 一個框和其余框比較 1XMidxs = idxs[ious[0] <= iou_threshold]keep = idxs.new(keep) # Tensorreturn keepbox = torch.tensor([[2,3.1,7,5],[3,4,8,4.8],[4,4,5.6,7],[0.1,0,8,1]]) score = torch.tensor([0.5, 0.3, 0.2, 0.4])output = nms(boxes=box, scores=score, iou_threshold=0.3) print('IOU of bboxes:') iou = torchvision.ops.box_iou(box,box) print(iou) print(output)

用numpy同樣可以實現(xiàn)

總結

以上是生活随笔為你收集整理的Pytorch实现nms (torchvision.ops.nms torchvision.ops.boxes.batched_nms)的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。