fastdev.geom.knn

Module Contents

fastdev.geom.knn.knn_points(p1: torch.Tensor, p2: torch.Tensor, lengths1: torch.Tensor | None = None, lengths2: torch.Tensor | None = None, norm: int = 2, K: int = 1, version: int = -1, return_nn: bool = False, return_sorted: bool = True) → _KNN[source]

K-Nearest neighbors on point clouds.

Parameters:

  • p1 (torch.Tensor) – Tensor of shape (N, P1, D) giving a batch of N point clouds, each containing up to P1 points of dimension D.

  • p2 (torch.Tensor) – Tensor of shape (N, P2, D) giving a batch of N point clouds, each containing up to P2 points of dimension D.

  • lengths1 (Union[torch.Tensor, None]) – LongTensor of shape (N,) of values in the range [0, P1], giving the length of each pointcloud in p1. Or None to indicate that every cloud has length P1.

  • lengths2 (Union[torch.Tensor, None]) – LongTensor of shape (N,) of values in the range [0, P2], giving the length of each pointcloud in p2. Or None to indicate that every cloud has length P2.

  • norm (int) – Integer indicating the norm of the distance. Supports only 1 for L1, 2 for L2.

  • K (int) – Integer giving the number of nearest neighbors to return.

  • version (int) – Which KNN implementation to use in the backend. If version=-1, the correct implementation is selected based on the shapes of the inputs.

  • return_nn (bool) – If set to True returns the K nearest neighbors in p2 for each point in p1.

  • return_sorted (bool) – (bool) whether to return the nearest neighbors sorted in ascending order of distance.

Returns:

Tensor of shape (N, P1, K) giving the squared distances to

the nearest neighbors. This is padded with zeros both where a cloud in p2 has fewer than K points and where a cloud in p1 has fewer than P1 points.

idx: LongTensor of shape (N, P1, K) giving the indices of the

K nearest neighbors from points in p1 to points in p2. Concretely, if p1_idx[n, i, k] = j then p2[n, j] is the k-th nearest neighbors to p1[n, i] in p2[n]. This is padded with zeros both where a cloud in p2 has fewer than K points and where a cloud in p1 has fewer than P1 points.

nn: Tensor of shape (N, P1, K, D) giving the K nearest neighbors in p2 for

each point in p1. Concretely, p2_nn[n, i, k] gives the k-th nearest neighbor for p1[n, i]. Returned if return_nn is True. The nearest neighbors are collected using knn_gather

p2_nn = knn_gather(p2, p1_idx, lengths2)

which is a helper function that allows indexing any tensor of shape (N, P2, U) with the indices p1_idx returned by knn_points. The output is a tensor of shape (N, P1, K, U).

Return type:

dists

fastdev.geom.knn.knn_gather(x: torch.Tensor, idx: torch.Tensor, lengths: torch.Tensor | None = None)[source]

A helper function for knn that allows indexing a tensor x with the indices idx returned by knn_points.

For example, if dists, idx = knn_points(p, x, lengths_p, lengths, K) where p is a tensor of shape (N, L, D) and x a tensor of shape (N, M, D), then one can compute the K nearest neighbors of p with p_nn = knn_gather(x, idx, lengths). It can also be applied for any tensor x of shape (N, M, U) where U != D.

Parameters:

  • x (torch.Tensor) – Tensor of shape (N, M, U) containing U-dimensional features to be gathered.

  • idx (torch.Tensor) – LongTensor of shape (N, L, K) giving the indices returned by knn_points.

  • lengths (Union[torch.Tensor, None]) – LongTensor of shape (N,) of values in the range [0, M], giving the length of each example in the batch in x. Or None to indicate that every example has length M.

Returns:

Tensor of shape (N, L, K, U) resulting from gathering the elements of x

with idx, s.t. x_out[n, l, k] = x[n, idx[n, l, k]]. If k > lengths[n] then x_out[n, l, k] is filled with 0.0.

Return type:

x_out