torchvision.ops.nms与batched_nms

batched_nms() 和 nms()内置区别在于

torchvision.ops.nms

不区分类别对所有bbox进行过滤。如果有不同类别的bbox重叠的话会导致被过滤掉并不会分开计算。  

boxes (Tensor[N, 4])) – bounding boxes坐标. 格式：(x1, y1, x2, y2)
scores (Tensor[N]) – bounding boxes得分
iou_threshold (float) – IoU过滤阈值
返回NMS过滤后的bouding boxes索引（降序排列）
 

torchvision.ops.batched_nms

根据每个类别进行过滤，只对同一种类别进行计算IOU和阈值过滤。

boxes: Tensor, 预测框
scores: Tensor, 预测置信度
idxs: Tensor, 预测框类别
iou_threshold: float, IOU阈值

demo

import torchvision.ops as ops import torchb = torch.Tensor([[2,2,4,4], [1,1,5,5], [3,3,3.5,3.9]]) # bbox c = torch.Tensor([0,1,0]) # classes s = torch.Tensor([0.8,0.8,0.8]) # scoresops.batched_nms(b, s, c, 0.001) #运行结果 tensor([1, 2]) #[1,1,5,5], [3,3,3.5,3.9] bbox实际上是有包含关系的，但是类别不一样ops.nms(b, s, 0.001) # 运行结果 tensor([0]) # 可以看到 [1,1,5,5] 类别为1 但是被过滤掉了，只留下0号类别的[2,2,4,4]

总结

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

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