Compute the intersection between grid cells and the ray. Args: ray_origins: (b, m, 3) ray_directions: (b, m, 3) ray_radius: (b,) or float grid_size: (b, 3) long, xyz grid_center: (b, 3), xyz
(
ray_origins: torch.Tensor, # (b, m, 3)
ray_directions: torch.Tensor, # (b, m, 3)
ray_radius: T.Union[torch.Tensor, float], # (b, )
grid_size: torch.Tensor, # (b, 3)
grid_center: torch.Tensor, # (b, 3) long
grid_width: torch.Tensor, # (b, 3) long
)
| 190 | |
| 191 | |
| 192 | def grid_ray_intersection( |
| 193 | ray_origins: torch.Tensor, # (b, m, 3) |
| 194 | ray_directions: torch.Tensor, # (b, m, 3) |
| 195 | ray_radius: T.Union[torch.Tensor, float], # (b, ) |
| 196 | grid_size: torch.Tensor, # (b, 3) |
| 197 | grid_center: torch.Tensor, # (b, 3) long |
| 198 | grid_width: torch.Tensor, # (b, 3) long |
| 199 | ): |
| 200 | """ |
| 201 | Compute the intersection between grid cells and the ray. |
| 202 | |
| 203 | Args: |
| 204 | ray_origins: |
| 205 | (b, m, 3) |
| 206 | ray_directions: |
| 207 | (b, m, 3) |
| 208 | ray_radius: |
| 209 | (b,) or float |
| 210 | grid_size: |
| 211 | (b, 3) long, xyz |
| 212 | grid_center: |
| 213 | (b, 3), xyz |
| 214 | grid_width: |
| 215 | (b, 3), xyz |
| 216 | |
| 217 | Returns: |
| 218 | list of list of list: b -> ray_idx -> grid_idx (including -1, outside the grid) |
| 219 | |
| 220 | Algorithm: |
| 221 | for each ray (px, py, pz, dx, dy, dz) |
| 222 | if dx.abs() < 1e-8: |
| 223 | # use y=c plane |
| 224 | else: |
| 225 | # use x=c plane |
| 226 | |
| 227 | - find plane-ray intersection point on each grid plane (xi, yi, zi) |
| 228 | - for each (xi, yi, zi) |
| 229 | x_from = discretize(xi - radius) |
| 230 | x_to = discretize(xi + radius) |
| 231 | (same for y and z) |
| 232 | |
| 233 | find grid idx for all of combinations, ie, get all grid cells |
| 234 | surrounding the point. |
| 235 | """ |
| 236 | |
| 237 | batch_size, n_rays, _ = ray_origins.shape |
| 238 | device = ray_origins.device |
| 239 | |
| 240 | if isinstance(ray_radius, float): |
| 241 | ray_radius = torch.ones(batch_size, device=device) * ray_radius # (b,) |
| 242 | |
| 243 | grid_from = (grid_center - grid_width / 2) # (b, 3) |
| 244 | cell_width = (grid_width / grid_size) # (b, 3) |
| 245 | total_cells = torch.prod(grid_size, dim=-1) # (b,) |
| 246 | |
| 247 | # determine whether to intersect with x=c plane, y=c plane, or z=c plane. |
| 248 | # It is important to select a direction so that the intersection point on two nearby planes |
| 249 | # do not exceed one grid in one of the rest of the two directions. |
no test coverage detected