-
Notifications
You must be signed in to change notification settings - Fork 0
/
fast_global_registration.py
106 lines (87 loc) · 4.43 KB
/
fast_global_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
import open3d as o3d
import numpy as np
import copy
import time
SOURCE_PCD = "bunny/data/bun000.ply"
TARGET_PCD = "bunny/data/bun045.ply"
# Visualization
def draw_registration_result(source, target, transformation):
source_temp = copy.deepcopy(source)
target_temp = copy.deepcopy(target)
source_temp.paint_uniform_color([1, 0.706, 0])
target_temp.paint_uniform_color([0, 0.651, 0.929])
source_temp.transform(transformation)
o3d.visualization.draw_geometries([source_temp, target_temp])
# Extract geometric feature
def preprocess_point_cloud(pcd, voxel_size):
print(":: Downsample with a voxel size %.3f." % voxel_size)
pcd_down = pcd.voxel_down_sample(voxel_size)
radius_normal = voxel_size * 2
print(":: Estimate normal with search radius %.3f." % radius_normal)
pcd_down.estimate_normals(
o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30))
# pcd.estimate_normals(
# o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30))
radius_feature = voxel_size * 5
print(":: Compute FPFH feature with search radius %.3f." % radius_feature)
pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
pcd_down,
o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100))
# pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
# pcd,
# o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100))
return pcd_down, pcd_fpfh
# return pcd, pcd_fpfh
# Input
def prepare_dataset(voxel_size):
print(":: Load two point clouds and disturb initial pose.")
source = o3d.io.read_point_cloud(SOURCE_PCD)
target = o3d.io.read_point_cloud(TARGET_PCD)
print(f":: Source points amount: {len(source.points)}")
print(f":: Target points amount: {len(target.points)}")
trans_init = np.asarray([[0.0, 0.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]])
source.transform(trans_init)
draw_registration_result(source, target, np.identity(4))
source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)
target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)
print(f":: Downsample source points amount: {len(source_down.points)}")
print(f":: Downsample target points amount: {len(target_down.points)}")
return source, target, source_down, target_down, source_fpfh, target_fpfh
# Fast Global Registration
def execute_fast_global_registration(source_down, target_down, source_fpfh, target_fpfh, voxel_size):
distance_threshold = voxel_size * 0.5
print(":: Apply fast global registration with distance threshold %.3f" \
% distance_threshold)
result = o3d.pipelines.registration.registration_fgr_based_on_feature_matching(
source_down, target_down, source_fpfh, target_fpfh,
o3d.pipelines.registration.FastGlobalRegistrationOption(
maximum_correspondence_distance=distance_threshold))
return result
# Local refinement
def refine_registration(source, target, source_fpfh, target_fpfh, voxel_size):
distance_threshold = voxel_size * 0.4
print(":: Point-to-plane ICP registration is applied on original point")
print(" clouds to refine the alignment. This time we use a strict")
print(" distance threshold %.3f." % distance_threshold)
result = o3d.pipelines.registration.registration_icp(
source, target, distance_threshold, result_fast.transformation,
o3d.pipelines.registration.TransformationEstimationPointToPlane())
return result
'''main code'''
# Input
voxel_size = 0.005 # 0.05 means 5cm for the dataset
source, target, source_down, target_down, source_fpfh, target_fpfh = prepare_dataset(voxel_size)
# Fast Global Registration
start = time.time()
result_fast = execute_fast_global_registration(source_down, target_down,
source_fpfh, target_fpfh,
voxel_size)
print("Fast global registration took %.3f sec.\n" % (time.time() - start))
print(result_fast)
# draw_registration_result(source_down, target_down, result_fast.transformation)
# Local refinement
result_icp = refine_registration(source_down, target_down, source_fpfh, target_fpfh, voxel_size)
print("Fast global registration plus Local refinement took %.3f sec.\n" % (time.time() - start))
print(result_icp)
draw_registration_result(source, target, result_icp.transformation)