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()