grid.py 1.07 KB
Newer Older
rusty1s's avatar
update  
rusty1s committed
1
from typing import Optional
rusty1s's avatar
new api  
rusty1s committed
2

rusty1s's avatar
update  
rusty1s committed
3
import torch
rusty1s's avatar
rusty1s committed
4
5


rusty1s's avatar
update  
rusty1s committed
6
7
8
9
@torch.jit.script
def grid_cluster(pos: torch.Tensor, size: torch.Tensor,
                 start: Optional[torch.Tensor] = None,
                 end: Optional[torch.Tensor] = None) -> torch.Tensor:
rusty1s's avatar
rusty1s committed
10
11
    """A clustering algorithm, which overlays a regular grid of user-defined
    size over a point cloud and clusters all points within a voxel.
rusty1s's avatar
rusty1s committed
12
13
14
15

    Args:
        pos (Tensor): D-dimensional position of points.
        size (Tensor): Size of a voxel in each dimension.
rusty1s's avatar
new api  
rusty1s committed
16
        start (Tensor, optional): Start position of the grid (in each
rusty1s's avatar
rusty1s committed
17
            dimension). (default: :obj:`None`)
rusty1s's avatar
new api  
rusty1s committed
18
        end (Tensor, optional): End position of the grid (in each
rusty1s's avatar
rusty1s committed
19
20
            dimension). (default: :obj:`None`)

rusty1s's avatar
docs  
rusty1s committed
21
22
    :rtype: :class:`LongTensor`

rusty1s's avatar
update  
rusty1s committed
23
    .. code-block:: python
rusty1s's avatar
rusty1s committed
24

rusty1s's avatar
update  
rusty1s committed
25
26
        import torch
        from torch_cluster import grid_cluster
rusty1s's avatar
rusty1s committed
27

rusty1s's avatar
update  
rusty1s committed
28
29
30
31
32
        pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
        size = torch.Tensor([5, 5])
        cluster = grid_cluster(pos, size)
    """
    return torch.ops.torch_cluster.grid(pos, size, start, end)