| 13 | class TestFindNeighborPoints(unittest.TestCase): |
| 14 | |
| 15 | def _test( |
| 16 | self, |
| 17 | type: str, # 'cpp' or 'cuda' |
| 18 | points: torch.Tensor, # (b, n, 3) |
| 19 | ray_origins: torch.Tensor, # (b, m, 3) |
| 20 | ray_directions: torch.Tensor, # (b, m, 3) |
| 21 | ray_radius: T.Union[torch.Tensor, float], # (b,) |
| 22 | grid_size: T.Union[torch.Tensor, int], # (b, 3) |
| 23 | grid_center: T.Union[torch.Tensor, float, None] = 0., # (b, 3) |
| 24 | grid_width: T.Union[torch.Tensor, float, None] = 1., # (b, 3) |
| 25 | k: int = 40, |
| 26 | ): |
| 27 | if type == 'cpp': |
| 28 | device = torch.device('cpu') |
| 29 | assert pr_utils.pr_cpp_loaded, f'pr_cpp is not loaded' |
| 30 | elif type == 'cuda': |
| 31 | if not torch.cuda.is_available(): |
| 32 | print('no available gpu, pass') |
| 33 | return |
| 34 | device = torch.device('cuda') |
| 35 | assert pr_utils.pr_cuda_loaded, f'pr_cuda is not loaded' |
| 36 | else: |
| 37 | raise NotImplementedError |
| 38 | |
| 39 | |
| 40 | batch_size, n_rays, _ = ray_origins.shape |
| 41 | |
| 42 | points = points.to(device=device) |
| 43 | ray_origins = ray_origins.to(device=device) |
| 44 | ray_directions = ray_directions.to(device=device) |
| 45 | ray_radius = ray_radius.to(device=device) if isinstance(ray_radius, torch.Tensor) else ray_radius |
| 46 | grid_size = grid_size.to(device=device) if isinstance(grid_size, torch.Tensor) else grid_size |
| 47 | grid_center = grid_center.to(device=device) if isinstance(grid_center, torch.Tensor) else grid_center |
| 48 | grid_width = grid_width.to(device=device) if isinstance(grid_width, torch.Tensor) else grid_width |
| 49 | |
| 50 | stime = timer() |
| 51 | all_ray2pidxs_gt = naive.find_neighbor_points_of_rays_brute_force( |
| 52 | #all_ray2pidxs_gt = naive.find_neighbor_points_of_rays( |
| 53 | points=points, |
| 54 | ray_origins=ray_origins, |
| 55 | ray_directions=ray_directions, |
| 56 | ray_radius=ray_radius, |
| 57 | grid_size=grid_size, |
| 58 | grid_center=grid_center, |
| 59 | grid_width=grid_width, |
| 60 | ) |
| 61 | total_time_gt = timer() - stime |
| 62 | |
| 63 | stime = timer() |
| 64 | all_ray2pidxs = pr_utils.find_neighbor_points_of_rays( |
| 65 | points=points, |
| 66 | ray_origins=ray_origins, |
| 67 | ray_directions=ray_directions, |
| 68 | ray_radius=ray_radius, |
| 69 | grid_size=grid_size, |
| 70 | grid_center=grid_center, |
| 71 | grid_width=grid_width, |
| 72 | ) |