import numpy as np import open3d as o3d import torch from tqdm import tqdm import torch.nn.functional as F def pts2normal(pts): h, w, _ = pts.shape # Compute differences in x and y directions dx = torch.cat([pts[2:, 1:-1] - pts[:-2, 1:-1]], dim=0) dy = torch.cat([pts[1:-1, 2:] - pts[1:-1, :-2]], dim=1) # Compute normal vectors using cross product normal_map = F.normalize(torch.cross(dx, dy, dim=-1), dim=-1) # Create padded normal map padded_normal_map = torch.zeros_like(pts) padded_normal_map[1:-1, 1:-1, :] = normal_map # Pad the borders padded_normal_map[0, 1:-1, :] = normal_map[0, :, :] # Top edge padded_normal_map[-1, 1:-1, :] = normal_map[-1, :, :] # Bottom edge padded_normal_map[1:-1, 0, :] = normal_map[:, 0, :] # Left edge padded_normal_map[1:-1, -1, :] = normal_map[:, -1, :] # Right edge # Pad the corners padded_normal_map[0, 0, :] = normal_map[0, 0, :] # Top-left corner padded_normal_map[0, -1, :] = normal_map[0, -1, :] # Top-right corner padded_normal_map[-1, 0, :] = normal_map[-1, 0, :] # Bottom-left corner padded_normal_map[-1, -1, :] = normal_map[-1, -1, :] # Bottom-right corner return padded_normal_map def point2mesh(pcd, depth=8, density_threshold=0.1, clean_mesh=True): print("\nPerforming Poisson surface reconstruction...") mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson( pcd, depth=depth, width=0, scale=1.1, linear_fit=False) print(f"Reconstructed mesh has {len(mesh.vertices)} vertices and {len(mesh.triangles)} triangles") # Normalize densities densities = np.asarray(densities) densities = (densities - densities.min()) / (densities.max() - densities.min()) # Remove low density vertices print("\nPruning low-density vertices...") vertices_to_remove = densities < np.quantile(densities, density_threshold) mesh.remove_vertices_by_mask(vertices_to_remove) print(f"Pruned mesh has {len(mesh.vertices)} vertices and {len(mesh.triangles)} triangles") if clean_mesh: print("\nCleaning the mesh...") mesh.remove_degenerate_triangles() mesh.remove_duplicated_triangles() mesh.remove_duplicated_vertices() mesh.remove_non_manifold_edges() print(f"Final cleaned mesh has {len(mesh.vertices)} vertices and {len(mesh.triangles)} triangles") mesh.compute_triangle_normals() return mesh def combine_and_clean_point_clouds(pcds, voxel_size): """ Combine, downsample, and clean a list of point clouds. Parameters: pcds (list): List of open3d.geometry.PointCloud objects to be processed. voxel_size (float): The size of the voxel for downsampling. Returns: o3d.geometry.PointCloud: The cleaned and combined point cloud. """ print("\nCombining point clouds...") pcd_combined = o3d.geometry.PointCloud() for p3d in pcds: pcd_combined += p3d print("\nDownsampling the combined point cloud...") pcd_combined = pcd_combined.voxel_down_sample(voxel_size) print(f"Downsampled from {len(pcd_combined.points)} to {len(pcd_combined.points)} points") print("\nCleaning the combined point cloud...") cl, ind = pcd_combined.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0) pcd_cleaned = pcd_combined.select_by_index(ind) print(f"Cleaned point cloud contains {len(pcd_cleaned.points)} points.") print(f"Removed {len(pcd_combined.points) - len(pcd_cleaned.points)} outlier points.") return pcd_cleaned def improved_multiway_registration(pcds, descriptors=None, voxel_size=0.05, max_correspondence_distance_coarse=None, max_correspondence_distance_fine=None, overlap=5, quadratic_overlap=False, use_colored_icp=False): if max_correspondence_distance_coarse is None: max_correspondence_distance_coarse = voxel_size * 1.5 if max_correspondence_distance_fine is None: max_correspondence_distance_fine = voxel_size * 0.15 def pairwise_registration(source, target, use_colored_icp, max_correspondence_distance_coarse, max_correspondence_distance_fine): current_transformation = np.identity(4) try: if use_colored_icp: icp_fine = o3d.pipelines.registration.registration_colored_icp( source, target, max_correspondence_distance_fine, current_transformation, o3d.pipelines.registration.TransformationEstimationForColoredICP(), o3d.pipelines.registration.ICPConvergenceCriteria(relative_fitness=1e-6, relative_rmse=1e-6, max_iteration=100)) else: icp_fine = o3d.pipelines.registration.registration_icp( source, target, max_correspondence_distance_fine, current_transformation, o3d.pipelines.registration.TransformationEstimationPointToPlane()) fitness = icp_fine.fitness FITNESS_THRESHOLD = 0.01 if fitness >= FITNESS_THRESHOLD: current_transformation = icp_fine.transformation information_icp = o3d.pipelines.registration.get_information_matrix_from_point_clouds( source, target, max_correspondence_distance_fine, current_transformation) return current_transformation, information_icp, True else: print(f"Registration failed. Fitness {fitness} is below threshold {FITNESS_THRESHOLD}") return None, None, False except RuntimeError as e: print(f" ICP registration failed: {str(e)}") return None, None, False def detect_loop_closure(descriptors, min_interval=3, similarity_threshold=0.9): n_pcds = len(descriptors) loop_edges = [] for i in range(n_pcds): for j in range(i + min_interval, n_pcds): similarity = torch.dot(descriptors[i], descriptors[j]) if similarity > similarity_threshold: loop_edges.append((i, j)) return loop_edges def generate_pairs(n_pcds, overlap, quadratic_overlap): pairs = [] for i in range(n_pcds - 1): for j in range(i + 1, min(i + overlap + 1, n_pcds)): pairs.append((i, j)) if quadratic_overlap: q = 2**(j-i) if q > overlap and i + q < n_pcds: pairs.append((i, i + q)) return pairs def full_registration(pcds_down, pairs, loop_edges): pose_graph = o3d.pipelines.registration.PoseGraph() n_pcds = len(pcds_down) for i in range(n_pcds): pose_graph.nodes.append(o3d.pipelines.registration.PoseGraphNode(np.identity(4))) print("\nPerforming pairwise registration:") for source_id, target_id in tqdm(pairs): transformation_icp, information_icp, success = pairwise_registration( pcds_down[source_id], pcds_down[target_id], use_colored_icp, max_correspondence_distance_coarse, max_correspondence_distance_fine) if success: uncertain = abs(target_id - source_id) == 1 pose_graph.edges.append( o3d.pipelines.registration.PoseGraphEdge(source_id, target_id, transformation_icp, information_icp, uncertain=uncertain)) else: print(f" Skipping edge between {source_id} and {target_id} due to ICP failure") # Add loop closure edges print("\nAdding loop closure edges:") for source_id, target_id in tqdm(loop_edges): transformation_icp, information_icp, success = pairwise_registration( pcds_down[source_id], pcds_down[target_id], use_colored_icp, max_correspondence_distance_coarse, max_correspondence_distance_fine) if success: pose_graph.edges.append( o3d.pipelines.registration.PoseGraphEdge(source_id, target_id, transformation_icp, information_icp, uncertain=True)) else: print(f" Skipping loop closure edge between {source_id} and {target_id} due to ICP failure") return pose_graph print("\n--- Improved Multiway Registration Process ---") print(f"Number of point clouds: {len(pcds)}") print(f"Voxel size: {voxel_size}") print(f"Max correspondence distance (coarse): {max_correspondence_distance_coarse}") print(f"Max correspondence distance (fine): {max_correspondence_distance_fine}") print(f"Overlap: {overlap}") print(f"Quadratic overlap: {quadratic_overlap}") print("\nPreprocessing point clouds...") pcds_down = pcds print(f"Preprocessing complete. {len(pcds_down)} point clouds processed.") print("\nGenerating initial graph pairs...") pairs = generate_pairs(len(pcds), overlap, quadratic_overlap) print(f"Generated {len(pairs)} pairs for initial graph.") if descriptors is None: print("\nNo descriptors provided. Skipping loop closure detection.") loop_edges = [] else: print(descriptors[0].shape) print("\nDetecting loop closures...") loop_edges = detect_loop_closure(descriptors) print(f"Detected {len(loop_edges)} loop closures.") print("\nPerforming full registration...") pose_graph = full_registration(pcds_down, pairs, loop_edges) print("\nOptimizing PoseGraph...") option = o3d.pipelines.registration.GlobalOptimizationOption( max_correspondence_distance=max_correspondence_distance_fine, edge_prune_threshold=0.25, reference_node=0) o3d.pipelines.registration.global_optimization( pose_graph, o3d.pipelines.registration.GlobalOptimizationLevenbergMarquardt(), o3d.pipelines.registration.GlobalOptimizationConvergenceCriteria(), option) # Count edges for each node edge_count = {i: 0 for i in range(len(pcds))} for edge in pose_graph.edges: edge_count[edge.source_node_id] += 1 edge_count[edge.target_node_id] += 1 # Filter nodes with more than 3 edges valid_nodes = [count > 3 for count in edge_count.values()] print("\nTransforming and combining point clouds...") pcd_combined = o3d.geometry.PointCloud() for point_id, is_valid in enumerate(valid_nodes): if is_valid: pcds[point_id].transform(pose_graph.nodes[point_id].pose) pcd_combined += pcds[point_id] else: print(f"Skipping point cloud {point_id} as it has {edge_count[point_id]} edges (<=3)") print("\nDownsampling the combined point cloud...") # pcd_combined.orient_normals_consistent_tangent_plane(k=30) pcd_combined = pcd_combined.voxel_down_sample(voxel_size * 0.1) print(f"Downsampled from {len(pcd_combined.points)} to {len(pcd_combined.points)} points") print("\nCleaning the combined point cloud...") cl, ind = pcd_combined.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0) pcd_cleaned = pcd_combined.select_by_index(ind) print(f"Cleaned point cloud contains {len(pcd_cleaned.points)} points.") print(f"Removed {len(pcd_combined.points) - len(pcd_cleaned.points)} outlier points.") print("\nMultiway registration complete.") print(f"Included {len(valid_nodes)} out of {len(pcds)} point clouds (with >3 edges).") return pcd_cleaned, pose_graph, valid_nodes