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]:  # sup and inf faces
                points.append([x, y, z])

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

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

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

    return point_cloud


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

    points = []

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

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

    # Walls
    for z in np.linspace(ground_z, ceiling_z, resolution):
        for y in np.linspace(-length / 2, length / 2, resolution):
            # Left
            x = shelf1_x + np.random.normal(0, noise_level)
            points.append([x, y, z])
            # Right
            x = shelf2_x + np.random.normal(0, noise_level)
            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 = generate_corridor_cloud()
    source_cloud = create_cube_cloud()
    source_cloud.paint_uniform_color([1, 0, 0])  # Red

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

    print("Real transformation = ", transformation)

    source_copy = create_cube_cloud()
    #source_copy = generate_corridor_cloud()

    # Generate target cloud
    target_cloud = source_copy.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:
        source_copy = create_cube_cloud()
        #source_copy = generate_corridor_cloud()
        aligned_cloud = source_copy.transform(result.transformation)
        aligned_cloud.paint_uniform_color([0, 0, 1])  # Blue

        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.")
~                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 
~