Initial commit: workspace configuration and src

src/FAST-LIVO2/scripts/mesh.py

@@ -0,0 +1,114 @@
+import os
+import glob
+import vdbfusion
+import numpy as np
+import open3d as o3d
+from scipy.spatial import KDTree
+
+# ------------------------- Key Parameters -------------------------
+# Dataset path
+SOURCE_DIR = "/home/chunran/Downloads/"  # Replace with your dataset path
+
+# VDBVolume parameters
+VOXEL_SIZE = 0.02  # Voxel size (smaller values increase precision but require more memory)
+SDF_TRUNC = 0.1    # Truncation distance for SDF (affects surface thickness, typically a multiple of voxel size)
+MIN_WEIGHT = 0.1   # Minimum weight for mesh extraction (filters out noisy voxels)
+
+# ------------------------- Dataset Class -------------------------
+class Dataset:
+    def __init__(self, folder: str):
+        super().__init__()
+        # Get all .pcd files in the folder
+        self.scan_files = glob.glob(os.path.join(folder, "*.pcd")) 
+        # Initialize poses as identity matrices
+        self.poses = np.array([np.eye(4) for _ in range(len(self.scan_files))]) 
+
+    def __getitem__(self, idx):
+        if idx >= len(self.scan_files):
+            raise IndexError("Index out of range")
+        
+        # Compute relative pose
+        pose = np.linalg.inv(self.poses[0]) @ self.poses[idx]
+        # Read point cloud
+        points, colors = self.read_pcd(self.scan_files[idx])
+        points = np.array(points, dtype=np.float64)
+        return points, colors, pose
+
+    def __len__(self):
+        return len(self.scan_files)
+    
+    def read_pcd(self, pcd_file):
+        # Read .pcd file using Open3D
+        pcd = o3d.io.read_point_cloud(pcd_file)
+        # Extract point cloud coordinates
+        points = np.asarray(pcd.points)
+        # Extract colors (if available)
+        if pcd.has_colors():
+            colors = np.asarray(pcd.colors)  # Open3D colors are in range [0, 1]
+        else:
+            colors = np.zeros_like(points)  # If no colors, fill with zeros
+        return points, colors
+
+# ------------------------- Main Program -------------------------
+if __name__ == '__main__':
+    # Initialize VDBVolume
+    print("Initializing VDBVolume...")
+    vdb_volume = vdbfusion.VDBVolume(voxel_size=VOXEL_SIZE, sdf_trunc=SDF_TRUNC)
+
+    # Load dataset
+    print("Loading dataset from", SOURCE_DIR)
+    dataset = Dataset(SOURCE_DIR)
+
+    # Integrate all point clouds into the VDBVolume
+    print("Integrating point clouds into VDBVolume...")
+    for i in range(len(dataset)):
+        scan, colors, origin = dataset[i]
+        vdb_volume.integrate(scan, origin)
+
+    print("Point cloud integration complete!")
+
+    # Extract triangle mesh
+    print("Extracting triangle mesh...")
+    vert, tri = vdb_volume.extract_triangle_mesh(min_weight=MIN_WEIGHT)
+
+    # Create Open3D mesh object
+    print("Creating Open3D mesh object...")
+    mesh = o3d.geometry.TriangleMesh(
+        o3d.utility.Vector3dVector(vert),
+        o3d.utility.Vector3iVector(tri),
+    )
+    
+    # Save the mesh
+    print("Saving the mesh to output_mesh.ply...")
+    o3d.io.write_triangle_mesh("mesh.ply", mesh)
+    print("Mesh saved successfully.")
+
+    # ------------------------- Colorize Mesh Vertices -------------------------
+    print("Starting mesh colorization...")
+    if dataset[0][1] is not None:  # Check if color information exists
+        # Combine all point cloud points and colors
+        pcd_points = np.vstack([dataset[i][0] for i in range(len(dataset))])
+        pcd_colors = np.vstack([dataset[i][1] for i in range(len(dataset))])
+
+        # Use KDTree to find the nearest point for each vertex
+        kdtree = KDTree(pcd_points)
+        _, indices = kdtree.query(vert)  # Find the nearest point cloud point for each vertex
+        vertex_colors = pcd_colors[indices]  # Assign colors
+
+        # Set mesh vertex colors
+        mesh.vertex_colors = o3d.utility.Vector3dVector(vertex_colors)
+
+    print("Mesh colorization complete!")
+
+    # Compute vertex normals
+    print("Computing vertex normals...")
+    mesh.compute_vertex_normals()
+
+    # Save the textured mesh
+    print("Saving the textured mesh to textured_mesh.ply...")
+    o3d.io.write_triangle_mesh("textured_mesh.ply", mesh)
+    print("Textured mesh saved successfully.")
+
+    # Visualize the final colorized mesh
+    print("Visualizing the colorized mesh...")
+    o3d.visualization.draw_geometries([mesh])