paddle.vision 与 torchvision 中的box NMS使用方式

最新推荐文章于 2023-06-16 14:46:01 发布

氵文大师

最新推荐文章于 2023-06-16 14:46:01 发布

阅读量890

点赞数 1

分类专栏： PaddleDetection paddlepaddle历险记文章标签： paddle numpy python

本文链接：https://blog.csdn.net/HaoZiHuang/article/details/129219828

版权

paddlepaddle历险记同时被 2 个专栏收录

46 篇文章 11 订阅

订阅专栏

PaddleDetection

13 篇文章 1 订阅

订阅专栏

文章比较了torchvision中的batched_nms函数和paddle.vision中的nms操作在处理边界框非极大值抑制时的用法。两者都用于去除检测框重叠部分，但paddle.vision的nms额外需要categories参数。通过示例代码展示了如何在两个库中使用这些函数。

摘要由CSDN通过智能技术生成

torchvision 中有多个用于计算 BBox NMS 的 API, 在本篇氵文中, 使用

torchvision.ops.boxes.batched_nms

paddle.vision 中通过 paddle.vision.ops.nms 来进行多个 Box 的 NMS 操作

1. torchvision 中 batched_nms 操作

torchvision batched_nms

def batched_nms(
    boxes: torch.Tensor,
    scores: torch.Tensor,
    idxs: torch.Tensor,
    iou_threshold: float,
) -> torch.Tensor

传入的参数分别为

边界框boxes, 格式[x1, y1, x2, y2]，shape 为 [num, 4]，dtype 为 float
置信度scores, shape 为 [num,]，dtype 为 float
类别idxs, shape 为 [num,]，dtype 为 int

来举个例子:

import numpy as np
import torch
from torchvision.ops import boxes as box_ops

seed = 1107
iou_threshold = 0.35
box_num = 100000
cls_num = 80

np.random.seed(seed)

boxes = np.random.rand(box_num, 4).astype("float32")
boxes = torch.from_numpy(boxes)

scores = np.random.rand(box_num).astype("float32")
scores = torch.from_numpy(scores)

idxs = np.random.randint(0, cls_num, size=(box_num,))
idxs = torch.from_numpy(idxs)


assert boxes.shape[-1] == 4

keep = box_ops.batched_nms(boxes.float(), scores, idxs, iou_threshold)

2. paddle.vision.ops.nms 操作

paddle.vision.ops.nms(
			boxes, 
			iou_threshold=0.3, 
			scores=None, 
			category_idxs=None, 
			categories=None, 
			top_k=None)

boxes、iou_threshold、scores 和 category_idxs 等参数和上述 torchvision 中 batched_nms 参数一样
不同的是 paddle 中还需要 categories 参数，(其实没什么必要)

category_idxs 是每个 bbox 的类别，而 categories 是一共的类别

比如 COCO 一共80类，则：

categories = paddle.arange(80)

Paddle 中的例子：

import numpy as np
import paddle

seed = 1107
iou_threshold = 0.35
box_num = 100000
cls_num = 80

np.random.seed(seed)

boxes = np.random.rand(box_num, 4).astype("float32")
boxes = paddle.to_tensor(boxes)

scores = np.random.rand(box_num).astype("float32")
scores = paddle.to_tensor(scores)

idxs = np.random.randint(0, cls_num, size=(box_num,))
idxs = paddle.to_tensor(idxs)

cls_list = paddle.arange(0, cls_num)


assert boxes.shape[-1] == 4

keep = paddle.vision.ops.nms(boxes, iou_threshold, scores, idxs, cls_list)