File size: 1,332 Bytes
26225c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from torch_geometric.nn.pool.consecutive import consecutive_cluster
from src.utils.point import is_xyz_tensor
__all__ = ['xy_partition']
def xy_partition(pos, grid, consecutive=True):
"""Partition a point cloud based on a regular XY grid. Returns, for
each point, the index of the grid cell it falls into.
:param pos: Tensor
Point cloud
:param grid: float
Grid size
:param consecutive: bool
Whether the grid cell indices should be consecutive. That is to
say all indices in [0, idx_max] are used. Note that this may
prevent trivially mapping an index value back to the
corresponding XY coordinates
:return:
"""
assert is_xyz_tensor(pos)
# Compute the (i, j) coordinates on the XY grid size
i = pos[:, 0].div(grid, rounding_mode='trunc').long()
j = pos[:, 1].div(grid, rounding_mode='trunc').long()
# Shift coordinates to positive integer to avoid negatives
# clashing with our downstream indexing mechanism
i -= i.min()
j -= j.min()
# Compute a "manual" partition based on the grid coordinates
super_index = i * (max(i.max(), j.max()) + 1) + j
# If required, update the used indices to be consecutive
if consecutive:
super_index = consecutive_cluster(super_index)[0]
return super_index
|