Lite3DReg / examples /python /example_python.py
duanbotu123
Initial commit: add index.html
f6dd1c2
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()