Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import open3d as o3d | |
| import numpy as np | |
| import pandas as pd | |
| from PIL import Image | |
| from gradio.components import File, Slider, Image, Dataframe, Model3D | |
| import tempfile | |
| def point_cloud_to_image(point_cloud): | |
| # Create Open3D visualization object | |
| vis = o3d.visualization.Visualizer() | |
| vis.create_window() | |
| vis.add_geometry(point_cloud) | |
| # Render the visualization and get the image | |
| vis.poll_events() | |
| vis.update_renderer() | |
| img = vis.capture_screen_float_buffer() | |
| # Convert image to PIL Image format | |
| pil_img = Image.fromarray(np.uint8(np.asarray(img)*255)) | |
| return pil_img | |
| def create_file_mesh_from_pcd(pcd, alpha): | |
| temp_file = tempfile.NamedTemporaryFile(suffix=".obj", delete=False) | |
| file_path = temp_file.name | |
| temp_file.close() | |
| # print(f"alpha={alpha:.3f}") | |
| mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_alpha_shape(pcd, alpha) | |
| mesh.compute_vertex_normals() | |
| o3d.io.write_triangle_mesh(file_path, mesh, write_triangle_uvs=True) | |
| return file_path | |
| def plane_detection(pointcloud_path, voxel_size=0.05, distance_threshold=0.01, num_iterations=1000, alpha=2): | |
| # Load point cloud from file | |
| pcd = o3d.io.read_point_cloud(pointcloud_path.name) | |
| # Downsample the point cloud to reduce computational load | |
| pcd = pcd.voxel_down_sample(voxel_size=voxel_size) | |
| # Find plane using RANSAC algorithm | |
| plane_model, inliers = pcd.segment_plane(distance_threshold=distance_threshold, | |
| ransac_n=3, | |
| num_iterations=num_iterations) | |
| # Get inlier and outlier point cloud | |
| inlier_cloud = pcd.select_by_index(inliers) | |
| outlier_cloud = pcd.select_by_index(inliers, invert=True) | |
| # extract the coefficients of the plane | |
| a, b, c, d = plane_model | |
| plane_model = np.array([[a], [b], [c], [d]]) | |
| df = pd.DataFrame(plane_model.reshape(1, 4), columns=["a", "b", "c", "d"]) | |
| input_path = create_file_mesh_from_pcd(pcd, alpha) | |
| inlier_path = create_file_mesh_from_pcd(inlier_cloud, alpha) | |
| outlier_path = create_file_mesh_from_pcd(outlier_cloud, alpha) | |
| # Return inlier point cloud, outlier point cloud, and plane model | |
| # return point_cloud_to_image(inlier_cloud), point_cloud_to_image(outlier_cloud), df | |
| return input_path, inlier_path, outlier_path, df | |
| outputs = [ | |
| # show pcd inlier point cloud | |
| Model3D(label="Input Cloud", clear_color=[1.0, 1.0, 1.0, 1.0]), | |
| # show pcd inlier point cloud | |
| Model3D(label="Inlier Cloud", clear_color=[1.0, 1.0, 1.0, 1.0]), | |
| # show pcd outlier point cloud | |
| Model3D(label="Outlier Cloud", clear_color=[1.0, 1.0, 1.0, 1.0]), | |
| # show the centroids and counts, which is a numpy array | |
| Dataframe(label="Coefficients of the plane model", type="pandas") | |
| ] | |
| # Create Gradio interface | |
| iface = gr.Interface(plane_detection, | |
| inputs=[ | |
| File(label="Point cloud file (.ply or .pcd format)"), | |
| Slider(label="Voxel size", minimum=0.001, maximum=50, step=1, value=2), | |
| Slider(label="Distance threshold", minimum=0.001, maximum=50, step=0.01, value=5), | |
| Slider(label="Number of iterations", minimum=1, maximum=10000, step=1, value=100), | |
| Slider(label="Alpha for surface reconstruction", minimum=0.02, maximum=100, step=0.01, value=2), | |
| ], | |
| outputs=outputs, | |
| title="Plane Detection using RANSAC", | |
| description="This app takes as input a point cloud file (.ply or .pcd format), voxel size, distance threshold, and number of iterations, finds a plane using RANSAC algorithm, displays the inlier and outlier point clouds, and returns the inlier point cloud, outlier point cloud, and the plane model.", | |
| allow_flagging="never", | |
| examples=[ | |
| # ["Pointclouds/p1.ply", 2, 5, 100, 2] | |
| ["Pointclouds/cloud_bin_0.ply", 0.001, 0.03, 100, 0.05], | |
| ["Pointclouds/cloud_bin_20.ply", 0.001, 0.03, 100, 0.05], | |
| ["Pointclouds/cloud_bin_30.ply", 0.001, 0.03, 100, 0.05], | |
| ["Pointclouds/cloud_bin_40.ply", 0.001, 0.03, 100, 0.05], | |
| ["Pointclouds/cloud_bin_50.ply", 0.001, 0.03, 100, 0.05], | |
| ], | |
| ) | |
| # Launch the interface | |
| iface.launch() | |