-
Notifications
You must be signed in to change notification settings - Fork 5
/
global_registration.py
374 lines (323 loc) · 16.5 KB
/
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
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
import open3d as o3d
import time
import numpy as np
from sklearn.cluster import KMeans
import hdbscan
from sklearn.neighbors import BallTree
from scipy.spatial.transform import Rotation as R
import pandas as pd
import extract
from datasets.lidar_dataset import LidarSynthetic
# define arguments needed to extract feautures from images
class Args:
def __init__(self, model=None):
self.model = model
self.top_k = 300
self.scale_f = 2 ** 0.25
self.min_size = 1024
self.max_size = 2048
self.min_scale = 1
self.max_scale = 1 # only orignal resolution to extract from
self.reliability_thr = 0.7
self.repeatability_thr = 0.7
self.gpu = [0]
# draw two point clouds transfromed into the same frame with transformation
def draw_registration_result(source, target, transformation):
import copy
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])
# estimate transformation based on key points and descriptors of two point clouds with RANSAC
def execute_global_registration(source_down, target_down, reference_desc, target_desc, distance_threshold):
result = o3d.registration.registration_ransac_based_on_feature_matching(
source_down, target_down, reference_desc, target_desc,
distance_threshold,
o3d.registration.TransformationEstimationPointToPoint(False), 4,
[o3d.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold),
o3d.registration.CorrespondenceCheckerBasedOnEdgeLength(0.7)],
o3d.registration.RANSACConvergenceCriteria(800000, 1000))
return result
# transformation estimate refined with ICP
def pairwise_registration(pc_source, pc_target, desc_source, desc_target, kpts_source, kpts_target, max_correspondence_distance):
ref = o3d.registration.Feature()
ref.data = desc_source.T
test = o3d.registration.Feature()
test.data = desc_target.T
ref_key = o3d.geometry.PointCloud()
ref_key.points = o3d.utility.Vector3dVector(kpts_source)
test_key = o3d.geometry.PointCloud()
test_key.points = o3d.utility.Vector3dVector(kpts_target)
result_ransac = execute_global_registration(ref_key, test_key, ref, test, 0.5)
tf = result_ransac.transformation
fitness = result_ransac.fitness
# run icp on point clouds to refine rigid transformation
icp = o3d.registration.registration_icp(
pc_source, pc_target, max_correspondence_distance, tf, o3d.registration.TransformationEstimationPointToPoint())
transformation_icp = icp.transformation
information_icp = o3d.registration.get_information_matrix_from_point_clouds(
pc_source, pc_target, max_correspondence_distance, transformation_icp)
return transformation_icp, information_icp, fitness
def global_registration(kpts_source, kpts_target, desc_source, desc_target, pc_source, pc_target, ransac_dist_threshold, icp_dist_threshold):
result_ransac = execute_global_registration(kpts_source, kpts_target, desc_source, desc_target, ransac_dist_threshold)
tf = result_ransac.transformation
fitness = result_ransac.fitness
reg_p2p = o3d.registration.registration_icp(
pc_source, pc_target, icp_dist_threshold, tf,
o3d.registration.TransformationEstimationPointToPoint(),
o3d.registration.ICPConvergenceCriteria(max_iteration=100))
tf = reg_p2p.transformation
information_icp = o3d.registration.get_information_matrix_from_point_clouds(
pc_source, pc_target, icp_dist_threshold, tf)
return tf, information_icp, fitness
# remove invalid key points from extracted xys, scores and desc
def remove_invalid_keypoints(xys, scores, desc, mask):
valid_mask = mask[xys[:, 1].astype(int), xys[:, 0].astype(int)]
xys = xys[valid_mask]
scores = scores[valid_mask]
desc = desc[valid_mask]
return xys, scores, desc
# cluster descriptor list in n_words clusters and calculate the normalized BoVW histogram for each scan
def build_bag_of_words_hist(desc_list, n_words=200):
n_candidates = len(desc_list)
desc = np.concatenate(desc_list, axis=0)
kmeans = KMeans(n_clusters=n_words, max_iter=10000)
labels = kmeans.fit_predict(desc)
# build histogram
hist = np.zeros((n_candidates, n_words), dtype=int)
# print(hist.shape)
candidate_it = 0
for idx in range(n_candidates):
n_landmarks = desc_list[idx].shape[0]
candidate_labels = labels[candidate_it:candidate_it+n_landmarks]
for candidate_label in candidate_labels:
hist[idx, candidate_label] = hist[idx, candidate_label] + 1
candidate_it = candidate_it + n_landmarks
#
hist = np.array(hist)
hist = hist/np.linalg.norm(hist, axis=1, keepdims=True)
return hist
# build search tree with all histograms and query a nearest neighbour search for every histogram(scan).
# returns ids of nearest neighbours for every scan with min frame dist and max histogram distance
def get_bag_of_words_nearest_neighbours(hist, min_frame_dist, max_hist_dist):
ball_tree = BallTree(hist)
ids = ball_tree.query_radius(hist, max_hist_dist)
n_candidates = len(ids)
for i in range(n_candidates):
dist_thresh_neighbour_mask = np.logical_or(i - min_frame_dist > ids[i], ids[i] > i + min_frame_dist)
ids[i] = ids[i][dist_thresh_neighbour_mask]
return ids
# run full pose graph estimation algorithm
def full_registration(db, net, net_args, args):
# initialize
icp_max_corr_dist = args["icp_max_corr_dist"]
ransac_max_dist = args["ransac_max_dist"]
min_frame_dist = args["min_frame_dist"]
max_hist_dist = args["max_hist_dist"]
n_words = args["n_words"]
pose_graph = o3d.registration.PoseGraph()
odometry = np.identity(4)
pose_graph.nodes.append(o3d.registration.PoseGraphNode(odometry))
n_pcds = len(db)
pcd_list = []
kpts_list = []
desc_list = []
# extract features from first image
img = db.get_image(0)
mask = db.get_valid_range_mask(0)
xys, scores, desc = extract.extract_keypoints(img, net_args, net)
xys, scores, desc = remove_invalid_keypoints(xys, scores, desc, mask)
xyz = db.get_xyz(0)
kpts = xyz[xys[:, 1].astype(int), xys[:, 0].astype(int)]
desc_list.append(desc)
kpts_list.append(kpts)
# extract features for every scan image and add transfromation to pose graph and key points anddscriptors to list
print("Calculating prior odometry with ICP and Collecting landmarks with r2d2 net")
for source_id in range(n_pcds-1):
progress = source_id / (n_pcds-1)
print("\rProgress {:2.1%}".format(progress), end="")
target_id = source_id + 1
xyzA = db.get_xyz(source_id)
maskA = db.get_valid_range_mask(source_id).reshape(-1)
pc_source = o3d.geometry.PointCloud()
pc_source.points = o3d.utility.Vector3dVector(xyzA.reshape(-1, 3)[maskA, :])
xyzB = db.get_xyz(target_id)
maskB = db.get_valid_range_mask(target_id).reshape(-1)
pc_target = o3d.geometry.PointCloud()
pc_target.points = o3d.utility.Vector3dVector(xyzB.reshape(-1, 3)[maskB, :])
# collect pointclouds
pcd_list.append(pc_source)
# collect landmarks with r2d2 network
img = db.get_image(target_id)
mask = db.get_valid_range_mask(target_id)
xys, scores, desc = extract.extract_keypoints(img, net_args, net)
xys, scores, desc = remove_invalid_keypoints(xys, scores, desc, mask)
kpts = xyzB[xys[:, 1].astype(int), xys[:, 0].astype(int)]
desc_list.append(desc)
kpts_list.append(kpts)
# get odometry constraints by icp
transformation_icp, information_icp, fitness = pairwise_registration(pc_source,
pc_target,
desc_list[source_id],
desc_list[target_id],
kpts_list[source_id],
kpts_list[target_id],
icp_max_corr_dist)
# draw_registration_result(pc_source, pc_target, transformation_icp)
odometry = np.dot(transformation_icp, odometry)
pose_graph.nodes.append(o3d.registration.PoseGraphNode(np.linalg.inv(odometry)))
pose_graph.edges.append(o3d.registration.PoseGraphEdge(source_id,
target_id,
transformation_icp,
information_icp,
confidence=fitness*0.1,
uncertain=True))
pcd_list.append(pc_target)
print("\rPorgress 100.0%")
# get loop closure candidates based on BoVW hist neighbours
print("Building BoVW")
hist = build_bag_of_words_hist(desc_list, n_words=n_words)
print("Calculating lc candidates from BoVW")
nn_ids = get_bag_of_words_nearest_neighbours(hist, min_frame_dist=min_frame_dist, max_hist_dist=max_hist_dist)
# calculate transformation between lc candidates
print("Calculating transformation between lc candidates")
for source_id in range(len(nn_ids)):
progress = source_id / len(nn_ids)
print("\rProgress {:2.1%}".format(progress), end="")
if nn_ids[source_id].shape[0] == 0:
continue
desc_source = o3d.registration.Feature()
desc_source.data = desc_list[source_id].T
kpts_source = o3d.geometry.PointCloud()
kpts_source.points = o3d.utility.Vector3dVector(kpts_list[source_id])
pc_source = pcd_list[source_id]
for target_id in nn_ids[source_id]:
desc_target = o3d.registration.Feature()
desc_target.data = desc_list[target_id].T
kpts_target = o3d.geometry.PointCloud()
kpts_target.points = o3d.utility.Vector3dVector(kpts_list[target_id])
pc_target = pcd_list[target_id]
tf, icp_info, fitness = global_registration(kpts_source=kpts_source,
kpts_target=kpts_target,
desc_source=desc_source,
desc_target=desc_target,
pc_source=pc_source,
pc_target=pc_target,
ransac_dist_threshold=ransac_max_dist,
icp_dist_threshold=icp_max_corr_dist)
# add found loop closure to pose graph if stable enough
if fitness < 0.4:
continue
pose_graph.edges.append(o3d.registration.PoseGraphEdge(source_id,
target_id,
tf,
icp_info,
confidence=fitness,
uncertain=False))
print("\rPorgress 100.0%")
return pose_graph, pcd_list
# vizulaization helper: get lines which connect the poses (trajectory)
def get_line_set_from_poses(pose_graph, color):
poses = pose_graph.nodes[0].pose[0:3, 3:]
for i in range(1, len(pose_graph.nodes)):
poses = np.concatenate((poses, pose_graph.nodes[i].pose[0:3, 3:]), axis=1)
poses = poses.T
lines = np.arange(len(pose_graph.nodes) - 1)[..., np.newaxis]
lines = np.concatenate((lines, lines + 1), axis=1)
lines.astype(int)
colors = [color for i in range(len(lines))]
line_set = o3d.geometry.LineSet()
line_set.points = o3d.utility.Vector3dVector(poses)
line_set.lines = o3d.utility.Vector2iVector(lines)
line_set.colors = o3d.utility.Vector3dVector(colors)
return line_set
# vizulaization helper: get loop closure edges as lines
def get_line_set_from_loop_closures(pose_graph, color):
poses = pose_graph.nodes[0].pose[0:3, 3:]
for i in range(1, len(pose_graph.nodes)):
poses = np.concatenate((poses, pose_graph.nodes[i].pose[0:3, 3:]), axis=1)
poses = poses.T
lc_edges = list()
for i in range(len(pose_graph.nodes), len(pose_graph.edges)):
edge = pose_graph.edges[i]
lc_edges.append([edge.source_node_id, edge.target_node_id])
colors = [color for i in range(len(lc_edges))]
line_set = o3d.geometry.LineSet()
line_set.points = o3d.utility.Vector3dVector(poses)
line_set.lines = o3d.utility.Vector2iVector(lc_edges)
line_set.colors = o3d.utility.Vector3dVector(colors)
return line_set
# vizulaization helper: draw point cloud alignment for every loop closure edge
def draw_loop_closure_matching(pose_graph, pcd_list):
for i in range(len(pose_graph.nodes), len(pose_graph.edges)):
edge = pose_graph.edges[i]
pc_source = pcd_list[edge.source_node_id]
pc_target = pcd_list[edge.target_node_id]
tf = edge.transformation
draw_registration_result(pc_source, pc_target, tf)
# write poses to csv (with timestamps)
def write_stamped_pose_graph_to_csv(pose_graph, timestamps, file_path):
pose = pose_graph.nodes[0].pose
poses = np.concatenate(([timestamps[0]], pose[0:3, 3:].squeeze(), R.from_matrix(pose[0:3, 0:3]).as_quat()))[np.newaxis, ...]
for i in range(1, len(pose_graph.nodes)):
pose = pose_graph.nodes[i].pose
poses = np.concatenate((poses, np.concatenate(([timestamps[i]],
pose[0:3, 3:].squeeze(),
R.from_matrix(pose[0:3, 0:3]).as_quat()))[np.newaxis, ...]), axis=0)
df = pd.DataFrame(poses, columns=['#ts', 'tx', 'ty', 'tz', 'qx', 'qy', 'qz', 'qw'])
df.to_csv(file_path, index=False, sep=" ")
# initilaize arguments
args = dict()
args["icp_max_corr_dist"] = 0.25
args["ransac_max_dist"] = 0.3
args["min_frame_dist"] = 30
args["max_hist_dist"] = 0.85
args["n_words"] = 180
# args["root"] = '/media/dominic/Extreme SSD/datasets/asl_koze/data_both_side/'
args["root"] = '/media/dominic/Extreme SSD/datasets/lidarmace_data/ethz_outside/'
# args["root"] = '/media/dominic/Extreme SSD/datasets/new_college/long_experiment/'
# # asl_both_side
# args["db_start_idx"] = 32
# args["db_stop_idx"] = -1
# args["db_idx_skip"] = 5
# lee_outside
args["db_start_idx"] = 3
args["db_stop_idx"] = -516
args["db_idx_skip"] = 5
# newer_college
# args["db_start_idx"] = 100
# args["db_stop_idx"] = -1
# args["db_idx_skip"] = 20
args["type"] = "OS0"
args["nn_model_path"] = "models/true_32.pt"
args["csv_export_path"] = args["root"] + "stamped_traj_estimate.txt"
args["csv_export_path_lc"] = args["root"] + "stamped_traj_estimate_w-lc.txt"
db = LidarSynthetic(args["root"] + "/data/", crop=False, skip=(args["db_start_idx"], args["db_stop_idx"], args["db_idx_skip"]), type=args["type"])
# load the network...
net_args = Args(args["nn_model_path"])
net = extract.load_network(args["nn_model_path"])
net = net.cuda()
print("Full registration ...")
pose_graph, pcd_list = full_registration(db, net, net_args=net_args, args=args)
# save poses without loop closure
write_stamped_pose_graph_to_csv(pose_graph, db.img_folders, args["csv_export_path"])
line_set_1 = get_line_set_from_poses(pose_graph, [1, 0, 0])
lc_edges_1 = get_line_set_from_loop_closures(pose_graph, [0, 1, 0])
# draw_loop_closure_matching(pose_graph, pcd_list)
print("Optimizing PoseGraph ...")
option = o3d.registration.GlobalOptimizationOption(
max_correspondence_distance=0.25,
edge_prune_threshold=0.1,
reference_node=0)
o3d.registration.global_optimization(
pose_graph, o3d.registration.GlobalOptimizationLevenbergMarquardt(),
o3d.registration.GlobalOptimizationConvergenceCriteria(), option)
# save poses optimized with loop closures
write_stamped_pose_graph_to_csv(pose_graph, db.img_folders, args["csv_export_path_lc"])
line_set_2 = get_line_set_from_poses(pose_graph, [0, 0, 1])
# draw results
o3d.visualization.draw_geometries([line_set_1, lc_edges_1, line_set_2])
o3d.visualization.draw_geometries([line_set_1, lc_edges_1])
o3d.visualization.draw_geometries([line_set_2])