torch_npu.npu_nms_with_mask

torch_npu.npu_nms_with_mask(input, iou_threshold) -> (Tensor, Tensor, Tensor)

生成值0或1，用于nms算子确定有效位。

• 参数解释：
  • input (Tensor) - 输入张量
  • iou_threshold (Scalar) - 阈值。如果超过此阈值，则值为1，否则值为0。
• 返回值:
  • selected_boxes (Tensor) - shape为[N,5]的2D张量，表示filtered box，包括proposal box和相应的置信度分数。
  • selected_idx (Tensor) - shape为[N]的1D张量，表示输入建议框的index。
  • selected_mask (Tensor) - shape为[N]的1D张量，判断输出建议框是否有效。
• 约束条件：

  输入box_scores的2nd-dim必须等于8。

• 示例：

  >>> input = torch.tensor([[0.0, 1.0, 2.0, 3.0, 0.6], [6.0, 7.0, 8.0, 9.0, 0.4]], dtype=torch.float16).to("npu")
  >>> iou_threshold = 0.5
  >>> output1, output2, output3, = torch_npu.npu_nms_with_mask(input, iou_threshold)
  >>> output1
  tensor([[0.0000, 1.0000, 2.0000, 3.0000, 0.6001],
          [6.0000, 7.0000, 8.0000, 9.0000, 0.3999]], device='npu:0',      dtype=torch.float16)
  >>> output2
  tensor([0, 1], device='npu:0', dtype=torch.int32)
  >>> output3
  tensor([1, 1], device='npu:0', dtype=torch.uint8)