Some Scraps for RGB to Point Cloud Registration

rgb_pointcloud_registration_scraps.py

# I never got to finish this fully, but it's most of the way there for full non-obstructed charuco board views
# Maybe one day will be useful... 
# Although retrospectively, I kind of wish I just projected the RGB points to 3D space and just
# did ICP with the point cloud. Not sure if I'd need to segment out non-plane points in the cloud for that still though...
import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
import cv2
from typing import List, Dict, Optional
from copy import deepcopy
from scipy.optimize import minimize
from scipy.spatial import ConvexHull
import spatialmath.base as smb
import spatialmath as sm
from spatialmath import SO3, SE3

def add_random_color_to_pcd(pcd):
    """Add random color to a point cloud.
    
    Args:
        pcd (o3d.geometry.PointCloud): The input point cloud.
    
    Returns:
        o3d.geometry.PointCloud: The colored point cloud.
    """
    color = np.random.rand(3)
    colors = np.tile(color, (np.asarray(pcd.points).shape[0], 1))
    pcd.colors = o3d.utility.Vector3dVector(colors)
    return pcd

def create_polyline(points, close_line=True):
    if len(points) < 2:
        raise ValueError("Polyline must have at least two points.")
    
    # Ensure points are in numpy array format
    points = np.array(points, dtype=np.float64)
    if points.shape[1] != 3:
        raise ValueError("Each point must have 3 coordinates (x, y, z).")
    
    # Create the lines connecting each point to the next
    lines = [[i, i+1] for i in range(len(points) - 1)]
    
    # If close_line is True, add a line connecting the last point to the first
    if close_line:
        lines.append([len(points) - 1, 0])
    
    colors = [[1, 0, 0] for _ in range(len(lines))]  # Red color for the lines
    
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(points)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    line_set.colors = o3d.utility.Vector3dVector(colors)
    
    return line_set

def draw_frame(transform, size=0.1):
    """Create an Open3D geometry representing an SE(3) frame."""
    mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size)
    mesh_frame.transform(transform)
    return mesh_frame

def visualize_multiple_point_clouds(pcds, polyline_points=None, frames=None):
    """Visualize multiple point clouds with random colors, an optional polyline, and optional SE(3) frames.
    
    Args:
        pcds (list of o3d.geometry.PointCloud): List of point clouds to visualize.
        polyline_points (list of list or np.array, optional): List of points [x, y, z] that make up the optional polyline.
        frames (list of np.array, optional): List of SE(3) frames (4x4 transformation matrices) to draw.
    """
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    
    for pcd in pcds:
        colored_pcd = add_random_color_to_pcd(pcd)
        vis.add_geometry(colored_pcd)
    
    if polyline_points is not None:
        
        polyline_shape = np.array(polyline_points).shape
        
        if polyline_shape[0] == 1:
            line_set = create_polyline(polyline_points[0])
            vis.add_geometry(line_set)
        elif polyline_shape[0] > 1:
            for line in polyline_points:
                line_set = create_polyline(line)
                vis.add_geometry(line_set)
    
    if frames is not None:
        for frame in frames:
            frame_mesh = draw_frame(frame)
            vis.add_geometry(frame_mesh)
    
    vis.run()
    vis.destroy_window()

def rotation_matrix_to_align_vectors(vec1, vec2):
    v1 = np.array(vec1) / np.linalg.norm(vec1)
    v2 = np.array(vec2) / np.linalg.norm(vec2)
    cross_product = np.cross(v1, v2)
    dot_product = np.dot(v1, v2)
    cross_product_matrix = np.array([
        [0, -cross_product[2], cross_product[1]],
        [cross_product[2], 0, -cross_product[0]],
        [-cross_product[1], cross_product[0], 0]
    ])
    rotation_matrix = np.eye(3) + cross_product_matrix + \
                      np.dot(cross_product_matrix, cross_product_matrix) * \
                      (1 / (1 + dot_product))

    return rotation_matrix

def numpy_to_pcd(xyz):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(xyz)

    return pcd

def pcd_to_numpy(pcd):
    return np.asarray(pcd.points)

def remove_noise_radius(pcd: o3d.geometry.PointCloud, 
                        nb_points: int = 16, 
                        radius: float = 0.05) -> o3d.geometry.PointCloud:
        """Remove point clouds noise using radius outlier removal method."""
        cl, ind = pcd.remove_radius_outlier(nb_points=nb_points, radius=radius)
        return cl
    
def cautiously_clean_pointcloud(pcd: o3d.geometry.PointCloud,
                          number_neighbors_for_radius_removal: int = 16, 
                          radius: float = .05):
    def remove_nan(pcd: o3d.geometry.PointCloud) -> o3d.geometry.PointCloud:
        """Remove nan values from point clouds."""
        points = np.asarray(pcd.points)
        pcd.points = o3d.utility.Vector3dVector(points[~np.isnan(points[:, 0])])
        return pcd
    
    no_nan_pcd = remove_nan(pcd)
    no_isolated_points_pcd = remove_noise_radius(no_nan_pcd)
    
    return no_isolated_points_pcd

def remove_noise_statistical(pcd: o3d.geometry.PointCloud, nb_neighbors: int = 20, std_ratio: float = 2.0) -> o3d.geometry.PointCloud:
    cl, ind = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
    return cl

def voxel_down_sample(pcd: o3d.geometry.PointCloud, voxel_size=0.003) -> o3d.geometry.PointCloud:
    return pcd.voxel_down_sample(voxel_size=voxel_size)

def random_down_sample(pcd: o3d.geometry.PointCloud, preservation_ratio=0.5) -> o3d.geometry.PointCloud:
    return pcd.random_down_sample(sampling_ratio=preservation_ratio)

def compute_plane_equation_from_obox(obox):
    # Get the normal vector (z-axis of the bounding box rotation matrix)
    normal = obox.R[:, 2]
    # Get a point on the plane (center of the bounding box)
    point_on_plane = obox.center
    # Plane equation: ax + by + cz + d = 0
    d = -np.dot(normal, point_on_plane)
    plane_equation = np.append(normal, d)
    return plane_equation

def get_bbox_vertices_above_plane(bbx, plane_eq=[0, 0, 1, 0]): 
    vertices = bbx.get_box_points()
    
    def check_if_point_above_plane(point, plane=plane_eq):
        A, B, C, D = plane
        x, y, z = point

        result = A * x + B * y + C * z + D

        return result >= 0
    above_plane_vertices = list(filter(lambda x: check_if_point_above_plane(x), vertices))

    above_plane_vertices = sorted(above_plane_vertices, 
                                        key=lambda x: np.linalg.norm(x - above_plane_vertices[0]))
    bbx = numpy_to_pcd(above_plane_vertices)
    return pcd_to_numpy(bbx)

def detect_multi_planes(pcd: o3d.geometry.PointCloud, min_ratio: float = 0.05, 
                        normal_variance_threshold_deg: float = 30, 
                        coplanarity_deg: float = 60, 
                        outlier_ratio: float = 0.75, 
                        min_plane_edge_length: float = 0,#min(rectangle_height, rectangle_width), 
                        min_num_points: int = 0,#MIN_PLANE_PTS, 
                        knn: int = 30):
    """    
    Detect multiple planes from given point clouds using Open3D's detect_planar_patches.
    """
    # Ensure normals are estimated
    if not pcd.has_normals():
        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=knn))
    
    # Detect planar patches
    oboxes = pcd.detect_planar_patches(
        normal_variance_threshold_deg=normal_variance_threshold_deg,
        coplanarity_deg=coplanarity_deg,
        outlier_ratio=outlier_ratio,
        min_plane_edge_length=min_plane_edge_length,
        min_num_points=min_num_points,
        search_param=o3d.geometry.KDTreeSearchParamKNN(knn=knn)
    )
    
    plane_list = []
    for obox in oboxes:
        inlier_cloud = pcd.crop(obox)        
        plane_equation = compute_plane_equation_from_obox(obox)
        plane_list.append((plane_equation, inlier_cloud, obox))
    
    return plane_list

  def project_point_to_plane(point, plane):
    A, B, C, D = plane
    point = np.array(point)
    normal = np.array([A, B, C])
    distance = (np.dot(normal, point) + D) / np.linalg.norm(normal)
    projected_point = point - distance * (normal / np.linalg.norm(normal))
    return projected_point
 
def select_checkerboard_from_candidates(plane_eqs_and_inliers_and_bboxs,
                                        rectangle_height: float = rectangle_height,
                                        rectangle_width: float = rectangle_width,
                                        fit_tolerance: float = RECT_FIT_TOL,
                                        min_num_points: int = MIN_PLANE_PTS,
                                       unit_vec = [0, 0, 1], 
                                       dim_tol = .02):
    results = []

    for (plane_eq, plane_inliers, bbox) in plane_eqs_and_inliers_and_bboxs:
        plane_vec = plane_eq[:3]
        align_plane_normal_to_z_axis_rot_mat = rotation_matrix_to_align_vectors(plane_vec, unit_vec)
        bbox_vertices_above_plane = get_bbox_vertices_above_plane(bbox, plane_eq)
        
        bbox_vertices_projected_to_plane = [project_point_to_plane(p, plane_eq) for p in bbox_vertices_above_plane]
        
        projected_corners = numpy_to_pcd(bbox_vertices_above_plane)
        projected_corners = projected_corners.rotate(align_plane_normal_to_z_axis_rot_mat)
        projected_corners = pcd_to_numpy(projected_corners)[:, :2]
        
        projected_corners_without_opp = projected_corners[:-1]
        dim = sorted([np.linalg.norm(projected_corners_without_opp[-1] - \
                                     projected_corners_without_opp[0]),
                     np.linalg.norm(projected_corners_without_opp[-2] - \
                                     projected_corners_without_opp[0])])
                                     
        results.append({"corners_3d": bbox_vertices_projected_to_plane,
                        "center_3d": np.asarray(bbox.center),
                        "dim": dim,
                        "align_plane_normal_to_z_axis_rot_mat": align_plane_normal_to_z_axis_rot_mat,
                        "plane_eq": plane_eq,
                        "inliers": plane_inliers,
                        "corners_2d": projected_corners,
                        "center_2d": np.mean(projected_corners, axis=0)})
        
    ref_dim = sorted([rectangle_height, rectangle_width])
    compare_dim = lambda x: ((ref_dim[0] - x[0])**2 + (ref_dim[1] - x[1])**2) ** .5
    dim_tol = (dim_tol ** 2 + dim_tol ** 2) ** .5
    
    results = sorted(results, key = lambda x: compare_dim(x['dim']))
    
    if len(results) == 0 or compare_dim(results[0]['dim']) > dim_tol:
        return None
    else:
        return results[0]
      
# cloud_calib = {}

# for idx, image_data in enumerate(image_data_list[play_slice]):
#     tof_pcd = o3d.geometry.PointCloud.create_from_depth_image(depth=image_data.depth,
#                                                               depth_scale=1000,
#                                     intrinsic=image_data.unregistered_depth_intrinsic.to_o3d())
#     cleaned_tof_pcd = cautiously_clean_pointcloud(tof_pcd, 
#                                                   number_neighbors_for_radius_removal=4,
#                                                   radius=.03)
    
#     plane_eqs_and_inliers_and_bboxs = detect_multi_planes(cleaned_tof_pcd)
#     result = select_checkerboard_from_candidates(plane_eqs_and_inliers_and_bboxs)
    
#     if result is not None:
#         result["cloud"] = cleaned_tof_pcd
#         cloud_calib[idx] = result
def assign_pose_to_cloud(cloud_calib, 
                         rgb_pose_est,
                         alignment_heuristic = RPY_HEURISTIC,
                         t_guess = T_HEURISTIC):
    
    rgb_t_depth = SE3.Rt(R=SO3.RPY(*alignment_heuristic),
                         t=t_guess)
    rgb_t_depth_rot = SO3.RPY(*alignment_heuristic)
    rgb_t_board = SE3(rgb_pose_est['pose'])
    rgb_t_board_rot = SO3(rgb_pose_est['pose'][:3,:3])
    
    
    depth_t_board_est_from_rgb = rgb_t_depth.inv() * rgb_t_board

    cloud_plane_normal = cloud_calib['plane_eq'][:-1]
    cloud_plane_normal /= np.linalg.norm(cloud_plane_normal) # normalize jic
    rgb_plane_normal = rgb_pose_est['pose'][:3,2]
    
    
    if np.dot(cloud_plane_normal, rgb_plane_normal) < 0:
        cloud_plane_normal *= -1 # reverse the direction to line things up 
    
    
    random_corner = cloud_calib['corners_3d'][0]
    distances = [np.linalg.norm(random_corner - c) for c in cloud_calib['corners_3d']]
    hypot_idx = np.argmax(np.array(distances))
    
    bordering_corners = []
    
    for idx in range(1, 4):
        if idx != hypot_idx:
            bordering_corners.append(cloud_calib['corners_3d'][idx])
    
    base_vectors = [c - random_corner for c in bordering_corners]
    
    base_vectors = [np.array(vec) / np.linalg.norm(vec) for vec in base_vectors] # normalize jic

    possible_basis_vectors = [base_vectors, # all possible basis to combine with normal for SO3
                             [base_vectors[0] * -1, base_vectors[1]],
                             [base_vectors[0], base_vectors[1] * -1],
                             [base_vectors[0] * -1, base_vectors[1] * -1]]
    
    rotations = []
    for vector_set in possible_basis_vectors:
        for order in [vector_set, vector_set[::-1]]:
            possible_rot_mat = np.zeros((3, 3))

            possible_rot_mat[:, 0] = order[0]
            possible_rot_mat[:, 1] = order[1]
            possible_rot_mat[:, 2] = cloud_plane_normal

            orthagonalized, upper_triangular = np.linalg.qr(possible_rot_mat) 

            for idx in range(3):
                if np.dot(orthagonalized[:, idx], possible_rot_mat[:, idx]) < 0:
                    orthagonalized[:, idx] *= -1.0

            if np.linalg.det(orthagonalized) > 0:
                try:
                    rotations.append(SO3(orthagonalized))
                except ValueError as e:
                    print(f"Rotation matrix is not properly formed: {e} ")
    
    depth_t_board_est_from_rgb_rot = SO3(depth_t_board_est_from_rgb.A[:3,:3]) # operator overload goes brr

    relative_rots = [(depth_t_board_est_from_rgb_rot.inv() * depth_t_board_rot).angvec()[0] for \
                     depth_t_board_rot in rotations]
    
    most_likely_rot = rotations[np.argmin(relative_rots)] # relative rotation here
    
    depth_camera_origin_t_charuco_center = SE3.Rt(R=most_likely_rot, t=cloud_calib["center_3d"].reshape((3,)))
    
    return {
        "pose": depth_camera_origin_t_charuco_center.A,
        "rgb_t_board_in_depth_frame_est": depth_t_board_est_from_rgb,
        "visu_rotation_matrix": cv2.Rodrigues(depth_camera_origin_t_charuco_center.A[:-1, :-1])[0], # visu 
        "visu_translation_vector": depth_camera_origin_t_charuco_center.A[:-1, -1].reshape((3, 1)) 
    }
 

# poses = {}
# for idx in calib_ids:
#     cloud_data = cloud_calib[idx]
#     rgb_pose = rgb_calib['poses'][idx]
#     cloud_pose = assign_pose_to_cloud(cloud_data, rgb_pose, alignment_heuristic = RPY_HEURISTIC,
#                                      t_guess = T_HEURISTIC)
#     poses[idx] = [rgb_pose['pose'], cloud_pose['pose'], cloud_pose["rgb_t_board_in_depth_frame_est"]]
# for jdx in calib_ids:
#     visualize_multiple_point_clouds([cloud_calib[jdx]['cloud']],
#                                    polyline_points=[[*cloud_calib[jdx]['corners_3d'],
#                                                     cloud_calib[jdx]['center_3d']]],
#                                    frames=[poses[jdx][1], poses[jdx][2]])
    
#     x = input("press c to continue")
    
#     if x != "c":
#         break

def remove_outliers_based_on_transform(poses, threshold=3):
    """
    Remove outliers from a list of 4x4 transformation matrices based on the last column values ([:-1, -1]).
    
    Parameters:
    poses (list of np.ndarray): List of 4x4 transformation matrices.
    threshold (float): Z-score threshold to determine outliers.
    
    Returns:
    np.ndarray: Cleaned list of transformation matrices with outliers removed.
    """
    # Extract the relevant values to determine outliers ([:-1, -1])
    extracted_values = np.array([pose[:-1, -1] for pose in poses], dtype=np.float32)
    
    # Calculate the mean and standard deviation for each column
    mean = np.mean(extracted_values, axis=0)
    std = np.std(extracted_values, axis=0)
    
    # Identify outliers
    z_scores = (extracted_values - mean) / std
    outliers = np.any(np.abs(z_scores) > threshold, axis=1)
    
    # Remove outliers
    cleaned_poses = np.array(poses)[~outliers]
    cleaned_values = extracted_values[~outliers]
    
    # Calculate mean and standard deviation of the cleaned data
    cleaned_mean = np.mean(cleaned_values, axis=0)
    cleaned_std = np.std(cleaned_values, axis=0)
    
    # Print results
    print("Mean of Extracted Values:\n", mean)
    print("Standard Deviation of Extracted Values:\n", std)
    print("Mean of Cleaned Extracted Values:\n", cleaned_mean)
    print("Standard Deviation of Cleaned Extracted Values:\n", cleaned_std)
    
    return cleaned_poses