| 148 | class TestGridRayIntersection(unittest.TestCase): |
| 149 | |
| 150 | def test(self, b=1, m=100, max_size=20): |
| 151 | |
| 152 | ray_origins = torch.randn(b, m, 3) # (b, m, 3) |
| 153 | ray_directions = torch.nn.functional.normalize(torch.randn(b, m, 3), dim=-1) # (b, m, 3) |
| 154 | ray_radius = torch.rand(b,)*2 + 1e-3 # (b,) |
| 155 | grid_size = (torch.rand(b, 3) * (max_size-1) + 1).floor().long() # (b, 3) |
| 156 | total_cells = torch.prod(grid_size, dim=-1) |
| 157 | grid_center = torch.randn(b, 3) |
| 158 | grid_width = torch.rand(b, 3) * 20 + 10 |
| 159 | |
| 160 | stime = timer() |
| 161 | all_grid_idxs_gt = naive.grid_ray_intersection( |
| 162 | ray_origins=ray_origins, |
| 163 | ray_directions=ray_directions, |
| 164 | ray_radius=ray_radius, |
| 165 | grid_size=grid_size, |
| 166 | grid_center=grid_center, |
| 167 | grid_width=grid_width, |
| 168 | ) |
| 169 | time_python = timer() - stime |
| 170 | |
| 171 | stime = timer() |
| 172 | _all_grid_idxs = pr_cpp.grid_ray_intersection( |
| 173 | ray_origins, |
| 174 | ray_directions, |
| 175 | ray_radius, |
| 176 | grid_size, |
| 177 | grid_center, |
| 178 | grid_width, |
| 179 | ) |
| 180 | time_cpp = timer() - stime |
| 181 | |
| 182 | # convert list of idxmap to list of list of list |
| 183 | _all_grid_idxs = [c.get_table() for c in _all_grid_idxs] |
| 184 | all_grid_idxs = [] |
| 185 | for b in range(len(_all_grid_idxs)): |
| 186 | grid_idxs = [[] for i in range(m)] |
| 187 | for midx in _all_grid_idxs[b]: |
| 188 | grid_idxs[midx] = list(_all_grid_idxs[b][midx]) |
| 189 | all_grid_idxs.append(grid_idxs) |
| 190 | |
| 191 | print('grid_ray_intersection:') |
| 192 | print(f'python: {time_python * 1000.:.3f} ms') |
| 193 | print(f'cpp: {time_cpp * 1000.:.3f} ms ({time_python / time_cpp:.2f} speed up)') |
| 194 | |
| 195 | assert len(all_grid_idxs_gt) == len(all_grid_idxs) |
| 196 | for b in range(len(all_grid_idxs_gt)): |
| 197 | grid_idxs_gt = all_grid_idxs_gt[b] |
| 198 | grid_idxs = all_grid_idxs[b] |
| 199 | assert len(grid_idxs_gt) == len(grid_idxs) |
| 200 | for midx in range(len(grid_idxs_gt)): |
| 201 | gidxs_gt = np.array(sorted(list(grid_idxs_gt[midx]))) |
| 202 | gidxs = np.array(sorted(list(grid_idxs[midx]))) |
| 203 | assert np.allclose(gidxs_gt, gidxs) |
| 204 | |
| 205 | |
| 206 | # class TestGridRayIntersectionV2(unittest.TestCase): |