Untitled

import open3d as o3d
import numpy as np

def create_cube_cloud():
    # Create pointcloud filling a cube
    points = []
    size = 1
    resolution = 10

    for x in np.linspace(-size/2, size/2, resolution):
        for y in np.linspace(-size/2, size/2, resolution):
            for z in [-size/2, size/2]:
                points.append([x, y, z])  # sup and inf faces
                points.append([x, z, y])  # front and back
                points.append([z, y, x])  # left and right

    point_cloud = o3d.geometry.PointCloud()
    point_cloud.points = o3d.utility.Vector3dVector(np.array(points))

    return point_cloud


def generate_corridor_cloud(length=6, width=2, height=3, resolution=10, noise_level=0):
    shelf1_x = -width / 2  # Left wall
    shelf2_x = width / 2   # Right wall
    ground_z = 0          # Ground
    ceiling_z = height    # Ceilling

    points = []

    for x in np.linspace(shelf1_x, shelf2_x, resolution * width):
        for y in np.linspace(-length / 2, length / 2, resolution * length):
            z = ground_z + np.random.normal(0, noise_level)  # Ground
            points.append([x, y, z])
            z = ceiling_z + np.random.normal(0, noise_level)  # Ceiling
            points.append([x, y, z])

    # Walls
    for z in np.linspace(ground_z, ceiling_z, resolution * height):
        for y in np.linspace(-length / 2, length / 2, resolution * length):
            x = shelf1_x + np.random.normal(0, noise_level)  # Left
            points.append([x, y, z])
            x = shelf2_x + np.random.normal(0, noise_level)  # Right
            points.append([x, y, z])

    cloud = o3d.geometry.PointCloud()
    cloud.points = o3d.utility.Vector3dVector(np.array(points))

    return cloud

def estimate_motion_with_ransac(source_cloud, target_cloud, voxel_size=0.05, max_correspondence_distance=0.2):
    print("Downsampling the point clouds")
    source_down = source_cloud.voxel_down_sample(voxel_size)
    target_down = target_cloud.voxel_down_sample(voxel_size)

    print("Estimating normals")
    source_down.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))
    target_down.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))

    print("Computing FPFH features")
    source_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        source_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5, max_nn=100)
    )
    target_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        target_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5, max_nn=100)
    )

    print("Applying RANSAC")
    try:
        result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
            source_down,
            target_down,
            source_fpfh,
            target_fpfh,
            mutual_filter=True,
            max_correspondence_distance=max_correspondence_distance,
            estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(with_scaling=False),
            ransac_n=4,
            checkers=[
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(max_correspondence_distance)
            ],
            criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(
                max_iteration=1000000,
                confidence=0.9999
            )
        )
    except RuntimeError as e:
        print(f"RANSAC failed: {e}")
        return None

    print("RANSAC result:")
    print(result)
    print("Estimated transformation:\n", result.transformation)
    return result

if __name__ == "__main__":
    print("Simulating point clouds...")

    #source_cloud = create_cube_cloud()
    source_cloud = generate_corridor_cloud()
    source_cloud.paint_uniform_color([1, 0, 0])  # Red

    # Define tranformation between the 2 frames
    angle = np.pi / 9
    transformation = np.array([
        [np.cos(angle), -np.sin(angle), 0, 0],
        [np.sin(angle), np.cos(angle), 0, 0.2],  # +transl
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ])

    print("Real transformation", transformation, sep='\n')

    # Generate target cloud
    target_cloud = o3d.geometry.PointCloud(source_cloud).transform(transformation)
    #target_points = np.asarray(target_cloud.points)
    #target_cloud.points = o3d.utility.Vector3dVector(target_points)
    target_cloud.paint_uniform_color([0, 1, 0])  # Green

    # Apply RANSAC to estimate movement
    result = estimate_motion_with_ransac(source_cloud, target_cloud)

    # Display
    if result is not None:
        aligned_cloud = o3d.geometry.PointCloud(source_cloud).transform(result.transformation)
        aligned_cloud.paint_uniform_color([0, 0, 1])  # Blue

        #o3d.visualization.draw_geometries([source_cloud, target_cloud, aligned_cloud])
        print("Visualizing source and target")
        o3d.visualization.draw_geometries([source_cloud, target_cloud])
        print("Visualizing source and aligned clouds...")
        o3d.visualization.draw_geometries([source_cloud, aligned_cloud])
        print("Visualizing target and aligned clouds...")
        o3d.visualization.draw_geometries([target_cloud, aligned_cloud])

    else:
        print("RANSAC did not converge or failed.")