| 54 | |
| 55 | class TestFindNeighborPointsOfRays(unittest.TestCase): |
| 56 | def _test( |
| 57 | self, |
| 58 | points: torch.Tensor, # (b, n, 3) |
| 59 | ray_origins: torch.Tensor, # (b, m, 3) |
| 60 | ray_directions: torch.Tensor, # (b, m, 3) |
| 61 | ray_radius: T.Union[torch.Tensor, float], # (b,) |
| 62 | grid_size: T.Union[torch.Tensor, int], # (b, 3) |
| 63 | grid_center: T.Union[torch.Tensor, float, None] = 0., # (b, 3) |
| 64 | grid_width: T.Union[torch.Tensor, float, None] = 1., # (b, 3) |
| 65 | # include_outside: bool = True, |
| 66 | ): |
| 67 | batch_size, n_rays, _ = ray_origins.shape |
| 68 | |
| 69 | stime = timer() |
| 70 | all_ray2pidxs = naive.find_neighbor_points_of_rays( |
| 71 | points=points, |
| 72 | ray_origins=ray_origins, |
| 73 | ray_directions=ray_directions, |
| 74 | ray_radius=ray_radius, |
| 75 | grid_size=grid_size, |
| 76 | grid_center=grid_center, |
| 77 | grid_width=grid_width, |
| 78 | # include_outside=include_outside, |
| 79 | ) |
| 80 | total_time = timer() - stime |
| 81 | |
| 82 | stime = timer() |
| 83 | all_ray2pidxs_gt = naive.find_neighbor_points_of_rays_brute_force( |
| 84 | points=points, |
| 85 | ray_origins=ray_origins, |
| 86 | ray_directions=ray_directions, |
| 87 | ray_radius=ray_radius, |
| 88 | grid_size=grid_size, |
| 89 | grid_center=grid_center, |
| 90 | grid_width=grid_width, |
| 91 | # include_outside=include_outside, |
| 92 | ) |
| 93 | total_time_gt = timer() - stime |
| 94 | print(f'new/gt = {total_time:1f}/{total_time_gt:1f} = {total_time/total_time_gt*100:.3f}%') |
| 95 | |
| 96 | assert len(all_ray2pidxs) == batch_size |
| 97 | assert len(all_ray2pidxs_gt) == batch_size |
| 98 | |
| 99 | for b in range(batch_size): |
| 100 | ray2pidxs = all_ray2pidxs[b] |
| 101 | ray2pidxs_gt = all_ray2pidxs_gt[b] |
| 102 | assert len(ray2pidxs) == n_rays |
| 103 | assert len(ray2pidxs_gt) == n_rays |
| 104 | |
| 105 | for m in range(n_rays): |
| 106 | pidxs = ray2pidxs[m] |
| 107 | pidxs_gt = ray2pidxs_gt[m] |
| 108 | pidxs = np.sort(np.array(pidxs)) |
| 109 | pidxs_gt = np.sort(np.array(pidxs_gt)) |
| 110 | assert len(pidxs) == len(pidxs_gt), f"[{b},{m}], {pidxs}, {pidxs_gt}" |
| 111 | assert np.allclose(pidxs, pidxs_gt), f"[{b},{m}], {pidxs}, {pidxs_gt}" |
| 112 | |
| 113 | def test( |