Compute the grid index given xyz_w. Args: points: (*, n, 3) grid_size: (*, 3) long. number of grid cells in x y z. center: (*, 3) center of the grid grid_width: (*, 3) length (full width) of the grid in xyz
(
points: torch.Tensor, # (*, n, 3)
grid_size: T.Union[torch.Tensor, int], # (*, 3)
center: T.Union[torch.Tensor, float] = 0., # (*, 3)
grid_width: T.Union[torch.Tensor, float] = 1., # (*, 3)
mode: str = 'ind',
)
| 73 | |
| 74 | |
| 75 | def get_grid_idx( |
| 76 | points: torch.Tensor, # (*, n, 3) |
| 77 | grid_size: T.Union[torch.Tensor, int], # (*, 3) |
| 78 | center: T.Union[torch.Tensor, float] = 0., # (*, 3) |
| 79 | grid_width: T.Union[torch.Tensor, float] = 1., # (*, 3) |
| 80 | mode: str = 'ind', |
| 81 | ) -> T.Tuple[torch.Tensor, torch.Tensor]: |
| 82 | """ |
| 83 | Compute the grid index given xyz_w. |
| 84 | Args: |
| 85 | points: |
| 86 | (*, n, 3) |
| 87 | grid_size: |
| 88 | (*, 3) long. number of grid cells in x y z. |
| 89 | center: |
| 90 | (*, 3) center of the grid |
| 91 | grid_width: |
| 92 | (*, 3) length (full width) of the grid in xyz |
| 93 | include_all: |
| 94 | whether to include any points outside the grid |
| 95 | If True, all out-of-bound points will be assigned to -1 or (-1, -1, -1) |
| 96 | mode: |
| 97 | 'subidx': return sub idx |
| 98 | 'ind': return linear index |
| 99 | |
| 100 | Returns: |
| 101 | grid_idx: |
| 102 | if mode == 'subidx': (*, n, 3) long |
| 103 | elif mode == 'ind': (*, n) long |
| 104 | valid_mask: |
| 105 | (*, n) bool |
| 106 | Algorithm: |
| 107 | Let |
| 108 | x_from = center_x - grid_length_x / 2 |
| 109 | x_to = center_x + grid_length_x / 2 |
| 110 | cell_width_x = grid_length / grid_size_x |
| 111 | |
| 112 | We divide x = [x_from, x_to] into grid_size cells, each cell is of width x_cell. |
| 113 | |
| 114 | x_idx = ((x - x_from) / cell_width_x).floor().clamp(0, grid_size_x-1) |
| 115 | y_idx = ((y - y_from) / cell_width_y).floor().clamp(0, grid_size_y-1) |
| 116 | z_idx = ((z - z_from) / cell_width_z).floor().clamp(0, grid_size_z-1) |
| 117 | |
| 118 | grid_idx = x_idx + y_dix * grid_size_x + z_idx * (grid_size_x * grid_size_y) |
| 119 | """ |
| 120 | |
| 121 | if isinstance(center, float): |
| 122 | center = torch.tensor(center, dtype=points.dtype, device=points.device) |
| 123 | center = center.view(*([1] * (points.ndim - 1))).expand(*([1] * (points.ndim - 2) + [3])) |
| 124 | if isinstance(grid_size, (float, int)): |
| 125 | grid_size = torch.tensor(grid_size, dtype=torch.long, device=points.device) |
| 126 | grid_size = grid_size.view(*([1] * (points.ndim - 1))).expand(*([1] * (points.ndim - 2) + [3])) |
| 127 | if isinstance(grid_width, (float, int)): |
| 128 | grid_width = torch.tensor(grid_width, dtype=points.dtype, device=points.device) |
| 129 | grid_width = grid_width.view(*([1] * (points.ndim - 1))).expand(*([1] * (points.ndim - 2) + [3])) |
| 130 | |
| 131 | grid_size = grid_size.to(dtype=torch.long, device=points.device) |
| 132 | center = center.to(device=points.device) |
no test coverage detected