|
|
|
import torch
|
|
|
|
from ..utils import ext_loader
|
|
|
|
ext_module = ext_loader.load_ext('_ext', [
|
|
'iou3d_boxes_iou_bev_forward', 'iou3d_nms_forward',
|
|
'iou3d_nms_normal_forward'
|
|
])
|
|
|
|
|
|
def boxes_iou_bev(boxes_a, boxes_b):
|
|
"""Calculate boxes IoU in the Bird's Eye View.
|
|
|
|
Args:
|
|
boxes_a (torch.Tensor): Input boxes a with shape (M, 5).
|
|
boxes_b (torch.Tensor): Input boxes b with shape (N, 5).
|
|
|
|
Returns:
|
|
ans_iou (torch.Tensor): IoU result with shape (M, N).
|
|
"""
|
|
ans_iou = boxes_a.new_zeros(
|
|
torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
|
|
|
|
ext_module.iou3d_boxes_iou_bev_forward(boxes_a.contiguous(),
|
|
boxes_b.contiguous(), ans_iou)
|
|
|
|
return ans_iou
|
|
|
|
|
|
def nms_bev(boxes, scores, thresh, pre_max_size=None, post_max_size=None):
|
|
"""NMS function GPU implementation (for BEV boxes). The overlap of two
|
|
boxes for IoU calculation is defined as the exact overlapping area of the
|
|
two boxes. In this function, one can also set ``pre_max_size`` and
|
|
``post_max_size``.
|
|
|
|
Args:
|
|
boxes (torch.Tensor): Input boxes with the shape of [N, 5]
|
|
([x1, y1, x2, y2, ry]).
|
|
scores (torch.Tensor): Scores of boxes with the shape of [N].
|
|
thresh (float): Overlap threshold of NMS.
|
|
pre_max_size (int, optional): Max size of boxes before NMS.
|
|
Default: None.
|
|
post_max_size (int, optional): Max size of boxes after NMS.
|
|
Default: None.
|
|
|
|
Returns:
|
|
torch.Tensor: Indexes after NMS.
|
|
"""
|
|
assert boxes.size(1) == 5, 'Input boxes shape should be [N, 5]'
|
|
order = scores.sort(0, descending=True)[1]
|
|
|
|
if pre_max_size is not None:
|
|
order = order[:pre_max_size]
|
|
boxes = boxes[order].contiguous()
|
|
|
|
keep = torch.zeros(boxes.size(0), dtype=torch.long)
|
|
num_out = ext_module.iou3d_nms_forward(boxes, keep, thresh)
|
|
keep = order[keep[:num_out].cuda(boxes.device)].contiguous()
|
|
if post_max_size is not None:
|
|
keep = keep[:post_max_size]
|
|
return keep
|
|
|
|
|
|
def nms_normal_bev(boxes, scores, thresh):
|
|
"""Normal NMS function GPU implementation (for BEV boxes). The overlap of
|
|
two boxes for IoU calculation is defined as the exact overlapping area of
|
|
the two boxes WITH their yaw angle set to 0.
|
|
|
|
Args:
|
|
boxes (torch.Tensor): Input boxes with shape (N, 5).
|
|
scores (torch.Tensor): Scores of predicted boxes with shape (N).
|
|
thresh (float): Overlap threshold of NMS.
|
|
|
|
Returns:
|
|
torch.Tensor: Remaining indices with scores in descending order.
|
|
"""
|
|
assert boxes.shape[1] == 5, 'Input boxes shape should be [N, 5]'
|
|
order = scores.sort(0, descending=True)[1]
|
|
|
|
boxes = boxes[order].contiguous()
|
|
|
|
keep = torch.zeros(boxes.size(0), dtype=torch.long)
|
|
num_out = ext_module.iou3d_nms_normal_forward(boxes, keep, thresh)
|
|
return order[keep[:num_out].cuda(boxes.device)].contiguous()
|
|
|
