Find all the points within ray_radius of a ray. Args: points: (b, n, 3) the xyz_w of points in world coord ray_origins: (b, m, 3) the ray origin in the world coord ray_directions: (b, m, 3) the ray direction in the world coord
(
points: torch.Tensor, # (b, n, 3)
ray_origins: torch.Tensor, # (b, m, 3)
ray_directions: torch.Tensor, # (b, m, 3)
ray_radius: T.Union[torch.Tensor, float], # (b,)
grid_size: T.Union[torch.Tensor, int], # (b, 3)
grid_center: T.Union[torch.Tensor, float, None] = None, # (b, 3)
grid_width: T.Union[torch.Tensor, float, None] = None, # (b, 3)
)
| 57 | |
| 58 | |
| 59 | def find_neighbor_points_of_rays( |
| 60 | points: torch.Tensor, # (b, n, 3) |
| 61 | ray_origins: torch.Tensor, # (b, m, 3) |
| 62 | ray_directions: torch.Tensor, # (b, m, 3) |
| 63 | ray_radius: T.Union[torch.Tensor, float], # (b,) |
| 64 | grid_size: T.Union[torch.Tensor, int], # (b, 3) |
| 65 | grid_center: T.Union[torch.Tensor, float, None] = None, # (b, 3) |
| 66 | grid_width: T.Union[torch.Tensor, float, None] = None, # (b, 3) |
| 67 | ) -> T.List[T.List[T.List[int]]]: |
| 68 | """ |
| 69 | Find all the points within ray_radius of a ray. |
| 70 | |
| 71 | Args: |
| 72 | points: |
| 73 | (b, n, 3) the xyz_w of points in world coord |
| 74 | ray_origins: |
| 75 | (b, m, 3) the ray origin in the world coord |
| 76 | ray_directions: |
| 77 | (b, m, 3) the ray direction in the world coord |
| 78 | ray_radius: |
| 79 | (b,) the radius of each ray |
| 80 | grid_size: |
| 81 | (b, 3) the number of grid points in xyz direction |
| 82 | grid_center: |
| 83 | (b, 3) the center of the grid. |
| 84 | If None: average of points |
| 85 | grid_width: |
| 86 | (b, 3) the full width of the grid. |
| 87 | If None: (max-min) from points |
| 88 | |
| 89 | Returns: |
| 90 | list of list of list: b -> m -> n_idx |
| 91 | n_idx is the index of the points |
| 92 | """ |
| 93 | |
| 94 | batch_size = points.size(0) |
| 95 | n_rays = ray_origins.size(-2) |
| 96 | |
| 97 | if isinstance(ray_radius, (float, int)): |
| 98 | # ray radius should have only one dimension (batch) |
| 99 | # ray_radius = torch.ones(batch_size, 3, dtype=points.dtype, device=points.device) * ray_radius |
| 100 | ray_radius = torch.ones(batch_size, dtype=points.dtype, device=points.device) * ray_radius |
| 101 | |
| 102 | if isinstance(grid_size, int): |
| 103 | grid_size = torch.ones(batch_size, 3, dtype=torch.long, device=points.device) * grid_size |
| 104 | |
| 105 | if grid_center is None: |
| 106 | grid_center = torch.mean(points, dim=-2) # (b, 3) |
| 107 | elif isinstance(grid_center, (float, int)): |
| 108 | grid_center = torch.ones(batch_size, 3, dtype=points.dtype, device=points.device) * grid_center |
| 109 | |
| 110 | if grid_width is None: |
| 111 | grid_width = torch.max(points, dim=-2)[0] - torch.min(points, dim=-2)[0] # (b, 3) |
| 112 | elif isinstance(grid_width, (float, int)): |
| 113 | grid_width = torch.ones(batch_size, 3, dtype=points.dtype, device=points.device) * grid_width |
| 114 | |
| 115 | t_min = 0. |
| 116 | t_max = 1.e10 |