Download examples/python/example_python.py from USTC3DVer/Lite3DReg: direct link, hf CLI and curl.
- Browser
- Download file 6.34 kB
-
https://huggingface.co/spaces/USTC3DVer/Lite3DReg/resolve/main/examples/python/example_python.py
- Command line
-
hf download hf://spaces/USTC3DVer/Lite3DReg/examples/python/example_python.py
-
curl -L -o example_python.py https://huggingface.co/spaces/USTC3DVer/Lite3DReg/resolve/main/examples/python/example_python.py
6.34 kB
| 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") | |
| # Rigid registration (file mode) | |
| 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")) | |
| # Load into NumPy and run in-memory registration | |
| 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")) | |
| # Non-rigid registration (file mode) | |
| 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")) | |
| # Load into NumPy and run in-memory registration | |
| 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() | |