pointtorch.operations.torch

Point cloud processing operations for the use with PyTorch tensors.

pointtorch.operations.torch.knn_search( coords_support_points: Tensor, coords_query_points: Tensor, batch_indices_support_points: Tensor, batch_indices_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, k: int, return_sorted: bool = True, ) → Tuple[Tensor, Tensor][source]

Computes the indices of k nearest neighbors. Decides between different implementations:

Implementation from PyTorch3D: This implementation is always used if PyTorch3D is installed because its more efficient in terms of runtime and memory consumption than the other available implementations.
Implementation from torch-cluster: This implementation is used when PyTorch3D is not installed. It is similar to the PyTorch3D implementation but is slighlty slower.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
batch_indices_support_points – Indices indicating to which point cloud in the batch each support point belongs.
batch_indices_query_points – Indices indicating to which point cloud in the batch each query point belongs.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
k – The number of nearest neighbors to search.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to True. Setting it to False can improve performance for some implementations.

Returns:

Tuple of two tensors. The first tensor contains the indices of the neighbors of each query point. The second tensor contains the distances between the neighbors and the query points.

Raises:

ValueError – If the shapes of the input tensors are inconsistent.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
batch_indices_support_points: \((N)\)
batch_indices_query_points: \((N')\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: Tuple of two tensors, both with shape \((N', k)\) if k \(\leq n_{max}\), otherwise \((N', n_{max})\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.knn_search_cdist( coords_support_points: Tensor, coords_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, k: int, return_sorted: bool = True, ) → Tuple[Tensor, Tensor][source]

Computes the indices of k nearest neighbors. This implementation packs the point clouds into a regular batch structure and pads missing points so that all neighbor indices can be computed in parallel using PyTorch’s cdist function. The memory consumption of this implementation is quadratic with regard to the number of points and can therefore only be used for small point cloud sizes.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
k – The number of nearest neighbors to search.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to True. Setting it to False can improve performance for some implementations.

Returns:

Tuple of two tensors. The first tensor contains the indices of the neighbors of each query point. The second tensor contains the distances between the neighbors and the query points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: Tuple of two tensors, both with shape \((N', k)\) if k \(\leq n_{max}\), otherwise \((N', n_{max})\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

Raises:: ValueError – If the input point clouds contain too many points to compute a pairwise distance matrix.

pointtorch.operations.torch.knn_search_open3d( coords_support_points: Tensor, coords_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, k: int, ) → Tuple[Tensor, Tensor][source]

Computes the indices of k nearest neighbors. This implementation is based on Open3D’s knn_search function and currently only supports CPU devices.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
k – The number of nearest neighbors to search.

Returns:

Tuple of two tensors. The first tensor contains the indices of the neighbors of each query point. The second tensor contains the distances between the neighbors and the query points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: Tuple of two tensors, both with shape \((N', k)\) if k \(\leq n_{max}\), otherwise \((N', n_{max})\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.knn_search_pytorch3d( coords_support_points: Tensor, coords_query_points: Tensor, batch_indices_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, k: int, return_sorted: bool = True, ) → Tuple[Tensor, Tensor][source]

Computes the indices of k nearest neighbors. This implementation is based on PyTorch3D’s knn_points function.

The GPU-based KNN search implementation from PyTorch3D launches one CUDA thread per query point and each thread then loops through all the support points to find the k-nearest neighbors. It is similar to the torch-cluster implementation but it requires input batches of regular shape. Therefore, the variable size point cloud batches are packed into regular shaped batches before passing them to PyTorch3D.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
batch_indices_query_points – Indices indicating to which point cloud in the batch each query point belongs.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
k – The number of nearest neighbors to search.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to True. Setting it to False can improve performance for some implementations.

Returns:

Tuple of two tensors. The first tensor contains the indices of the neighbors of each query point. The second tensor contains the distances between the neighbors and the query points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
batch_indices_query_points: \((N')\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: Tuple of two tensors, both with shape \((N', k)\) if k \(\leq n_{max}\), otherwise \((N', n_{max})\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.knn_search_torch_cluster( coords_support_points: Tensor, coords_query_points: Tensor, batch_indices_support_points: Tensor, batch_indices_query_points: Tensor, point_cloud_sizes_support_points: Tensor, k: int, ) → Tuple[Tensor, Tensor][source]

Computes the indices of k nearest neighbors. This implementation is based on the knn method from torch-cluster.

The GPU-based KNN search implementation from torch-cluster launches one CUDA thread per query point and each thread then loops through all the support points to find the k-nearest neighbors. It is similar to the PyTorch3D implementation but can handle variable size point clouds directly.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
batch_indices_support_points – Indices indicating to which point cloud in the batch each support point belongs.
batch_indices_query_points – Indices indicating to which point cloud in the batch each query point belongs.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
k – The number of nearest neighbors to search.

Returns:

Tuple of two tensors. The first tensor contains the indices of the neighbors of each query point. The second tensor contains the distances between the neighbors and the query points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
batch_indices_support_points: \((N)\)
batch_indices_query_points: \((N')\)
point_cloud_sizes_support_points: \((B)\)
Output: Tuple of two tensors, both with shape \((N', k)\) if k \(\leq n_{max}\), otherwise \((N', n_{max})\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.make_labels_consecutive( labels: Tensor, start_id: int = 0, ignore_id: int | None = None, inplace: bool = False, return_unique_labels: bool = False, ) → Tensor | Tuple[Tensor, Tensor][source]

Transforms the input labels into consecutive integer labels starting from a given start_id.

Parameters:

labels – An array of original labels.
start_id – The starting ID for the consecutive labels. Defaults to zero.
ignore_id – A label ID that should not be changed when transforming the labels.
inplace – Whether the transformation should be applied inplace to the labels array. Defaults to False.
return_unique_labels – Whether the unique labels after applying the transformation (excluding ignore_id) should be returned. Defaults to False.

Returns:

A tensor with the transformed consecutive labels. If return_unique_labels is set to True, a tuple of two tensors is returned, where the second tensor contains the unique labels after the transformation.

pointtorch.operations.torch.max_pooling( x: Tensor, point_cloud_indices: Tensor, expand: bool = False, ) → Tensor[source]

Max pooling over all point features within a batch item.

Parameters:

x – Point features to be pooled.
point_cloud_indices – Indices indicating to which point cloud in the batch each point belongs.
expand – Whether the max-pooling result should be expanded to the size of the input point clouds by repeating the maximum value for each point. Defaults to False.

Returns:

Max-pooled features.

Shape:

x: \((N, D)\)
point_cloud_indices: \((N)\)
Output: \((B, D)\) if expand is False, otherwise \((N, D)\).

where

\(N = \text{ number of points}\)

\(B = \text{ batch size}\)

\(D = \text{ number of feature channels}\)

pointtorch.operations.torch.neighbor_search_cdist( coords_support_points: Tensor, coords_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, *, radius: float | None = None, k: int | None = None, return_sorted: bool = False, ) → Tuple[Tensor, Tensor][source]

Retrieves the neighbor points within a given search radius or the k nearest neighbors. The implementation of this function is based on PyTorch’s cdist function. The memory consumption of this implementation is quadratic in the number of points, so it can only be used for small point cloud sizes. Internally, this function packs the point clouds into a regular batch structure by padding all point clouds to the same size, so that all neighbor points can be computed in parallel using PyTorch’s cdist function.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
radius – Search radius in which to search for neighbors. Defaults to None, which means that an infinite search radius is used. Either radius or k must not be None.
k – The maximum number of neighbors to search. If radius is not None and the radius neighborhood of a point contains more than k points, the k nearest neighbors are selected. Defaults to None, which means that all neighbors within the specified radius are searched. Either radius or k must not be None.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to False.

Returns:

Tuple of two tensors. The first tensor contains the indices of the neighbors of each query point. The second tensor contains the distances between the neighbors and the query points. If a query point has less than \(n_{max}\) neighbors, where \(n_{max}\) is the maximum number of neighbors a query point has, the invalid neighbor indices in the first tensor are set to \(N + 1\) where \(N\) is the number of support points and the invalid distances in the second tensor are set to torch.inf.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: Tuple of two tensors, both with shape \((N', n_{max})\) if k is None or \(n_{max} <\) k, otherwise \((N', k)\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

Raises:: ValueError – If both radius and k are set to None or the input point clouds contain too many points to compute a pairwise distance matrix.

pointtorch.operations.torch.pack_batch( input_batch: Tensor, point_cloud_sizes: Tensor, fill_value: float = inf, ) → Tuple[Tensor, Tensor][source]

Packs a batch containing point clouds of varying size into a regular batch structure by padding all point clouds to the same size.

Parameters:

input_batch – Batch to be packed.
point_cloud_sizes – Number of points in each point cloud in the batch.
fill_value – Value to be used to pad point clouds that contain less points than the largest point cloud in the batch. Defaults to torch.inf.

Returns:

Tuple of two tensors. The first tensor is the packed batch. Point clouds containing less than \(N_{max}\) points are padded with fill_value. The second tensor is a boolean mask, which is True in all positions where the packed batch contains valid points and False in all positions filled with fill_value.

Shape:

input_batch: \((N_1 + ... + N_B, D)\)
point_cloud_sizes: \((B)\)
Output: Tuple of two tensors with shape \((B, N_{max}, D)\) and \((B, N_{max})\).

where

\(B = \text{ batch size}\)

\(D = \text{ number of feature channels}\)

\(N_i = \text{ number of points in the i-th point cloud}\)

\(N_{max} = \text{ number of points in the largest point cloud in the batch}\)

pointtorch.operations.torch.radius_search( coords_support_points: Tensor, coords_query_points: Tensor, batch_indices_support_points: Tensor, batch_indices_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, radius: float, voxel_size: float | None = None, k: int | None = None, return_sorted: bool = False, ) → Tensor[source]

Retrieves the indices of all neighbor points within a radius. Decides between different implementations:

Implementation from Open3D: The Open3D implementation is based on a spatial hash table and therefore achieves a high memory efficiency if all neighbors within the given radius are to be searched, i.e., if k is None. None of the available implementations provides options to sort the returned neighbors by distance or to select the k-nearest neighbors from the search radius if it contains more than k points. Therefore, if return_sorted is True, all neighbors have to be searched and sorted by distance afterwards. For this reason, the Open3D implementation is used when k is None or return_sorted is True and Open3D is installed. The Open3D implementation is only used when the input size is smaller than 1048576. This is because the memory allocation of the Open3D implementation potentially contains a bug and does not work for very large inputs.
Implementation from PyTorch3D: PyTorch3D implements a hybrid search that limits the maximum number of neighbors to k. The GPU-based implementation launches one CUDA thread per query point and then loops through all support points until at most k neighbors are found within the search radius. The implementation requires the point clouds to be packed into a regular batch structure. The implementation is quite efficient for small values of k. For a batch size of 1, it usually outperforms the torch-cluster implementation. For larger batch sizes, the torch-cluster implementation uses slightly less memory because it can handle variable-sized point clouds directly and does not require packing to a regular batch. Therefore, the PyTorch3D implementation is used when k is not None and the batch size is 1 and PyTorch3D is installed. The torch-cluster implementation is used when k is not None and the batch size is greater than 1 or when PyTorch3D is not installed.
Implementation from torch-cluster: The torch-cluster implementation is quite similar to the PyTorch3D implementation but it can handle variable size point clouds directly.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
batch_indices_support_points – Indices indicating to which point cloud in the batch each support point belongs.
batch_indices_query_points – Indices indicating to which point cloud in the batch each query point belongs.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
radius – Search radius in which to search for neighbors.
voxel_size – Voxel size that was used to downsample the support point clouds before passing them to this method. If specified, this information can be used to calculate the maximum possible number of points within the search radius, which may be used to reduce the memory consumption of the neighbor search. Defaults to None.
k – The maximum number of neighbors to search. If the radius neighborhood of a point contains more than k points, the returned neighbors are picked randomly if return_sorted is False. Otherwise, the k nearest neighbors are selected. Defaults to None, which means that all neighbors within the specified radius are returned.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to False.

Returns:

The indices of the neighbors of each query point. If a query point has less than \(n_{max}\) neighbors, where \(n_{max}\) is the maximum number of neighbors a query point has, the invalid neighbor indices are set to \(N + 1\) where \(N\) is the number of support points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
batch_indices_support_points: \((N)\)
batch_indices_query_points: \((N')\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: \((N', n_{max})\) if k is None or \(n_{max} <\) k, otherwise \((N', k)\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.radius_search_cdist( coords_support_points: Tensor, coords_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, radius: float, voxel_size: float | None = None, k: int | None = None, return_sorted: bool = False, ) → Tensor[source]

Computes the indices of all neighbor points within a radius. This implementation packs the point clouds into a regular batch structure and pads missing points so that all neighbor indices can be computed in parallel using PyTorch’s cdist function. The memory consumption of this implementation is quadratic with regard to the number of points and can therefore only be used for small point cloud sizes.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
radius – Search radius in which to search for neighbors.
voxel_size – Voxel size that was used to downsample the support point clouds before passing them to this method. If specified, this information can be used to calculate the maximum possible number of points within the search radius, which may be used to reduce the memory consumption of the neighbor search. Defaults to None.
k – The maximum number of neighbors to search. If the radius neighborhood of a point contains more than k points, the returned neighbors are picked randomly if return_sorted is False. Otherwise, the k nearest neighbors are selected. Defaults to None, which means that all neighbors within the specified radius are returned.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to False.

Returns:

The indices of the neighbors of each query point. If a query point has less than \(n_{max}\) neighbors, where \(n_{max}\) is the maximum number of neighbors a query point has, the invalid neighbor indices are set to \(N + 1\) where \(N\) is the number of support points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: \((N', n_{max})\) if k is None or \(n_{max} <\) k, otherwise \((N', k)\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

Raises:: ValueError – If the input point clouds contain too many points to compute a pairwise distance matrix.

pointtorch.operations.torch.radius_search_open3d( coords_support_points: Tensor, coords_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, radius: float, k: int | None = None, return_sorted: bool = False, ) → Tensor[source]

Computes the indices of all neighbor points within a radius. This implementation is based on Open3D and uses a spatial hash table to implement the neighborhood search. Computes the indices of all neighbor points within a radius. This implementation is based on Open3D’s fixed_radius_search function and currently only supports CPU devices.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
radius – Search radius in which to search for neighbors.
k – The maximum number of neighbors to search. If the radius neighborhood of a point contains more than k points, the returned neighbors are picked randomly if return_sorted is False. Otherwise, the k nearest neighbors are selected. Defaults to None, which means that all neighbors within the specified radius are returned.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to False.

Returns:

The indices of the neighbors of each query point. If a query point has less than \(n_{max}\) neighbors, where \(n_{max}\) is the maximum number of neighbors a query point has, the invalid neighbor indices are set to \(N + 1\) where \(N\) is the number of support points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: \((N', n_{max})\) if k is None or \(n_{max} <\) k, otherwise \((N', k)\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.radius_search_pytorch3d( coords_support_points: Tensor, coords_query_points: Tensor, point_cloud_sizes_support_points: Tensor, point_cloud_sizes_query_points: Tensor, radius: float, voxel_size: float | None = None, k: int | None = None, return_sorted: bool = False, ) → Tensor[source]

Computes the indices of all neighbors within a radius. This implementation is based on PyTorch3D’s ball_query function.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
point_cloud_sizes_query_points – Number of points in each point cloud in the batch of query points.
radius – Search radius in which to search for neighbors.
voxel_size – Voxel size that was used to downsample the support point clouds before passing them to this method. If specified, this information can be used to calculate the maximum possible number of points within the search radius, which may be used to reduce the memory consumption of the neighbor search. Defaults to None.
k – The maximum number of neighbors to search. If the radius neighborhood of a point contains more than k points, the returned neighbors are picked randomly if return_sorted is False. Otherwise, the k nearest neighbors are selected. Defaults to None, which means that all neighbors within the specified radius are returned.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to False.

Returns:

The indices of the neighbors of each query point. If a query point has less than \(n_{max}\) neighbors, where \(n_{max}\) is the maximum number of neighbors a query point has, the invalid neighbor indices are set to \(N + 1\) where \(N\) is the number of support points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
batch_indices_support_points: \((N)\)
batch_indices_query_points: \((N')\)
point_cloud_sizes_support_points: \((B)\)
point_cloud_sizes_query_points: \((B)\)
Output: \((N', n_{max})\) if k is None or \(n_{max} <\) k, otherwise \((N', k)\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.radius_search_torch_cluster( coords_support_points: Tensor, coords_query_points: Tensor, batch_indices_support_points: Tensor, batch_indices_query_points: Tensor, point_cloud_sizes_support_points: Tensor, radius: float, voxel_size: float | None = None, k: int | None = None, return_sorted: bool = False, ) → Tensor[source]

Computes the indices of all neighbor points within a radius. This implementation is based on the radius method from torch-cluster.

Parameters:

coords_support_points – Coordinates of the support points to be searched for neighbors.
coords_query_points – Coordinates of the query points.
batch_indices_support_points – Indices indicating to which point cloud in the batch each support point belongs.
batch_indices_query_points – Indices indicating to which point cloud in the batch each query point belongs.
point_cloud_sizes_support_points – Number of points in each point cloud in the batch of support points.
radius – Search radius in which to search for neighbors.
voxel_size – Voxel size that was used to downsample the support point clouds before passing them to this method. If specified, this information can be used to calculate the maximum possible number of points within the search radius, which may be used to reduce the memory consumption of the neighbor search. Defaults to None.
k – The maximum number of neighbors to search. If the radius neighborhood of a point contains more than k points, the returned neighbors are picked randomly if return_sorted is False. Otherwise, the k nearest neighbors are selected. Defaults to None, which means that all neighbors within the specified radius are returned.
return_sorted – Whether the returned neighbors should be sorted by their distance to the query point. Defaults to False.

Returns:

The indices of the neighbors of each query point. If a query point has less than \(n_{max}\) neighbors, where \(n_{max}\) is the maximum number of neighbors a query point has, the invalid neighbor indices are set to \(N + 1\) where \(N\) is the number of support points.

Shape:

coords_support_points: \((N, 3)\)
coords_query_points: \((N', 3)\)
batch_indices_support_points: \((N)\)
batch_indices_query_points: \((N')\)
point_cloud_sizes_support_points: \((B)\)
Output: \((N', n_{max})\) if k is None or \(n_{max} <\) k, otherwise \((N', k)\).

where

\(B = \text{ batch size}\)

\(N = \text{ number of support points}\)

\(N' = \text{ number of query points}\)

\(n_{max} = \text{ maximum number of neighbors a query point has}\)

pointtorch.operations.torch.random_sampling( points: Tensor, batch_indices: Tensor, point_cloud_sizes: Tensor, new_point_cloud_sizes: Tensor, ) → Tensor[source]

Randomly samples a given number of points from each point cloud in a batch. From each point cloud in a batch, as many points as specified by new_point_cloud_sizes are selected. The points are expected to be shuffled beforehand, so the first new_point_cloud_sizes[i] points are selected from the i-th point cloud and no shuffling is performed by this method.

Parameters:

points – Batch of point clouds from which to sample.
batch_indices – Indices indicating to which input point cloud each point in the batch belongs.
point_cloud_sizes – Number of points contained in each input point cloud.
new_point_cloud_sizes – Number of points to be sampled from each input point cloud.

Returns:

Sampling results.

Shape:

points: \((N, ...)\)
batch_indices: \((N)\)
point_cloud_sizes: \((B)\)
new_point_cloud_sizes: \((B)\)
Output: \((N', ...)\)

where

\(N = \text{ number of points before sampling}\)

\(N' = \text{ number of points after sampling}\)

\(B = \text{ batch size}\)

pointtorch.operations.torch.ravel_index( index: Tensor, input_tensor: Tensor, dim: int, ) → Tensor[source]

Converts index argument of a multi-dimensional torch.gather() or torch.scatter_add() operation to an index that can be used to apply the operation on the flattened input tensor.

Parameters:

index – Multi-dimensional index.
input_tensor – Input tensor of the torch.gather() or torch.scatter_add() operation.
dim – Dimension in which the torch.gather() or torch.scatter_add() operation is to be applied.

Returns:

Index for applying torch.gather() or torch.scatter_add() on the flattened input tensor.

pointtorch.operations.torch.ravel_multi_index( multi_index: Tensor, dims: Size | Tensor, ) → Tensor[source]

PyTorch implementation of numpy.ravel_multi_index. This operation is inverse to pointtorch.operations.torch.unravel_flat_index(). Note that, compared to the numpy implementation, the shape of multi_index is transposed.

Parameters:

multi_index – Tensor containing the indices for each dimension.
dims – The shape of the tensor into which the indices from multi_index apply.

Returns:

Indices for the flattened version of the tensor, referring to the same elements as referenced by multi_index for the non-flattened version of the tensor.

Shape:

multi_index: \((N, D)\)
dims: \((D)\)
Output: \((N)\)

where

\(N = \text{ number of items}\)

\(D = \text{ number of index dimensions}\)

pointtorch.operations.torch.shuffle( points: Tensor, point_cloud_sizes: Tensor, generator: Generator | None = None, ) → Tuple[Tensor, Tensor][source]

Shuffles points within a batch of point clouds. Each point cloud in the batch can contain a different number of points.

Parameters:

points – Batch of point clouds to shuffle.
point_cloud_sizes – Number of points contained in each input point cloud.
generator – Random generator to be used for shuffling. Defaults to None.

Returns:

Tuple of two tensors. The first is the shuffled tensor. The second contains the index of each point after shuffling.

Shape:

points: \((N, ...)\)
point_cloud_sizes: \((B)\)
Output: Tuple of two tensors. The first has the same shape as points. The second has shape (N).

where

\(N = \text{ number of points}\)

\(B = \text{ batch size}\)

pointtorch.operations.torch.unravel_flat_index( flat_index: Tensor, dims: Size | Tensor, ) → Tensor[source]

Converts an index for a 1-dimensional tensor into an index for an equivalent multi-dimensional tensor. This operation is inverse to pointtorch.operations.torch.ravel_multi_index().

Parameters:

flat_index – Tensor containing the indices for the flat array.
dims – The shape of the tensor into which the returned indices should apply.

Returns:

Indices for the multi-dimensional version of the tensor, referring to the same elements as referenced by flat_index for the flattened version of the tensor.

Shape:

flat_index: \((N)\)
dims: \((D)\)
Output: \((N, D)\)

where

\(N = \text{ number of items}\)

\(D = \text{ number of index dimensions}\)

pointtorch.operations.torch.voxel_downsampling( coords: Tensor, batch_indices: Tensor, point_cloud_sizes: Tensor, voxel_size: float, *, features: Tensor | None = None, feature_aggregation: Literal['max', 'mean', 'min', 'nearest_neighbor'] = 'mean', point_aggregation: Literal['mean', 'nearest_neighbor'] = 'mean', preserve_order: bool = True, start: Tensor | None = None, ) → Tuple[Tensor, Tensor | None, Tensor, Tensor, Tensor][source]

Voxel-based downsampling of a batch of point clouds.

Parameters:

coords – Coordinates of the points to be downsampled.
batch_indices – Indices indicating to which input point cloud each point in the batch belongs.
point_cloud_sizes – Number of points contained in each input point cloud.
voxel_size – The size of the voxels used for downsampling.
features – The features of the points to be downsampled. Defaults to None.
feature_aggregation – Method to be used to aggregate features of points within the same voxel: "max" | "mean" | "min" | "nearest_neighbor". "nearest_neighbor" means that the features of the point closest to the voxel center are selected.
point_aggregation – Method to be used to aggregate the point coordinates within the same voxel: "mean" | "nearest_neighbor". "nearest_neighbor" means that the coordinates of the point closest to the voxel center are selected.
preserve_order – If this is set to True and point_aggregation is set to "nearest_neighbor", the point order is preserved during downsampling. This means that for any two points included in the downsampled point cloud, the point that is first in the original point cloud is also first in the downsampled point cloud. Defaults to True.
start – Coordinates of a point at which the voxel grid is to be aligned, i.e., the grid is placed so that it starts at a corner point of a voxel. Defaults to None, which means that the grid is aligned at the coordinate origin.

Returns:

Tuple of five tensors. The first contains the coordinates of the downsampled points. The second tensor contains the features of the downsampled points and is None if no input features are provided. The third contains indices indicating to which point cloud each downsampled point belongs. The fourth contains the size of each downsampled point cloud. The fifth contains indices indicating in which voxel each point from the original point clouds is located.

Raises:

ValueError – If start is not None and has an invalid shape.

Shape:

coords: \((N, 3)\)
batch_indices: \((N)\)
point_cloud_sizes: \((B)\)
features: \((N, D)\)
start: \((B, 3)\)
Output: Tuple of five tensors. The first has shape \((N', 3)\). The second has shape \((N', D)\). The third has shape \((N')\). The fourth has shape \((B)\). The fifth has shape \((N)\).

where

\(N = \text{ number of points before downsampling}\)

\(N' = \text{ number of points after downsampling}\)

\(B = \text{ batch size}\)

\(D = \text{ number of feature channels}\)