nms3d

接口原型

mx_driving.nms3d(Tensor boxes, Tensor scores, float iou_threshold) -> Tensor

兼容：

mx_driving.detection.nms3d(Tensor boxes, Tensor scores, float iou_threshold) -> Tensor

mx_driving.detection.npu_nms3d(Tensor boxes, Tensor scores, float iou_threshold) -> Tensor

功能描述

3D非极大值抑制，在bev视角下剔除多个3d box交并比大于阈值的box。

参数说明

• boxes(Tensor)：框张量，数据类型为float32, float16。shape 为[N, 7]。7分别代表x, y, z, x_size, y_size, z_size, rz。
• scores(Tensor)：评分张量，数据类型为float32, float16。shape 为[N]。
• iou_threshold(float)：IoU阈值。

返回值

• output(Tensor)：NMS后的框张量，数据类型为int32。

支持的型号

• Atlas A2 训练系列产品

调用示例

import torch, torch_npu
from mx_driving import nms3d
boxes = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [3, 4, 5, 6, 7, 8, 9]], dtype=torch.float32).npu()
scores = torch.tensor([1, 2], dtype=torch.float32).npu()
out = nms3d(boxes, scores, 0.5)