Pytorch实现nms (torchvision.ops.nms torchvision.ops.boxes.batched

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)相关推荐

torchvision.ops.nms与batched_nms
batched_nms() 和 nms()内置区别在于 torchvision.ops.nms 不区分类别对所有bbox进行过滤.如果有不同类别的bbox重叠的话会导致被过滤掉并不会分开计算. box ...
解决AttributeError: module ‘torchvision‘ has no attribute ‘ops‘
解决AttributeError: module 'torchvision' has no attribute 'ops' 报错信息 AttributeError: module 'torchvisi ...
‘mmdet\ops\nms\src/soft_nms_cpu.pyx‘ doesn‘t match any files
'mmdet\ops\nms\src/soft_nms_cpu.pyx' doesn't match any files 试试? pip install mmdet 自动安装mmcv win10上可以 ...
PyTorch实现非极大值抑制(NMS)
NMS即non maximum suppression即非极大抑制,顾名思义就是抑制不是极大值的元素,搜索局部的极大值.在最近几年常见的物体检测算法(包括rcnn.sppnet.fast-rcnn.f ...
[Pytorch系列-42]：工具集 - torchvision常见预训练模型的下载地址
作者主页(文火冰糖的硅基工坊):文火冰糖(王文兵)的博客_文火冰糖的硅基工坊_CSDN博客本文网址:https://blog.csdn.net/HiWangWenBing/article/detai ...
用PyTorch实现CIoU NMS,DIoU NMS和GIoU NMS
在目标检测的后处理阶段我们需要用到NMS算法. 而在NMS算法里有一个步是需要计算当前score最大的框和其他框的IoU大小的. 针对这一步,我们可以进行优化,改变IoU的计算方式. 目前经典的IoU ...
NMS系列(NMS,Soft-NMS,Weighted-NMS,IOU-Guided NMS,Softer-NMS,Adaptive NMS,DIOU NMS,Cluster NMS)
文章目录 NMS ConvNMS(2016) Soft-NMS(2017) Weighted-NMS(2017) IOU-Guided NMS (2018) Pure NMS Network(2017 ...
NMS、soft NMS、softer NMS与IOU-Guided NMS
提示:文章写完后,目录可以自动生成,如何生成可参考右边的帮助文档 NMS.soft NMS.softer NMS与IOU-Guided NMS 一.NMS 二.soft NMS 三.softer NM ...
目标检测后处理：从nms到softer nms
文章目录 1 NMS 1.1 动机 1.2 步骤 2 Soft-NMS 2.1 动机 2.2 算法思想 2.3 步骤 3 Softer-NMS 3.1 动机 3.1.1 现有方法的问题 3.1.2 本 ...

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

torchvision.ops.nms

计算多类别nms

手动实现

Pytorch实现nms (torchvision.ops.nms torchvision.ops.boxes.batched_nms)相关推荐

最新文章

热门文章