115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
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])
|