比较与torchvision.ops.nms的差异

torchvision.ops.nms

torchvision.ops.nms(boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float)

更多内容详见torchvision.ops.nms。

mindspore.ops.NMSWithMask

class mindspore.ops.NMSWithMask(iou_threshold=0.5)(bboxes)

更多内容详见mindspore.ops.NMSWithMask。

差异对比

PyTorch：非极大值抑制算法。参数 boxes 与 scores 维度分别为(N, 4)和(N, 1)，表示bbox与score。

MindSpore：非极大值抑制算法，参数 bboxes 维度为(N, 5)，以(x0、y0、x1、y1, score)形式表示框与分数。Ascend平台最大支持2864个输入框。

分类	子类	PyTorch	MindSpore	差异
参数	参数1	boxes	-	边界框，此参数位于MindSpore算子的输入参数列表中
	参数2	scores	-	边界框的分数，此参数位于MindSpore算子的输入参数列表中
	参数3	iou_threshold	iou_threshold	指定删除框的IOU的阈值
输入	输入1	-	bboxes	边界框与分数
输出	输出1	indices	-	PyTorch表示为NMS后的bbox索引
	输出2	-	output_boxes	按score排序后的bbox列表
	输出3	-	output_idx	output_boxes 的索引列表
	输出4	-	selected_mask	表示NMS后的box掩码，True为剩余bbox，False为丢弃的bbox

代码示例

# PyTorch
import torch
import torchvision as tv
import numpy as np

boxes = np.array([
    [0, 0, 4, 4],
    [0, 0, 3, 3],
    [0, 0, 2, 2],
    [0, 0, 1, 1]
]).astype(np.float32)

scores = np.array([0.8, 0.7, 0.6, 0.5]).astype(np.float32)

iou_threshold = 0.4

boxes_t = torch.from_numpy(boxes)
scores_t = torch.from_numpy(scores)

remain_boxes = tv.ops.nms(boxes_t, scores_t, iou_threshold)
print(remain_boxes)
# Out: tensor([0, 2, 3])

# MindSpore
import mindspore as ms
from mindspore import ops

box_with_score = np.column_stack((boxes, scores))
box_with_score_m = ms.Tensor(box_with_score)

output_boxes, output_idx, selected_mask = ops.NMSWithMask(iou_threshold)(box_with_score_m)
print(selected_mask)
# Out: [True False True True]