134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
import torch
|
|
|
|
from ..utils import ext_loader
|
|
|
|
ext_module = ext_loader.load_ext('_ext', [
|
|
'points_in_boxes_part_forward', 'points_in_boxes_cpu_forward',
|
|
'points_in_boxes_all_forward'
|
|
])
|
|
|
|
|
|
def points_in_boxes_part(points, boxes):
|
|
"""Find the box in which each point is (CUDA).
|
|
|
|
Args:
|
|
points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate
|
|
boxes (torch.Tensor): [B, T, 7],
|
|
num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz] in
|
|
LiDAR/DEPTH coordinate, (x, y, z) is the bottom center
|
|
|
|
Returns:
|
|
box_idxs_of_pts (torch.Tensor): (B, M), default background = -1
|
|
"""
|
|
assert points.shape[0] == boxes.shape[0], \
|
|
'Points and boxes should have the same batch size, ' \
|
|
f'but got {points.shape[0]} and {boxes.shape[0]}'
|
|
assert boxes.shape[2] == 7, \
|
|
'boxes dimension should be 7, ' \
|
|
f'but got unexpected shape {boxes.shape[2]}'
|
|
assert points.shape[2] == 3, \
|
|
'points dimension should be 3, ' \
|
|
f'but got unexpected shape {points.shape[2]}'
|
|
batch_size, num_points, _ = points.shape
|
|
|
|
box_idxs_of_pts = points.new_zeros((batch_size, num_points),
|
|
dtype=torch.int).fill_(-1)
|
|
|
|
# If manually put the tensor 'points' or 'boxes' on a device
|
|
# which is not the current device, some temporary variables
|
|
# will be created on the current device in the cuda op,
|
|
# and the output will be incorrect.
|
|
# Therefore, we force the current device to be the same
|
|
# as the device of the tensors if it was not.
|
|
# Please refer to https://github.com/open-mmlab/mmdetection3d/issues/305
|
|
# for the incorrect output before the fix.
|
|
points_device = points.get_device()
|
|
assert points_device == boxes.get_device(), \
|
|
'Points and boxes should be put on the same device'
|
|
if torch.cuda.current_device() != points_device:
|
|
torch.cuda.set_device(points_device)
|
|
|
|
ext_module.points_in_boxes_part_forward(boxes.contiguous(),
|
|
points.contiguous(),
|
|
box_idxs_of_pts)
|
|
|
|
return box_idxs_of_pts
|
|
|
|
|
|
def points_in_boxes_cpu(points, boxes):
|
|
"""Find all boxes in which each point is (CPU). The CPU version of
|
|
:meth:`points_in_boxes_all`.
|
|
|
|
Args:
|
|
points (torch.Tensor): [B, M, 3], [x, y, z] in
|
|
LiDAR/DEPTH coordinate
|
|
boxes (torch.Tensor): [B, T, 7],
|
|
num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz],
|
|
(x, y, z) is the bottom center.
|
|
|
|
Returns:
|
|
box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0.
|
|
"""
|
|
assert points.shape[0] == boxes.shape[0], \
|
|
'Points and boxes should have the same batch size, ' \
|
|
f'but got {points.shape[0]} and {boxes.shape[0]}'
|
|
assert boxes.shape[2] == 7, \
|
|
'boxes dimension should be 7, ' \
|
|
f'but got unexpected shape {boxes.shape[2]}'
|
|
assert points.shape[2] == 3, \
|
|
'points dimension should be 3, ' \
|
|
f'but got unexpected shape {points.shape[2]}'
|
|
batch_size, num_points, _ = points.shape
|
|
num_boxes = boxes.shape[1]
|
|
|
|
point_indices = points.new_zeros((batch_size, num_boxes, num_points),
|
|
dtype=torch.int)
|
|
for b in range(batch_size):
|
|
ext_module.points_in_boxes_cpu_forward(boxes[b].float().contiguous(),
|
|
points[b].float().contiguous(),
|
|
point_indices[b])
|
|
point_indices = point_indices.transpose(1, 2)
|
|
|
|
return point_indices
|
|
|
|
|
|
def points_in_boxes_all(points, boxes):
|
|
"""Find all boxes in which each point is (CUDA).
|
|
|
|
Args:
|
|
points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate
|
|
boxes (torch.Tensor): [B, T, 7],
|
|
num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz],
|
|
(x, y, z) is the bottom center.
|
|
|
|
Returns:
|
|
box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0.
|
|
"""
|
|
assert boxes.shape[0] == points.shape[0], \
|
|
'Points and boxes should have the same batch size, ' \
|
|
f'but got {boxes.shape[0]} and {boxes.shape[0]}'
|
|
assert boxes.shape[2] == 7, \
|
|
'boxes dimension should be 7, ' \
|
|
f'but got unexpected shape {boxes.shape[2]}'
|
|
assert points.shape[2] == 3, \
|
|
'points dimension should be 3, ' \
|
|
f'but got unexpected shape {points.shape[2]}'
|
|
batch_size, num_points, _ = points.shape
|
|
num_boxes = boxes.shape[1]
|
|
|
|
box_idxs_of_pts = points.new_zeros((batch_size, num_points, num_boxes),
|
|
dtype=torch.int).fill_(0)
|
|
|
|
# Same reason as line 25-32
|
|
points_device = points.get_device()
|
|
assert points_device == boxes.get_device(), \
|
|
'Points and boxes should be put on the same device'
|
|
if torch.cuda.current_device() != points_device:
|
|
torch.cuda.set_device(points_device)
|
|
|
|
ext_module.points_in_boxes_all_forward(boxes.contiguous(),
|
|
points.contiguous(),
|
|
box_idxs_of_pts)
|
|
|
|
return box_idxs_of_pts
|