numpy-only point set registration (Arun et. al 1987)

recover_6D_transform.py

import numpy as np

def recover_6D_transform(pts_a, pts_b):
    """
    Returns least squares transform between the two point
    sets pts_a and pts_b: T_b2a (Arun et. al 1987)
    """

    assert np.shape(pts_a) == np.shape(pts_b), "Input data must have same shape"
    assert np.shape(pts_a)[1] == 3, "Expecting points as (N,3) matrix"
    
    p_centroid = np.mean(pts_a, axis=0).reshape(-1,1)
    pp_centroid = np.mean(pts_b, axis=0).reshape(-1,1)
    
    qs = [np.reshape(pi, (-1,1)) - p_centroid for pi in pts_a]
    qps = [np.reshape(ppi, (-1,1)) - pp_centroid for ppi in pts_b]
    
    H = np.zeros((3,3))
    for qi, qpi in zip(qs, qps):
        H += qi @ qpi.T
    
    U, _, Vt = np.linalg.svd(H)
    X = Vt.T @ U.T
    assert np.isclose(np.linalg.det(X), 1), "Determinant is off!"
    
    t = pp_centroid - X @ p_centroid
    
    T_b2a = np.eye(4)
    T_b2a[:3,:3] = X
    T_b2a[:3,3] = t.flatten()
    
    return T_b2a


### Validation
from pysixd import transform

## Input data
# create known point set
ps = np.random.uniform(size=(5, 3))

# apply some random transform, get second set
R_random = transform.random_rotation_matrix()[:3,:3]
t_random = np.array([1,2,3])
pps = (R_random @ ps.T).T + t_random

## Transformation recovery
T_pp2p = recover_6D_transform(ps, pps)

# convert first point set to homogenous coords, so we can apply T_pp2p
ps_hom = np.hstack([ps, np.ones((len(ps), 1))])

# check that computed transform is correct
pps_computed = (T_pp2p @ ps_hom.T).T
assert np.allclose(pps_computed[:,:3], pps)