423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-08-03 04:12:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

KNN return order documentation

Browse Source 
Summary: Fix documentation of KNN, issue #180

Reviewed By: gkioxari

Differential Revision: D21384761

fbshipit-source-id: 2b36ee496f2060d17827d2fd66c490cdfa766866
This commit is contained in:
Jeremy Reizenstein 2020-05-04 13:35:49 -07:00 committed by Facebook GitHub Bot
parent 0eca74fa5f
commit 06ad1fb6c4
3 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
pytorch3d/ops/knn.py
View File

@ -115,18 +115,18 @@ def knn_points(
return_nn: If set to True returns the K nearest neighors in p2 for each point in p1. return_nn: If set to True returns the K nearest neighors in p2 for each point in p1.
Returns: Returns:
p1_idx: LongTensor of shape (N, P1, K) giving the indices of the 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.
idx: LongTensor of shape (N, P1, K) giving the indices of the
K nearest neighbors from points in p1 to points in p2. 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 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 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 in p2 has fewer than K points and where a cloud in p1 has fewer than P1
points. points.
p1_dists: Tensor of shape (N, P1, K) giving the squared distances to nn: Tensor of shape (N, P1, K, D) giving the K nearest neighbors in p2 for
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.
p2_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 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. for `p1[n, i]`. Returned if `return_nn` is True.
The nearest neighbors are collected using `knn_gather` The nearest neighbors are collected using `knn_gather`

2
pytorch3d/ops/points_alignment.py
View File

@ -158,7 +158,7 @@ def iterative_closest_point(
for iteration in range(max_iterations): for iteration in range(max_iterations):
Xt_nn_points = knn_points( Xt_nn_points = knn_points(
Xt, Yt, lengths1=num_points_X, lengths2=num_points_Y, K=1, return_nn=True Xt, Yt, lengths1=num_points_X, lengths2=num_points_Y, K=1, return_nn=True
)[2][:, :, 0, :] ).knn[:, :, 0, :]
# get the alignment of the nearest neighbors from Yt with Xt_init # get the alignment of the nearest neighbors from Yt with Xt_init
R, T, s = corresponding_points_alignment( R, T, s = corresponding_points_alignment(

4
pytorch3d/ops/utils.py
View File

@ -126,14 +126,14 @@ def get_point_covariances(
of shape `(minibatch, num_points, neighborhood_size, dim)`. of shape `(minibatch, num_points, neighborhood_size, dim)`.
""" """
# get K nearest neighbor idx for each point in the point cloud # get K nearest neighbor idx for each point in the point cloud
_, _, k_nearest_neighbors = knn_points( k_nearest_neighbors = knn_points(
points_padded, points_padded,
points_padded, points_padded,
lengths1=num_points_per_cloud, lengths1=num_points_per_cloud,
lengths2=num_points_per_cloud, lengths2=num_points_per_cloud,
K=neighborhood_size, K=neighborhood_size,
return_nn=True, return_nn=True,
) ).knn
# obtain the mean of the neighborhood # obtain the mean of the neighborhood
pt_mean = k_nearest_neighbors.mean(2, keepdim=True) pt_mean = k_nearest_neighbors.mean(2, keepdim=True)
# compute the diff of the neighborhood and the mean of the neighborhood # compute the diff of the neighborhood and the mean of the neighborhood