def npu_batched_multiclass_nms(multi_bboxes,multi_scores,score_thr=0.05,nms_thr=0.45,max_num=50,score_factors=None):
使用NPU API的批量多类bbox NMS。
from torch_npu.contrib.function import npu_batched_multiclass_nms det_bboxes, det_labels = npu_batched_multiclass_nms(boxes, scores, score_thr=0.3, nms_thr=0.5, max_num=3)
在动态shape条件下,由于NPU op的限制,最多支持20个类别(nmsed_classes)和10000个框(nmsed_boxes)。
Tuple - (bboxes, labels),shape为(bs, k, 5)和(bs, k, 1)的张量。标签以0为基础。
>>> boxes = torch.randint(1, 255, size=(4, 200, 80, 4)) >>> scores = torch.randn(4, 200, 81) >>> boxes = boxes.npu().half() >>> scores = scores.npu().half() >>> det_bboxes, det_labels = npu_batched_multiclass_nms(boxes, scores, score_thr=0.3, nms_thr=0.5, max_num=3) >>> expedt_det_bboxes = torch.tensor([[[221.8750, 60.0000, 183.0000, 22.0000, 3.8867], [167.0000, 250.0000, 136.0000, 144.0000, 3.6445], [ 45.9688, 147.0000, 67.0000, 241.8750, 3.4844]], [[ 5.0000, 178.0000, 243.8750, 138.0000, 3.7344], [238.0000, 132.0000, 47.0000, 84.0000, 3.6836], [ 32.0000, 110.0000, 131.0000, 73.0000, 3.6309]], [[111.9375, 120.9375, 54.0000, 231.0000, 3.9219], [147.0000, 162.0000, 78.0000, 1.0010, 3.9219], [157.0000, 118.0000, 57.0000, 115.0000, 3.6523]], [[ 80.0000, 126.9375, 54.0000, 246.8750, 3.7344], [ 31.0000, 253.8750, 19.0000, 138.0000, 3.6328], [ 54.0000, 253.8750, 78.0000, 75.0000, 3.5586]]],dtype=torch.float16)