icp submission

icp.py

from numpy import *
import random as rndom
from plyfile import PlyData, PlyElement
from scipy import spatial #kd tree



def to_4d(R,T):
	rt = concatenate((R,T),axis=1)
	fourth = array([[ 0.,  0.,  0., 1.]])
	return list(concatenate((rt,fourth)).flat)
def to_OF(frd):
	print "move.transform(16,Matrix4x4({}))".format(", ".join(str(s) for s in frd))
def to_array(data):
	return array([[a['x'], a['y'], a['z']] for a in data])
def bary_center(points, center):
	"""center is a tuple, points is array of coords"""
	#R = [center]*len(points)
	#if len(r) !== len(center):
	#	raise Exception()
	return points - center
def select_rand_subset(points):
	"""Select a subset of points from p
	"""
	return array(rndom.sample(points, 20))

def select_from_q_nearest_to_points(p=None, q=None):
	"""Match each p to closest point q on other scan, returns closest  points from q in
	corresponding order """
	tree = spatial.KDTree(q)
	_, indices = tree.query(p)
	return array([tree.data[i] for i in indices])
def reject(p,q):
	"""
	reject worst 20% of pairs, based on point to point distance
	"""
	#array of distances
	dists = [(linalg.norm(p[i]-q[i]), i) for i in range(len(p))]
	#sort based on distance, largest first
	dists.sort(key=lambda x: x[0], reverse=True)
	#calculate 20% of length
	num_remove = int(.2 * len(q))
	#get rid of the rest
	dists = dists[:num_remove]
	#return sorted(dists, key=lambda x: x[1], reverse=True)

	for _ , i in sorted(dists, key=lambda x: x[1], reverse=True):
		q = delete(q, i, 0)
		p = delete(p, i, 0)
	return p, q
	#delete(q, 2, 0)
	#linalg.norm(p[i]-b[])
def centre_of_mass(data):
	"Returns the coordinates R of the center of mass of the point cloud in 'data'"
	x = data[:,0]; y = data[:,1]; z = data[:,2]
	return array([x.mean(),y.mean(),z.mean()], dtype=data.dtype)

def least_squares(p,q):
	"""Returns least squares linear transformation"""
	if len(p) != len(q):
		raise Exception("lengths dont match")
	sum1 = zeros((3,3))
	sum2 = zeros((3,3))
	for i in range(len(q)):
		#q is allready in "transposed" form, p needs to become a column vector
		sum1 += p[i].reshape(3,1) * q[i]
		sum2 += q[i].reshape(3,1) * q[i]
	sum2 = matrix(sum2).I
	A = sum1 * sum2
	return A
def polar_decomposition(A):
	"""returns the rotation by finding polar decomposition of A"""
	#first find the svd
	U, s, V = linalg.svd(A)
	R = U*V
	return R
def find_translation(R, p_center, q_center):
#	t = p - Rq
	return p_center.reshape(3,1) - (R*q_center.reshape(3,1))
def update_alignment(R=None,T=None,q=None):
	#new_alignment = Rq(i) +t with some format munging
	new_set = [array((R*(i.reshape(3,1)) + T).reshape(1,3)).reshape(-1) for i in q]
	return array(new_set)
def calculate_minimisation_metric(R,T,q=False, p=False):
	return sum([pow(linalg.norm(p[i].reshape(3,1) - R*q[i].reshape(3,1) - T),2) for i in range(len(q))])

def main(p_all=False, q_all=False):
	if q_all is False: #q is being registered to reference p
		q_all = PlyData.read(open('/Users/md/Documents/UCL 3/geometry processing/cwk1/python/data/bun270.ply')).elements[0].data
	if p_all is False:
		p_all = PlyData.read(open('/Users/md/Documents/UCL 3/geometry processing/cwk1/python/data/bun180.ply')).elements[0].data
	#just some data cleaning/munging
	q_all = to_array(q_all)
	p_all = to_array(p_all)

	R = identity(3)
	T = zeros((3,1))
	variance = []
	iterations = 20
	for i in range(iterations):
		p = select_rand_subset(p_all)
		q = select_from_q_nearest_to_points(p=p,q=q_all)
		p, q = reject(p,q)
		p_bar = centre_of_mass(p)
		p = bary_center(p, p_bar)
		#compute p bar and q bar
		q_bar = centre_of_mass(q)
		q = bary_center(q, q_bar)
		variance.append(calculate_minimisation_metric(R,T,q=q, p=p))
		R = polar_decomposition(least_squares(p,q))
		T = find_translation(R,p_bar, q_bar)
		#print "4d"
		to_OF(to_4d(R,T))

		q_all = update_alignment(R=R,T=T,q=q_all)
	print "Variance: {}".format(variance)

if __name__ == "__main__":
    main()