forked from koide3/small_gicp
-
Notifications
You must be signed in to change notification settings - Fork 0
/
basic_registration.py
executable file
·158 lines (122 loc) · 7.02 KB
/
basic_registration.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
#!/usr/bin/python3
# SPDX-FileCopyrightText: Copyright 2024 Kenji Koide
# SPDX-License-Identifier: MIT
import numpy
from scipy.spatial.transform import Rotation
import small_gicp
# Basic registation example with numpy arrays
def example_numpy1(target_raw_numpy : numpy.ndarray, source_raw_numpy : numpy.ndarray):
# Example A : Perform registration with numpy arrays
# Arguments
# - target_points : Nx4 or Nx3 numpy array of the target point cloud
# - source_points : Nx4 or Nx3 numpy array of the source point cloud
# Optional arguments
# - init_T_target_source : Initial guess of the transformation matrix (4x4 numpy array)
# - registration_type : Registration type ("ICP", "PLANE_ICP", "GICP", "VGICP")
# - voxel_resolution : Voxel resolution for VGICP
# - downsampling_resolution : Downsampling resolution
# - max_correspondence_distance : Maximum correspondence distance
# - num_threads : Number of threads
result = small_gicp.align(target_raw_numpy, source_raw_numpy, downsampling_resolution=0.25)
return result.T_target_source
# Example to perform preprocessing and registration separately
def example_numpy2(target_raw_numpy : numpy.ndarray, source_raw_numpy : numpy.ndarray):
# Example B : Perform preprocessing and registration separately
# Preprocess point clouds
# Arguments
# - points : Nx4 or Nx3 numpy array of the target point cloud
# Optional arguments
# - downsampling_resolution : Downsampling resolution
# - num_neighbors : Number of neighbors for normal and covariance estimation
# - num_threads : Number of threads
target, target_tree = small_gicp.preprocess_points(target_raw_numpy, downsampling_resolution=0.25)
source, source_tree = small_gicp.preprocess_points(source_raw_numpy, downsampling_resolution=0.25)
# Align point clouds
# Arguments
# - target : Target point cloud (small_gicp.PointCloud)
# - source : Source point cloud (small_gicp.PointCloud)
# - target_tree : KD-tree of the target point cloud
# Optional arguments
# - init_T_target_source : Initial guess of the transformation matrix (4x4 numpy array)
# - max_correspondence_distance : Maximum correspondence distance
# - num_threads : Number of threads
result = small_gicp.align(target, source, target_tree)
return result.T_target_source
# Basic registation example with small_gicp.PointCloud
def example_small1(target_raw_numpy : numpy.ndarray, source_raw_numpy : numpy.ndarray):
# Convert numpy arrays (Nx3 or Nx4) to small_gicp.PointCloud
target_raw = small_gicp.PointCloud(target_raw_numpy)
source_raw = small_gicp.PointCloud(source_raw_numpy)
# Preprocess point clouds
target, target_tree = small_gicp.preprocess_points(target_raw, downsampling_resolution=0.25)
source, source_tree = small_gicp.preprocess_points(source_raw, downsampling_resolution=0.25)
result = small_gicp.align(target, source, target_tree)
return result.T_target_source
# Example to perform each preprocessing and registration separately
def example_small2(target_raw_numpy : numpy.ndarray, source_raw_numpy : numpy.ndarray):
# Convert numpy arrays (Nx3 or Nx4) to small_gicp.PointCloud
target_raw = small_gicp.PointCloud(target_raw_numpy)
source_raw = small_gicp.PointCloud(source_raw_numpy)
# Downsampling
target = small_gicp.voxelgrid_sampling(target_raw, 0.25)
source = small_gicp.voxelgrid_sampling(source_raw, 0.25)
# KdTree construction
target_tree = small_gicp.KdTree(target)
source_tree = small_gicp.KdTree(source)
# Estimate covariances
small_gicp.estimate_covariances(target, target_tree)
small_gicp.estimate_covariances(source, source_tree)
# Align point clouds
result = small_gicp.align(target, source, target_tree)
return result.T_target_source
### Following functions are for testing ###
# Verity the estimated transformation matrix (for testing)
def verify_result(T_target_source, gt_T_target_source):
error = numpy.linalg.inv(T_target_source) @ gt_T_target_source
error_trans = numpy.linalg.norm(error[:3, 3])
error_rot = Rotation.from_matrix(error[:3, :3]).magnitude()
assert error_trans < 0.05
assert error_rot < 0.05
import pytest
# Load the point clouds and the ground truth transformation matrix
@pytest.fixture(scope='module', autouse=True)
def load_points():
gt_T_target_source = numpy.loadtxt('data/T_target_source.txt') # Load the ground truth transformation matrix
print('--- gt_T_target_source ---')
print(gt_T_target_source)
target_raw = small_gicp.read_ply(('data/target.ply')) # Read the target point cloud (small_gicp.PointCloud)
source_raw = small_gicp.read_ply(('data/source.ply')) # Read the source point cloud (small_gicp.PointCloud)
target_raw_numpy = target_raw.points() # Nx4 numpy array of the target point cloud
source_raw_numpy = source_raw.points() # Nx4 numpy array of the source point cloud
yield (gt_T_target_source, target_raw_numpy, source_raw_numpy)
# Check if the point clouds are loaded correctly
def test_load_points(load_points):
gt_T_target_source, target_raw_numpy, source_raw_numpy = load_points
assert gt_T_target_source.shape[0] == 4 and gt_T_target_source.shape[1] == 4
assert len(target_raw_numpy) > 0
assert len(source_raw_numpy) > 0
def test_example_numpy1(load_points):
gt_T_target_source, target_raw_numpy, source_raw_numpy = load_points
T_target_source = example_numpy1(target_raw_numpy, source_raw_numpy)
verify_result(T_target_source, gt_T_target_source)
def test_example_numpy2(load_points):
gt_T_target_source, target_raw_numpy, source_raw_numpy = load_points
T_target_source = example_numpy2(target_raw_numpy, source_raw_numpy)
verify_result(T_target_source, gt_T_target_source)
def test_example_small1(load_points):
gt_T_target_source, target_raw_numpy, source_raw_numpy = load_points
T_target_source = example_small1(target_raw_numpy, source_raw_numpy)
verify_result(T_target_source, gt_T_target_source)
def test_example_small2(load_points):
gt_T_target_source, target_raw_numpy, source_raw_numpy = load_points
T_target_source = example_small2(target_raw_numpy, source_raw_numpy)
verify_result(T_target_source, gt_T_target_source)
if __name__ == "__main__":
target_raw = small_gicp.read_ply(('data/target.ply')) # Read the target point cloud (small_gicp.PointCloud)
source_raw = small_gicp.read_ply(('data/source.ply')) # Read the source point cloud (small_gicp.PointCloud)
target_raw_numpy = target_raw.points() # Nx4 numpy array of the target point cloud
source_raw_numpy = source_raw.points() # Nx4 numpy array of the source point cloud
T_target_source = example_numpy1(target_raw_numpy, source_raw_numpy)
T_target_source = example_numpy2(target_raw_numpy, source_raw_numpy)
T_target_source = example_small1(target_raw_numpy, source_raw_numpy)
T_target_source = example_small2(target_raw_numpy, source_raw_numpy)