Source code for prepmd.lib.icp

"""
ICP method in python by Clay Fannigan. Improvement by Max Bazik. This version
originally had scaling, added by Alvin Wan. I have removed the scalilng but
kept some changes that allow for differently-sized point clouds to be aligned.
This code was originally licensed under the Apache License version 2.0.
"""

import numpy as np
from sklearn.neighbors import NearestNeighbors



def best_fit_transform(A, B):
    '''
    Calculates the least-squares best-fit transform between corresponding 3D points A->B
    Input:
      A: Nx3 numpy array of corresponding 3D points
      B: Nx3 numpy array of corresponding 3D points
    Returns:
      T: 4x4 homogeneous transformation matrix
      R: 3x3 rotation matrix
      t: 3x1 column vector for translation
      s: 3x1 column vector for scaling
    '''

    assert len(A) == len(B)

    # translate points to their centroids
    centroid_A = np.mean(A, axis=0)
    centroid_B = np.mean(B, axis=0)
    AA = A - centroid_A
    BB = B - centroid_B

    # rotation matrix
    H = np.dot(AA.T, BB)
    U, S, Vt = np.linalg.svd(H)
    R = np.dot(Vt.T, U.T)

    # special reflection case
    if np.linalg.det(R) < 0:
       Vt[2,:] *= -1
       R = np.dot(Vt.T, U.T)

    # translation
    t = centroid_B.T - np.dot(R, centroid_A.T)

    # homogeneous transformation
    T = np.identity(4)
    T[0:3, 0:3] = R
    T[0:3, 3] = t

    return T, R, t




def nearest_neighbor(src, dst):
    '''
    Find the nearest (Euclidean) neighbor in dst for each point in src
    Input:
        src: Nx3 array of points
        dst: Nx3 array of points
    Output:
        distances: Euclidean distances of the nearest neighbor
        indices: dst indices of the nearest neighbor
    '''

    neigh = NearestNeighbors(n_neighbors=1)
    neigh.fit(dst)
    distances, indices = neigh.kneighbors(src, return_distance=True)
    return distances.ravel(), indices.ravel()




def icp(A, B, init_pose=None, max_iterations=100, tolerance=1e-10):
    '''
    The Iterative Closest Point method
    Input:
        A: Nx3 numpy array of source 3D points
        B: Nx3 numpy array of destination 3D point
        init_pose: 4x4 homogeneous transformation
        max_iterations: exit algorithm after max_iterations
        tolerance: convergence criteria
    Output:
        T: final homogeneous transformation
        distances: Euclidean distances (errors) of the nearest neighbor
    '''

    # make points homogeneous, copy them so as to maintain the originals
    src = np.ones((4,A.shape[0]))
    dst = np.ones((4,B.shape[0]))
    src[0:3,:] = np.copy(A.T)
    dst[0:3,:] = np.copy(B.T)

    # apply the initial pose estimation
    if init_pose is not None:
        src = np.dot(init_pose, src)

    prev_error = 0

    try:

        for i in range(max_iterations):
            # find the nearest neighbours between the current source and destination points
            distances, indices = nearest_neighbor(src[0:3,:].T, dst[0:3,:].T)

            # compute the transformation between the current source and nearest destination points
            T, _, _ = best_fit_transform(src[0:3,:].T, dst[0:3,indices].T)

            # update the current source
            src = T.dot(src)

            # check error
            mean_error = np.sum(distances) / distances.size
            if abs(prev_error-mean_error) < tolerance:
                break
            prev_error = mean_error

        # calculate final transformation
        T, _, _ = best_fit_transform(A, src[0:3,:].T)

        return T, distances
    except ValueError as e:
        print(e)
        return np.eye(4), 1, np.array([np.inf])
    except np.linalg.linalg.LinAlgError as e:
        print(e)
        return np.eye(4), 1, np.array([np.inf])