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