当前位置：首页 > news >正文

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

news 2025/8/8 9:07:52

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_opsseed = 1107
iou_threshold = 0.35
box_num = 100000
cls_num = 80np.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] == 4keep = 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 paddleseed = 1107
iou_threshold = 0.35
box_num = 100000
cls_num = 80np.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] == 4keep = paddle.vision.ops.nms(boxes, iou_threshold, scores, idxs, cls_list)