| import os |
| import numpy as np |
| import trimesh |
| from plyfile import PlyData |
| import pyregister |
|
|
|
|
| def read_vertices_normals(file_path): |
| """ |
| Read vertices and normals from a point cloud or mesh file. |
| Supports formats such as .ply, .obj, .xyz, etc. |
| """ |
| ext = os.path.splitext(file_path)[1].lower() |
|
|
| if ext == ".ply": |
| plydata = PlyData.read(file_path) |
| vertex_data = plydata['vertex'] |
| vertices = np.vstack([vertex_data['x'], vertex_data['y'], vertex_data['z']]).astype(np.float64) |
|
|
| normals = None |
| if {'nx', 'ny', 'nz'}.issubset(vertex_data.data.dtype.names): |
| normals = np.vstack([vertex_data['nx'], vertex_data['ny'], vertex_data['nz']]).astype(np.float64) |
|
|
| else: |
| mesh = trimesh.load(file_path, process=False) |
| vertices = mesh.vertices.T.astype(np.float64) |
| normals = getattr(mesh, "vertex_normals", None) |
| if normals is not None and len(normals) > 0: |
| normals = normals.T.astype(np.float64) |
| else: |
| normals = np.empty((0, 0)) |
|
|
| return vertices, normals |
|
|
|
|
| def run_rigid_file(file_target, file_source, output_path): |
| """ |
| Perform rigid registration (FRICP method) by directly reading files. |
| |
| Parameters |
| ---------- |
| file_target : str |
| Path to the target point cloud or mesh file. |
| file_source : str |
| Path to the source point cloud or mesh file. |
| output_path : str |
| Output file path for the registered result. |
| """ |
| print(f"\n[RIGID-FILE] Registering {file_source} β {file_target}") |
| reg = pyregister.RigidFricpRegistration() |
| reg.Paras_init(useinit=False, maxiter=100, stop=1e-5) |
| reg.Reg(file_target, file_source, output_path) |
| print(f"[RIGID-FILE] Completed β {output_path}") |
| print("Deformed points shape:", reg.deformed_points_3X_.shape) |
|
|
|
|
| def run_rigid_numpy(target_pts, source_pts, target_n, source_n, output_path): |
| """ |
| Perform rigid registration (FRICP method) using NumPy matrices. |
| |
| Parameters |
| ---------- |
| target_pts : np.ndarray, shape (3, N) |
| Target point cloud coordinates. |
| source_pts : np.ndarray, shape (3, N) |
| Source point cloud coordinates. |
| target_n : np.ndarray, shape (3, N) or (0, 0) |
| Target normals. |
| source_n : np.ndarray, shape (3, N) or (0, 0) |
| Source normals. |
| output_path : str |
| Output file path for the registered result. |
| """ |
| print(f"\n[RIGID-NUMPY] Registering from NumPy matrices") |
| reg = pyregister.RigidFricpRegistration() |
| reg.Paras_init() |
| reg.Read_data(target_pts, source_pts, target_n, source_n) |
| reg.Register() |
| reg.Output_data(output_path, "FRICP") |
| print(f"[RIGID-NUMPY] Completed β {output_path}") |
| print("Deformed points shape:", reg.deformed_points_3X_.shape) |
|
|
|
|
| def run_nonrigid_file(file_target, file_source, output_path): |
| """ |
| Perform non-rigid registration (SPARE method) by directly reading files. |
| |
| Parameters |
| ---------- |
| file_target : str |
| Path to the target point cloud or mesh file. |
| file_source : str |
| Path to the source point cloud or mesh file. |
| output_path : str |
| Output file path for the registered result. |
| """ |
| print(f"\n[NONRIGID-FILE] Registering {file_source} β {file_target}") |
| reg = pyregister.NonrigidSpareRegistration() |
| reg.Paras_init(iters = 30 ,stopcoarse=1e-3,stopfine=1e-4, uselandmark = False); |
| reg.Reg(file_target, file_source, output_path) |
| print(f"[NONRIGID-FILE] Completed β {output_path}") |
| print("Deformed points shape:", reg.deformed_points_3X_.shape) |
|
|
|
|
|
|
| def run_nonrigid_numpy(target_pts, source_pts, target_n, source_n, output_path): |
| """ |
| Perform non-rigid registration (SPARE method) using NumPy matrices. |
| |
| Parameters |
| ---------- |
| target_pts : np.ndarray, shape (3, N) |
| Target point cloud coordinates. |
| source_pts : np.ndarray, shape (3, N) |
| Source point cloud coordinates. |
| target_n : np.ndarray, shape (3, N) or (0, 0) |
| Target normals. |
| source_n : np.ndarray, shape (3, N) or (0, 0) |
| Source normals. |
| output_path : str |
| Output file path for the registered result. |
| """ |
| print(f"\n[NONRIGID-NUMPY] Registering from NumPy matrices") |
| reg = pyregister.NonrigidSpareRegistration() |
| reg.Paras_init() |
| reg.Read_data(target_pts, source_pts, target_n, source_n) |
| reg.Register() |
| reg.Output_data(output_path, "SPARE") |
| print(f"[NONRIGID-NUMPY] Completed β {output_path}") |
| print("Deformed points shape:", reg.deformed_points_3X_.shape) |
|
|
|
|
|
|
| def main(): |
| folder = "data" |
| output_folder = os.path.join(folder, "results_register") |
| os.makedirs(output_folder, exist_ok=True) |
| target_ply = os.path.join(folder, "target.ply") |
|
|
| source_ply = os.path.join(folder, "source.ply") |
|
|
| target_obj = os.path.join(folder, "target.obj") |
|
|
| source_obj = os.path.join(folder, "source.obj") |
|
|
| |
| if os.path.exists(target_ply) and os.path.exists(source_ply): |
| run_rigid_file(target_ply, source_ply, os.path.join(output_folder, "rigid_file_result")) |
|
|
| |
| target_pts, target_n = read_vertices_normals(target_ply) |
| source_pts, source_n = read_vertices_normals(source_ply) |
| target_n = target_n if target_n is not None else np.empty((0, 0)) |
| source_n = source_n if source_n is not None else np.empty((0, 0)) |
| run_rigid_numpy(target_pts, source_pts, target_n, source_n, |
| os.path.join(output_folder, "rigid_numpy_result")) |
|
|
| |
| if os.path.exists( target_obj ) and os.path.exists(source_obj): |
| run_nonrigid_file( target_obj , source_obj, os.path.join(output_folder, "nonrigid_file_result")) |
|
|
| |
| target_pts, target_n = read_vertices_normals( target_obj ) |
| source_pts, source_n = read_vertices_normals(source_obj) |
| target_n = target_n if target_n is not None else np.empty((0, 0)) |
| source_n = source_n if source_n is not None else np.empty((0, 0)) |
| run_nonrigid_numpy(target_pts, source_pts, target_n, source_n, |
| os.path.join(output_folder, "nonrigid_numpy_result")) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|