111 lines
4.3 KiB
Python
111 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
import open3d as o3d
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
|
|
@dataclass
|
|
class RegistrationResult:
|
|
transformation: np.ndarray
|
|
fitness: float
|
|
rmse: float
|
|
target: o3d.geometry.PointCloud
|
|
recovered: o3d.geometry.PointCloud
|
|
|
|
|
|
class PointCloudRegistration:
|
|
"""RANSAC + ICP point-cloud registration."""
|
|
|
|
def rotation_matrix(self, roll: float = 10.0, pitch: float = 10.0, yaw: float = 10.0) -> np.ndarray:
|
|
return Rotation.from_euler("xyz", [roll, pitch, yaw], degrees=True).as_matrix()
|
|
|
|
def translation_vector(self, trans_range: float = 0.0, seed: int | None = None) -> np.ndarray:
|
|
rng = np.random.default_rng(seed)
|
|
return rng.uniform(-trans_range, trans_range, 3)
|
|
|
|
def transform_point_cloud(
|
|
self,
|
|
pointcloud: o3d.geometry.PointCloud,
|
|
rotation: np.ndarray,
|
|
translation: np.ndarray,
|
|
) -> o3d.geometry.PointCloud:
|
|
transformed = copy.deepcopy(pointcloud)
|
|
points = np.asarray(transformed.points)
|
|
transformed.points = o3d.utility.Vector3dVector((rotation @ points.T).T + translation)
|
|
return transformed
|
|
|
|
def register(
|
|
self,
|
|
source: o3d.geometry.PointCloud,
|
|
target: o3d.geometry.PointCloud,
|
|
voxel_size: float = 10.0,
|
|
) -> tuple[np.ndarray, float, float]:
|
|
source_down = source.voxel_down_sample(voxel_size)
|
|
target_down = target.voxel_down_sample(voxel_size)
|
|
|
|
normal_radius = voxel_size * 2
|
|
source_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=normal_radius, max_nn=30))
|
|
target_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=normal_radius, max_nn=30))
|
|
|
|
feature_radius = voxel_size * 5
|
|
source_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
|
|
source_down,
|
|
o3d.geometry.KDTreeSearchParamHybrid(radius=feature_radius, max_nn=100),
|
|
)
|
|
target_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
|
|
target_down,
|
|
o3d.geometry.KDTreeSearchParamHybrid(radius=feature_radius, max_nn=100),
|
|
)
|
|
|
|
distance_threshold = voxel_size * 1.5
|
|
result_ransac = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
|
|
source_down,
|
|
target_down,
|
|
source_fpfh,
|
|
target_fpfh,
|
|
mutual_filter=True,
|
|
max_correspondence_distance=distance_threshold,
|
|
estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
|
|
ransac_n=3,
|
|
checkers=[
|
|
o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
|
|
o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold),
|
|
],
|
|
criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999),
|
|
)
|
|
|
|
source.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
|
|
target.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
|
|
|
|
result_icp = o3d.pipelines.registration.registration_icp(
|
|
source,
|
|
target,
|
|
distance_threshold / 2,
|
|
result_ransac.transformation,
|
|
o3d.pipelines.registration.TransformationEstimationPointToPlane(),
|
|
o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200),
|
|
)
|
|
return result_icp.transformation, float(result_icp.fitness), float(result_icp.inlier_rmse)
|
|
|
|
def run_simulated_registration(
|
|
self,
|
|
source: o3d.geometry.PointCloud,
|
|
voxel_size: float = 10.0,
|
|
roll: float = 10.0,
|
|
pitch: float = 10.0,
|
|
yaw: float = 10.0,
|
|
trans_range: float = 0.0,
|
|
random_seed: int | None = None,
|
|
) -> RegistrationResult:
|
|
rotation = self.rotation_matrix(roll, pitch, yaw)
|
|
translation = self.translation_vector(trans_range, random_seed)
|
|
target = self.transform_point_cloud(source, rotation, translation)
|
|
transformation, fitness, rmse = self.register(source, target, voxel_size)
|
|
recovered = copy.deepcopy(target)
|
|
recovered.transform(np.linalg.inv(transformation))
|
|
return RegistrationResult(transformation, fitness, rmse, target, recovered)
|