def knn_points_idx(p1, p2, K, sorted=False, version=-1):
"""
K-Nearest neighbors on point clouds.
Args:
p1: Tensor of shape (N, P1, D) giving a batch of point clouds, each
containing P1 points of dimension D.
p2: Tensor of shape (N, P2, D) giving a batch of point clouds, each
containing P2 points of dimension D.
K: Integer giving the number of nearest neighbors to return
sorted: Whether to sort the resulting points.
version: Which KNN implementation to use in the backend. If version=-1,
the correct implementation is selected based on the shapes of
the inputs.
Returns:
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 idx[n, i, k] = j then p2[n, j] is one of the K
nearest neighbor to p1[n, i] in p2[n]. If sorted=True, then
p2[n, j] is the kth nearest neighbor to p1[n, i].
"""
idx, dists = _C.knn_points_idx(p1, p2, K, version)
if sorted:
dists, sort_idx = dists.sort(dim=2)
idx = idx.gather(2, sort_idx)
return idx, dists
def forward(
ctx, p1, p2, lengths1, lengths2, K, version, return_sorted: bool = True
):
"""
K-Nearest neighbors on point clouds.
Args:
p1: Tensor of shape (N, P1, D) giving a batch of N point clouds, each
containing up to P1 points of dimension D.
p2: Tensor of shape (N, P2, D) giving a batch of N point clouds, each
containing up to P2 points of dimension D.
lengths1: 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: 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.
K: Integer giving the number of nearest neighbors to return.
version: Which KNN implementation to use in the backend. If version=-1,
the correct implementation is selected based on the shapes of the inputs.
return_sorted: (bool) whether to return the nearest neighbors sorted in
ascending order of distance.
Returns:
p1_dists: 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.
p1_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.
"""
# pyre-fixme[16]: Module `pytorch3d` has no attribute `_C`.
idx, dists = _C.knn_points_idx(p1, p2, lengths1, lengths2, K, version)
# sort KNN in ascending order if K > 1
if K > 1 and return_sorted:
if lengths2.min() < K:
P1 = p1.shape[1]
mask = lengths2[:, None] <= torch.arange(K, device=dists.device)[None]
# mask has shape [N, K], true where dists irrelevant
mask = mask[:, None].expand(-1, P1, -1)
# mask has shape [N, P1, K], true where dists irrelevant
dists[mask] = float("inf")
dists, sort_idx = dists.sort(dim=2)
dists[mask] = 0
else:
dists, sort_idx = dists.sort(dim=2)
idx = idx.gather(2, sort_idx)
ctx.save_for_backward(p1, p2, lengths1, lengths2, idx)
ctx.mark_non_differentiable(idx)
return dists, idx
def knn_points_idx(
p1,
p2,
K: int,
lengths1=None,
lengths2=None,
sorted: bool = False,
version: int = -1,
):
"""
K-Nearest neighbors on point clouds.
Args:
p1: Tensor of shape (N, P1, D) giving a batch of point clouds, each
containing up to P1 points of dimension D.
p2: Tensor of shape (N, P2, D) giving a batch of point clouds, each
containing up to P2 points of dimension D.
K: Integer giving the number of nearest neighbors to return.
lengths1: 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: 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.
sorted: Whether to sort the resulting points.
version: Which KNN implementation to use in the backend. If version=-1,
the correct implementation is selected based on the shapes of the inputs.
Returns:
p1_neighbor_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 idx[n, i, k] = j then p2[n, j] is one of the K nearest
neighbors to p1[n, i] in p2[n]. If sorted=True, then p2[n, j] is the kth
nearest neighbor to p1[n, i]. 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.
If you want an (N, P1, K, D) tensor of the actual points, you can get it
using
p2[:, :, None].expand(-1, -1, K, -1).gather(1,
x_idx[:, :, :, None].expand(-1, -1, -1, D)
)
If K=1 and you want an (N, P1, D) tensor of the actual points, use
p2.gather(1, x_idx.expand(-1, -1, D))
p1_neighbor_dists: 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.
Warning: this is calculated outside of the autograd framework.
"""
P1 = p1.shape[1]
P2 = p2.shape[1]
if lengths1 is None:
lengths1 = torch.full((p1.shape[0],), P1, dtype=torch.int64, device=p1.device)
if lengths2 is None:
lengths2 = torch.full((p1.shape[0],), P2, dtype=torch.int64, device=p1.device)
idx, dists = _C.knn_points_idx(p1, p2, lengths1, lengths2, K, version)
if sorted:
if lengths2.min() < K:
device = dists.device
mask1 = lengths2[:, None] <= torch.arange(K, device=device)[None]
# mask1 has shape [N, K], true where dists irrelevant
mask2 = mask1[:, None].expand(-1, P1, -1)
# mask2 has shape [N, P1, K], true where dists irrelevant
dists[mask2] = float("inf")
dists, sort_idx = dists.sort(dim=2)
dists[mask2] = 0
else:
dists, sort_idx = dists.sort(dim=2)
idx = idx.gather(2, sort_idx)
return idx, dists
def knn():
_C.knn_points_idx(x, y, lengths, lengths, K, -1)
def knn():
_C.knn_points_idx(x, y, lengths1, lengths2, K, v)
torch.cuda.synchronize()
def knn():
_C.knn_points_idx(x, y, K, 0)
def knn():
_C.knn_points_idx(x, y, K, v)
torch.cuda.synchronize()