3D点云全局注册RANSAC算法Python实现

3D pointcloud global registration RANSAC 算法的python实现，除了基本的库函数外，没有实现现成的RANSAC函数。Python实现这个算法运行效率特别低，跑起来很慢，效果也没能比得上库函数。下面附上效果图和代码：

import cv2
import open3d as o3d
from matplotlib import pyplot as pyplot
import numpy as np
import copy
import scipy
from scipy import spatial 
import random
import sys
import math

#Kabsch Algorithm
def compute_transformation(source,target):
    #Normalization
    number = len(source)
    #the centroid of source points
    cs = np.zeros((3,1))
    #the centroid of target points
    ct = copy.deepcopy(cs)
    cs[0] = np.mean(source[:][0]);cs[1]=np.mean(source[:][1]);cs[2]=np.mean(source[:][2])
    ct[0] = np.mean(target[:][0]);cs[1]=np.mean(target[:][1]);cs[2]=np.mean(target[:][2])
    #covariance matrix
    cov = np.zeros((3,3))
    #translate the centroids of both models to the origin of the coordinate system (0,0,0)
    #subtract from each point coordinates the coordinates of its corresponding centroid
    for i in range(number):
        sources = source[i].reshape(-1,1)-cs
        targets = target[i].reshape(-1,1)-ct
        cov = cov + np.dot(sources,np.transpose(targets))
    #SVD (singular values decomposition)
    u,w,v = np.linalg.svd(cov)
    #rotation matrix
    R = np.dot(u,np.transpose(v))
    #Transformation vector
    T = ct - np.dot(R,cs)
    return R, T

#compute the transformed points from source to target based on the R/T found in Kabsch Algorithm
def _transform(source,R,T):
    points = []
    for point in source:
        points.append(np.dot(R,point.reshape(-1,1)+T))
    return points

#compute the root mean square error between source and target
def compute_rmse(source,target,R,T):
    rmse = 0
    number = len(target)
    points = _transform(source,R,T)
    for i in range(number):
        error = target[i].reshape(-1,1)-points[i]
        rmse = rmse + math.sqrt(error[0]**2+error[1]**2+error[2]**2)
    return rmse

def draw_registrations(source, target, transformation = None, recolor = False):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    if(recolor): # recolor the points
        source_temp.paint_uniform_color([1, 0.706, 0])
        target_temp.paint_uniform_color([0, 0.651, 0.929])
    if(transformation is not None): # transforma source to targets
        source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp])

def pc2array(pointcloud):
    return np.asarray(pointcloud.points)

def registration_RANSAC(source,target,source_feature,target_feature,ransac_n=3,max_iteration=100000,max_validation=100):
    #the intention of RANSAC is to get the optimal transformation between the source and target point cloud
    s = pc2array(source) #(4760,3)
    t = pc2array(target)
    #source features (33,4760)
    sf = np.transpose(source_feature.data)
    tf = np.transpose(target_feature.data)
    #create a KD tree
    tree = spatial.KDTree(tf)
    corres_stock = tree.query(sf)[1]
    for i in range(max_iteration):
        #take ransac_n points randomly
        idx = [random.randint(0,s.shape[0]-1) for j in range(ransac_n)]
        corres_idx = corres_stock[idx]
        source_point = s[idx,...]
        target_point = t[corres_idx,...]
        #estimate transformation
        #use Kabsch Algorithm
        R, T = compute_transformation(source_point,target_point)
        #calculate rmse for all points
        source_point = s
        target_point = t[corres_stock,...]
        rmse = compute_rmse(source_point,target_point,R,T)
        #compare rmse and optimal rmse and then store the smaller one as optimal values
        if not i:
            opt_rmse = rmse
            opt_R = R
            opt_T = T
        else:
            if rmse < opt_rmse:
                opt_rmse = rmse
                opt_R = R
                opt_T = T
    return opt_R, opt_T

#used for downsampling
voxel_size = 0.05
#this is to get the fpfh features, just call the library
def get_fpfh(cp):
    cp = cp.voxel_down_sample(voxel_size)
    cp.estimate_normals()
    return cp, o3d.registration.compute_fpfh_feature(cp, o3d.geometry.KDTreeSearchParamHybrid(radius=5, max_nn=100))

if "__name__" == "__main__":
    source = o3d.io.read_point_cloud('global_registration/r1.pcd')
    target = o3d.io.read_point_cloud('global_registration/r2.pcd')   
    #if we want to use RANSAC registration, get_fpfh features should be acquired firstly                 
    r1, f1 = get_fpfh(source)
    r2, f2 = get_fpfh(target)
    R, T = registration_RANSAC(r1,r2,f1,f2)
    #transformation matrix is formed by R, T based on np.hstack and np.vstack(corporate two matrices by rows)
    #Notice we need add the last row [0 0 0 1] to make it homogeneous 
    transformation = np.vstack((np.hstack((np.float64(R), np.float64(T))), np.array([0,0,0,1])))
    draw_registrations(r1, r2, transformation, True)