ICP

 avatar
unknown
python
a month ago
8.1 kB
5
Indexable
import open3d as o3d
import copy
import numpy as np

def draw_registration_result(source, target, transformation):
    """
    Visualizes how well two point clouds are aligned after ICP.
    This function creates a visualization window showing both point clouds:
    - Source points in yellow (after transformation)
    - Target points in blue
    
    Parameters:
        source: The point cloud we're trying to align (moving cloud)
        target: The point cloud we're trying to align to (static cloud)
        transformation: 4x4 matrix showing how to move source to align with target
    """
    # Create copies to avoid modifying the original point clouds
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    
    # Color the point clouds differently so we can tell them apart
    source_temp.paint_uniform_color([1, 0.706, 0])      # Yellow color for source
    target_temp.paint_uniform_color([0, 0.651, 0.929])  # Blue color for target
    
    # Apply the transformation to the source point cloud
    source_temp.transform(transformation)
    
    # Show both point clouds in a visualization window
    # zoom: How close to view the point clouds
    # front: Direction we're looking from
    # lookat: Point we're looking at
    # up: Which direction is up in the view
    o3d.visualization.draw_geometries([source_temp, target_temp],
                                    zoom=0.4459,
                                    front=[0.9288, -0.2951, -0.2242],
                                    lookat=[1.6784, 2.0612, 1.4451],
                                    up=[-0.3402, -0.9189, -0.1996])

def find_nearest_neighbors(source_pc, target_pc, nearest_neigh_num):
    """
    For each point in the target cloud, finds the closest point in the source cloud.
    This is a crucial step in ICP - we need to know which points correspond to each other.
    
    Parameters:
        source_pc: Point cloud we're moving
        target_pc: Point cloud we're trying to match
        nearest_neigh_num: How many nearest neighbors to find (usually 1)
    
    Returns:
        Array of closest points from source for each target point
    """
    # Create a KD-tree from the source points for efficient nearest neighbor search
    # A KD-tree is like a spatial index - it helps us find closest points much faster
    point_cloud_tree = o3d.geometry.KDTreeFlann(source_pc)
    
    # Array to store the matching points we find
    points_arr = []
    
    # For each point in the target cloud
    for point in target_pc.points:
        # Find the nearest neighbor in the source cloud
        # Returns: [number of neighbors found, indices of neighbors, squared distances]
        [_, idx, _] = point_cloud_tree.search_knn_vector_3d(point, nearest_neigh_num)
        
        # Add the closest point to our array
        points_arr.append(source_pc.points[idx[0]])
        
    return np.asarray(points_arr)

def icp(source, target):
    """
    Implements the Iterative Closest Point algorithm.
    This algorithm tries to find the best way to align two point clouds by:
    1. Finding matching points between the clouds
    2. Computing the best transformation to align these matches
    3. Applying the transformation and repeating until the alignment is good enough
    
    Parameters:
        source: Point cloud we want to move
        target: Point cloud we want to align to
    
    Returns:
        4x4 transformation matrix that best aligns source to target
    """
    # Set initial colors for visualization
    source.paint_uniform_color([0.5, 0.5, 0.5])  # Gray for source
    target.paint_uniform_color([0, 0, 1])        # Blue for target
    
    # Convert target points to numpy array for easier math
    target_points = np.asarray(target.points)
    
    # Apply an initial transformation to source
    # This can help avoid getting stuck in local minima
    # The matrix includes both rotation and translation
    initial_transform = np.asarray([
        [0.862, 0.011, -0.507, 0.5],   # First row  (rotation + x translation)
        [-0.139, 0.967, -0.215, 0.7],  # Second row (rotation + y translation)
        [0.487, 0.255, 0.835, -1.4],   # Third row  (rotation + z translation)
        [0.0, 0.0, 0.0, 1.0]           # Fourth row (always [0,0,0,1] for homogeneous coordinates)
    ])
    source = source.transform(initial_transform)

    # Initialize variables for the iterative process
    curr_iteration = 0                    # Keep track of how many iterations we've done
    cost_change_threshold = 0.001         # Stop when improvement is smaller than this
    curr_cost = 1000                      # Current alignment error
    prev_cost = 10000                     # Previous alignment error

    while True:
        # STEP 1: Find closest points between source and target
        new_source_points = find_nearest_neighbors(source, target, 1)

        # STEP 2: Find the centers (centroids) of both point sets
        source_centroid = np.mean(new_source_points, axis=0)
        target_centroid = np.mean(target_points, axis=0)
        
        # Center both point clouds by subtracting their centroids
        # This makes it easier to find the optimal rotation
        source_repos = np.asarray([new_source_points[ind] - source_centroid 
                                 for ind in range(len(new_source_points))])
        target_repos = np.asarray([target_points[ind] - target_centroid 
                                 for ind in range(len(target_points))])

        # STEP 3: Find the optimal rotation using SVD (Singular Value Decomposition)
        # First, compute the cross-covariance matrix
        cov_mat = target_repos.transpose() @ source_repos
        
        # Perform SVD decomposition
        # U and Vt contain the rotation information we need
        U, X, Vt = np.linalg.svd(cov_mat)
        
        # Compute the optimal rotation matrix
        R = U @ Vt
        
        # Compute the optimal translation
        # This moves the rotated source centroid to the target centroid
        t = target_centroid - R @ source_centroid
        t = np.reshape(t, (1,3))
        
        # Calculate how good our current alignment is
        # Lower cost means better alignment
        curr_cost = np.linalg.norm(target_repos - (R @ source_repos.T).T)
        print("Current alignment error =", curr_cost)

        # Check if we've improved enough to continue
        if ((prev_cost - curr_cost) > cost_change_threshold):
            # If yes, update everything and continue
            prev_cost = curr_cost
            
            # Create a 4x4 homogeneous transformation matrix
            # This combines rotation and translation into one matrix
            transform_matrix = np.hstack((R, t.T))  # Combine R and t horizontally
            transform_matrix = np.vstack((transform_matrix, np.array([0, 0, 0, 1])))
            
            # Apply the transformation to the source point cloud
            source = source.transform(transform_matrix)
            curr_iteration += 1
        else:
            # If we haven't improved enough, stop iterating
            break

    print("\nCompleted after", curr_iteration, "iterations")
    
    # Show the final result
    draw_registration_result(source, target, transform_matrix)
    return transform_matrix

# Test the ICP implementation on example point clouds

### PART A: Demo point clouds from Open3D ###
demo_icp_pcds = o3d.data.DemoICPPointClouds()
source = o3d.io.read_point_cloud(demo_icp_pcds.paths[0])
target = o3d.io.read_point_cloud(demo_icp_pcds.paths[1])
part_a = icp(source, target)

### PART B: KITTI dataset point clouds ###
# KITTI is a popular dataset for autonomous driving
source = o3d.io.read_point_cloud("kitti_frame1.pcd")
target = o3d.io.read_point_cloud("kitti_frame2.pcd")
print("Point cloud sizes:", len(source.points), len(target.points))
part_b = icp(source, target)
Leave a Comment