Pytorch实现nms (torchvision.ops.nms torchvision.ops.boxes.batched_nms)

torchvision.ops.nms

torchvision中已经有了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（）：根据每个类别进行过滤，只对同一种类别进行计算IOU和阈值过滤。
• nms（）：不区分类别对所有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]

 

 手动实现

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]， 此处传进来的框，是经过筛选（NMS之前选取过得分TopK）之后， 在传入之前处理好的；:param scores: [N]:param iou_threshold: 0.7:return:"""keep = [] # 最终保留的结果， 在boxes中对应的索引；idxs = scores.argsort() # 值从小到大的 索引while idxs.numel() > 0: # 循环直到null； numel()： 数组元素个数# 得分最大框对应的索引, 以及对应的坐标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同样可以实现

总结

以上是生活随笔为你收集整理的Pytorch实现nms (torchvision.ops.nms torchvision.ops.boxes.batched_nms)的全部内容，希望文章能够帮你解决所遇到的问题。

如果觉得生活随笔网站内容还不错，欢迎将生活随笔推荐给好友。