| 462 | |
| 463 | |
| 464 | class TestFindNeighborPointsOfRays(unittest.TestCase): |
| 465 | def _test( |
| 466 | self, |
| 467 | points: torch.Tensor, # (b, n, 3) |
| 468 | ray_origins: torch.Tensor, # (b, m, 3) |
| 469 | ray_directions: torch.Tensor, # (b, m, 3) |
| 470 | ray_radius: T.Union[torch.Tensor, float], # (b,) |
| 471 | grid_size: T.Union[torch.Tensor, int], # (b, 3) |
| 472 | grid_center: T.Union[torch.Tensor, float, None] = 0., # (b, 3) |
| 473 | grid_width: T.Union[torch.Tensor, float, None] = 1., # (b, 3) |
| 474 | ): |
| 475 | batch_size, n_rays, _ = ray_origins.shape |
| 476 | n_points = points.size(1) |
| 477 | |
| 478 | if isinstance(grid_center, float): |
| 479 | grid_center = torch.tensor(grid_center, dtype=points.dtype, device=points.device) |
| 480 | grid_center = grid_center.view(1, 1).expand(batch_size, 3) # (b, 3) |
| 481 | if isinstance(grid_width, (float, int)): |
| 482 | grid_width = torch.tensor(grid_width, dtype=points.dtype, device=points.device) |
| 483 | grid_width = grid_width.view(1, 1).expand(batch_size, 3) # (b, 3) |
| 484 | if isinstance(grid_size, (float, int)): |
| 485 | grid_size = torch.tensor(grid_size, dtype=torch.long, device=points.device) |
| 486 | grid_size = grid_size.view(1, 1).expand(batch_size, 3) # (b, 3) |
| 487 | grid_size = grid_size.long() # (b, 3) |
| 488 | |
| 489 | points = points.cuda() |
| 490 | ray_origins = ray_origins.cuda() |
| 491 | ray_directions = ray_directions.cuda() |
| 492 | ray_radius = ray_radius.cuda() |
| 493 | grid_size = grid_size.cuda() |
| 494 | grid_center = grid_center.cuda() |
| 495 | grid_width = grid_width.cuda() |
| 496 | |
| 497 | if n_rays * n_points < 5e9: |
| 498 | stime = timer() |
| 499 | all_ray2pidxs_gt = naive.find_neighbor_points_of_rays_brute_force( |
| 500 | points=points, |
| 501 | ray_origins=ray_origins, |
| 502 | ray_directions=ray_directions, |
| 503 | ray_radius=ray_radius, |
| 504 | grid_size=grid_size, |
| 505 | grid_center=grid_center, |
| 506 | grid_width=grid_width, |
| 507 | # include_outside=include_outside, |
| 508 | ) |
| 509 | total_time_brute_force = timer() - stime |
| 510 | else: |
| 511 | total_time_brute_force = np.inf |
| 512 | all_ray2pidxs_gt = None |
| 513 | |
| 514 | points = points.cpu() |
| 515 | ray_origins = ray_origins.cpu() |
| 516 | ray_directions = ray_directions.cpu() |
| 517 | ray_radius = ray_radius.cpu() |
| 518 | grid_size = grid_size.cpu() |
| 519 | grid_center = grid_center.cpu() |
| 520 | grid_width = grid_width.cpu() |
| 521 |
nothing calls this directly
no outgoing calls
no test coverage detected