Chamfer distance between two point clouds in tensorflow

Question

I am trying to implement chamfer distance in tensorflow.

But, my code is taking input as numpy array. To convert a numpy into a tensor, we need to run a session, but the process is already in another session. I think two sessions can't run in parallel.

So, can anyone help me with the implementation of chamfer distance in tensorflow or help me with this problem of two simultaneous sessions?

my code is:

def chamfer_distance(array1,array2):
    # final = 0
    # final = tf.cast(final,tf.float32)
    batch_size = array1.get_shape()[0].value
    num_point = array1.get_shape()[1].value
    sess = tf.Session()
    arr1,arr2 = sess.run([array1,array2])
    del sess
    dist = 0
    for i in range(batch_size):
        tree1 = KDTree(arr1[i], leafsize=num_point+1)
        tree2 = KDTree(arr2[i], leafsize=num_point+1)
        distances1, _ = tree1.query(arr2[i])
        distances2, _ = tree2.query(arr1[i])
        distances1 = tf.convert_to_tensor(distances1)
        distances2 = tf.convert_to_tensor(distances2)
        av_dist1 = tf.reduce_mean(distances1)
        av_dist2 = tf.reduce_mean(distances2)
        dist = dist + (av_dist1+av_dist2)/batch_size
    return dist

Answer 1

I've implemented TF version of chamfer distance:

def distance_matrix(array1, array2):
    """
    arguments: 
        array1: the array, size: (num_point, num_feature)
        array2: the samples, size: (num_point, num_feature)
    returns:
        distances: each entry is the distance from a sample to array1
            , it's size: (num_point, num_point)
    """
    num_point, num_features = array1.shape
    expanded_array1 = tf.tile(array1, (num_point, 1))
    expanded_array2 = tf.reshape(
            tf.tile(tf.expand_dims(array2, 1), 
                    (1, num_point, 1)),
            (-1, num_features))
    distances = tf.norm(expanded_array1-expanded_array2, axis=1)
    distances = tf.reshape(distances, (num_point, num_point))
    return distances

def av_dist(array1, array2):
    """
    arguments:
        array1, array2: both size: (num_points, num_feature)
    returns:
        distances: size: (1,)
    """
    distances = distance_matrix(array1, array2)
    distances = tf.reduce_min(distances, axis=1)
    distances = tf.reduce_mean(distances)
    return distances

def av_dist_sum(arrays):
    """
    arguments:
        arrays: array1, array2
    returns:
        sum of av_dist(array1, array2) and av_dist(array2, array1)
    """
    array1, array2 = arrays
    av_dist1 = av_dist(array1, array2)
    av_dist2 = av_dist(array2, array1)
    return av_dist1+av_dist2

def chamfer_distance_tf(array1, array2):
    batch_size, num_point, num_features = array1.shape
    dist = tf.reduce_mean(
               tf.map_fn(av_dist_sum, elems=(array1, array2), dtype=tf.float64)
           )
    return dist

And for validation purpose, I also implemented a sklearn version:

def chamfer_distance_sklearn(array1,array2):
    batch_size, num_point = array1.shape[:2]
    dist = 0
    for i in range(batch_size):
        tree1 = KDTree(array1[i], leaf_size=num_point+1)
        tree2 = KDTree(array2[i], leaf_size=num_point+1)
        distances1, _ = tree1.query(array2[i])
        distances2, _ = tree2.query(array1[i])
        av_dist1 = np.mean(distances1)
        av_dist2 = np.mean(distances2)
        dist = dist + (av_dist1+av_dist2)/batch_size
    return dist

Also a numpy version:

def array2samples_distance(array1, array2):
    """
    arguments: 
        array1: the array, size: (num_point, num_feature)
        array2: the samples, size: (num_point, num_feature)
    returns:
        distances: each entry is the distance from a sample to array1 
    """
    num_point, num_features = array1.shape
    expanded_array1 = np.tile(array1, (num_point, 1))
    expanded_array2 = np.reshape(
            np.tile(np.expand_dims(array2, 1), 
                    (1, num_point, 1)),
            (-1, num_features))
    distances = LA.norm(expanded_array1-expanded_array2, axis=1)
    distances = np.reshape(distances, (num_point, num_point))
    distances = np.min(distances, axis=1)
    distances = np.mean(distances)
    return distances

def chamfer_distance_numpy(array1, array2):
    batch_size, num_point, num_features = array1.shape
    dist = 0
    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])
        dist = dist + (av_dist1+av_dist2)/batch_size
    return dist

You can validate the result using following script:

batch_size = 8
num_point = 20
num_features = 4
np.random.seed(1)
array1 = np.random.randint(0, high=4, size=(batch_size, num_point, num_features))
array2 = np.random.randint(0, high=4, size=(batch_size, num_point, num_features))

print('sklearn: ', chamfer_distance_sklearn(array1, array2))
print('numpy: ', chamfer_distance_numpy(array1, array2))

array1_tf = tf.constant(array1, dtype=tf.float64)
array2_tf = tf.constant(array2, dtype=tf.float64)
dist_tf = chamfer_distance_tf(array1_tf, array2_tf)

with tf.Session() as sess:
    print('tf: ', sess.run(dist_tf))