Tensorflow2计算二维点集的K近邻

# -*- coding: utf-8 -*-
"""
Created on 2021.06.17
@author: xi'an Li
"""
import tensorflow as tf
import numpy as np


def pairwise_distance(point_set):
    """Compute pairwise distance of a point cloud.
        Args:
          (x-y)^2 = x^2 - 2xy + y^2
          point_set: tensor (num_points, dims2point)
        Returns:
          pairwise distance: (num_points, num_points)
    """
    point_set_shape = point_set.get_shape().as_list()
    assert(len(point_set_shape)) == 2

    point_set_transpose = tf.transpose(point_set, perm=[1, 0])
    point_set_inner = tf.matmul(point_set, point_set_transpose)
    point_set_inner = -2 * point_set_inner
    point_set_square = tf.reduce_sum(tf.square(point_set), axis=-1, keepdims=True)
    point_set_square_transpose = tf.transpose(point_set_square, perm=[1, 0])
    return point_set_square + point_set_inner + point_set_square_transpose


def knn_includeself(dist_matrix, k=20):
    """Get KNN based on the pairwise distance.
        How to use tf.nn.top_k(): https://blog.youkuaiyun.com/wuguangbin1230/article/details/72820627
      Args:
        pairwise distance: (num_points, num_points)
        k: int

      Returns:
        nearest neighbors: (num_points, k)
      """
    neg_dist = -1.0*dist_matrix
    _, nn_idx = tf.nn.top_k(neg_dist, k=k)  # 这个函数的作用是返回 input 中每行最大的 k 个数，并且返回它们所在位置的索引
    return nn_idx


def knn_excludeself(dist_matrix, k=20):
    """Get KNN based on the pairwise distance.
      Args:
        pairwise distance: (num_points, num_points)
        k: int

      Returns:
        nearest neighbors index: (num_points, k)
      """
    neg_dist = -1.0*dist_matrix
    k_neighbors = k+1
    _, knn_idx = tf.nn.top_k(neg_dist, k=k_neighbors)  # 这个函数的作用是返回 input 中每行最大的 k 个数，并且返回它们所在位置的索引
    nn_idx = knn_idx[:, 1: k_neighbors]
    return nn_idx

def get_kneighbors_2DTensor(point_set, nn_idx):
    """Construct neighbors feature for each point
        Args:
        point_set: (num_points, dim)
        nn_idx: (num_points, k_num)
        num_points: the number of point
        k_num: the number of neighbor

        Returns:
        neighbors features: (num_points, k_num, dim)
      """
    shape2point_set = point_set.get_shape().as_list()
    assert(len(shape2point_set) == 2)
    point_set_neighbors = tf.gather(point_set, nn_idx)
    return point_set_neighbors

if __name__ == "__main__":
    num2points = 50
    kneighbor = 3
    indim = 2
    outdim = 1
    hiddens = (4, 8, 16)
    # x0 = np.array([[1, 2],
    #               [2, 3],
    #               [3, 5],
    #               [8, 9],
    #               [5, 2],
    #               [4, 1],
    #               [3, 2],
    #               [9, 3],
    #               [8, 5],
    #               [7, 1]], dtype=np.float)
    x0 = np.array([[1, 2],
                   [2, 3],
                   [3, 5],
                   [8, 9],
                   [5, 2]], dtype=np.float)
    print(x0)
    # x_point = np.random.rand(num2points, indim)
    # X = tf.Variable(initial_value=x_point, trainable=True, name='X')
    X = tf.Variable(initial_value=x0, trainable=True, name='X')
    print(X)
    adj_mat = pairwise_distance(X)
    print('adj_mat:', adj_mat)
    k_index = knn_excludeself(adj_mat, k=kneighbor)
    print('k_index:', k_index)
    neighbors = get_kneighbors_2DTensor(X, k_index)
    print('neighbors[0]', neighbors[0, :, :])

结果：
<tf.Variable ‘X:0’ shape=(5, 2) dtype=float64, numpy=
array([[1., 2.],
[2., 3.],
[3., 5.],
[8., 9.],
[5., 2.]])>
adj_mat: tf.Tensor(
[[ 0. 2. 13. 98. 16.]
[ 2. 0. 5. 72. 10.]
[13. 5. 0. 41. 13.]
[98. 72. 41. 0. 58.]
[16. 10. 13. 58. 0.]], shape=(5, 5), dtype=float64)
k_index: tf.Tensor(
[[1 2 4]
[0 2 4]
[1 0 4]
[2 4 1]
[1 2 0]], shape=(5, 3), dtype=int32)
neighbors[0] tf.Tensor(
[[2. 3.]
[3. 5.]
[5. 2.]], shape=(3, 2), dtype=float64)