def register_FCGF(self, xyz0, xyz1, inlier_thr=0.00):
'''
Main algorithm of DeepGlobalRegistration
'''
self.reg_timer.tic()
with torch.no_grad():
# Step 0: voxelize and generate sparse input
xyz0, coords0, feats0 = self.preprocess(xyz0)
xyz1, coords1, feats1 = self.preprocess(xyz1)
# Step 1: Feature extraction
self.feat_timer.tic()
fcgf_feats0 = self.fcgf_feature_extraction(feats0, coords0)
fcgf_feats1 = self.fcgf_feature_extraction(feats1, coords1)
self.feat_timer.toc()
# Step 2: Coarse correspondences
corres_idx0, corres_idx1 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1)
# > Case 1: Safeguard RANSAC + (Optional) ICP
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
T = self.safeguard_registration(pcd0,
pcd1,
corres_idx0,
corres_idx1,
feats0,
feats1,
2 * self.voxel_size,
num_iterations=80000)
safeguard_time = self.reg_timer.toc()
print(f'=> Safeguard takes {safeguard_time:.2} s')
return T
def do_single_pair_evaluation(feature_path,
set_name,
traj,
voxel_size,
tau_1=0.1,
tau_2=0.05,
num_rand_keypoints=-1):
trans_gth = np.linalg.inv(traj.pose)
i = traj.metadata[0]
j = traj.metadata[1]
name_i = "%s_%03d" % (set_name, i)
name_j = "%s_%03d" % (set_name, j)
# coord and feat form a sparse tensor.
data_i = np.load(os.path.join(feature_path, name_i + ".npz"))
coord_i, points_i, feat_i = data_i['xyz'], data_i['points'], data_i[
'feature']
data_j = np.load(os.path.join(feature_path, name_j + ".npz"))
coord_j, points_j, feat_j = data_j['xyz'], data_j['points'], data_j[
'feature']
# use the keypoints in 3DMatch
if num_rand_keypoints > 0:
# Randomly subsample N points
Ni, Nj = len(points_i), len(points_j)
inds_i = np.random.choice(Ni,
min(Ni, num_rand_keypoints),
replace=False)
inds_j = np.random.choice(Nj,
min(Nj, num_rand_keypoints),
replace=False)
sample_i, sample_j = points_i[inds_i], points_j[inds_j]
key_points_i = ME.utils.fnv_hash_vec(np.floor(sample_i / voxel_size))
key_points_j = ME.utils.fnv_hash_vec(np.floor(sample_j / voxel_size))
key_coords_i = ME.utils.fnv_hash_vec(np.floor(coord_i / voxel_size))
key_coords_j = ME.utils.fnv_hash_vec(np.floor(coord_j / voxel_size))
inds_i = np.where(np.isin(key_coords_i, key_points_i))[0]
inds_j = np.where(np.isin(key_coords_j, key_points_j))[0]
coord_i, feat_i = coord_i[inds_i], feat_i[inds_i]
coord_j, feat_j = coord_j[inds_j], feat_j[inds_j]
coord_i = make_open3d_point_cloud(coord_i)
coord_j = make_open3d_point_cloud(coord_j)
hit_ratio = evaluate_feature_3dmatch(coord_i, coord_j, feat_i, feat_j,
trans_gth, tau_1)
# logging.info(f"Hit ratio of {name_i}, {name_j}: {hit_ratio}, {hit_ratio >= tau_2}")
if hit_ratio >= tau_2:
return True
else:
return False
def register(self, xyz0, xyz1, inlier_thr=0.00):
'''
Main algorithm of DeepGlobalRegistration
'''
self.reg_timer.tic()
T = o3d.registration.registration_icp(
make_open3d_point_cloud(xyz0),
make_open3d_point_cloud(xyz1), self.voxel_size * 2, np.identity(4),
o3d.registration.TransformationEstimationPointToPoint()).transformation
safeguard_time = self.reg_timer.toc()
print(f'=> ICP takes {safeguard_time:.2} s')
return T
def read_data(feature_path, name):
data = np.load(os.path.join(feature_path, name + ".npz"))
xyz = make_open3d_point_cloud(data['xyz'])
feat = make_open3d_feature_from_numpy(data['feature'])
print("In benchmark_util.py - data['feature'] : ", data['feature'])
print("In benchmark_util.py - data['feature'] shape : ", data['feature'].shape)
print("In benchmark_util.py - feat : ", feat)
print('-'*20)
return data['points'], xyz, feat
def register(self, xyz0, xyz1, T_gt=None):
'''
Main algorithm of DeepGlobalRegistration
'''
self.reg_timer.tic()
with torch.no_grad():
# Step 0: voxelize and generate sparse input
xyz0, coords0, feats0 = self.preprocess(xyz0)
xyz1, coords1, feats1 = self.preprocess(xyz1)
# Step 1: Feature extraction
self.feat_timer.tic()
## generate fpfh features
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
pcd0.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*2, max_nn=30))
fpfh0 = o3d.registration.compute_fpfh_feature(pcd0, o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*5, max_nn=100))
fpfh_np0 = np.array(fpfh0.data).T
pcd1.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*2, max_nn=30))
fpfh1 = o3d.registration.compute_fpfh_feature(pcd1, o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*5, max_nn=100))
fpfh_np1 = np.array(fpfh1.data).T
fcgf_feats0 = torch.from_numpy(fpfh_np0).cuda()
fcgf_feats1 = torch.from_numpy(fpfh_np1).cuda()
import torch.nn.functional as F
fcgf_feats0 = F.normalize(fcgf_feats0, dim=-1)
fcgf_feats1 = F.normalize(fcgf_feats1, dim=-1)
# fcgf_feats0 = self.fcgf_feature_extraction(feats0, coords0)
# fcgf_feats1 = self.fcgf_feature_extraction(feats1, coords1)
self.feat_timer.toc()
# np.savez_compressed(
# os.path.join('KITTI_trainval/', f'kitti_pair_{i}'),
# xyz0=xyz0.detach().cpu().numpy(),
# features0=fcgf_feats0.detach().cpu().numpy(),
# xyz1=xyz1.detach().cpu().numpy(),
# features1=fcgf_feats1.detach().cpu().numpy(),
# gt_trans=T_gt
# )
# Step 2: Coarse correspondences
corres_idx0, corres_idx1 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1)
# _, corres_idx0 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1)
# _, corres_idx1 = self.fcgf_feature_matching(fcgf_feats1, fcgf_feats0)
# mutual_nearest = (corres_idx1[corres_idx0] == torch.arange(corres_idx0.shape[0]).cuda() )
# # corr = []
# # for i in range(len(corres_idx0)):
# # if corres_idx1[corres_idx0[i]] == i:
# # corr.append([i, corres_idx0[i]])
# mutual_corres_idx0 = torch.where(mutual_nearest == 1)[0]
# mutual_corres_idx1 = corres_idx0[mutual_nearest]
# corres_idx0, corres_idx1 = mutual_corres_idx0, mutual_corres_idx1
# Step 3: Inlier feature generation
# coords[corres_idx0]: 1D temporal + 3D spatial coord
# coords[corres_idx1, 1:]: 3D spatial coord
# => 1D temporal + 6D spatial coord
inlier_coords = torch.cat((coords0[corres_idx0], coords1[corres_idx1, 1:]),
dim=1).int()
inlier_feats = self.inlier_feature_generation(xyz0, xyz1, coords0, coords1,
fcgf_feats0, fcgf_feats1,
corres_idx0, corres_idx1)
# Step 4: Inlier likelihood estimation and truncation
import time
s_time = time.time()
logit = self.inlier_prediction(inlier_feats.contiguous(), coords=inlier_coords)
weights = logit.sigmoid()
if self.clip_weight_thresh > 0:
weights[weights < self.clip_weight_thresh] = 0
wsum = weights.sum().item()
# Step 5: Registration. Note: torch's gradient may be required at this stage
# > Case 0: Weighted Procrustes + Robust Refinement
wsum_threshold = max(200, len(weights) * 0.05)
# wsum_threshold = -1
sign = '>=' if wsum >= wsum_threshold else '<'
print(f'=> Weighted sum {wsum:.2f} {sign} threshold {wsum_threshold}')
T = np.identity(4)
safe_guard = 0
if wsum >= wsum_threshold:
try:
rot, trans, opt_output = GlobalRegistration(xyz0[corres_idx0],
xyz1[corres_idx1],
weights=weights.detach().cpu(),
break_threshold_ratio=1e-4,
quantization_size=2 *
self.voxel_size,
verbose=False)
T[0:3, 0:3] = rot.detach().cpu().numpy()
T[0:3, 3] = trans.detach().cpu().numpy()
dgr_time = self.reg_timer.toc()
print(f'=> DGR takes {dgr_time:.2} s')
except RuntimeError:
# Will directly go to Safeguard
print('###############################################')
print('# WARNING: SVD failed, weights sum: ', wsum)
print('# Falling back to Safeguard')
print('###############################################')
else:
safe_guard = 1
# > Case 1: Safeguard RANSAC + (Optional) ICP
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
T = self.safeguard_registration(pcd0,
pcd1,
corres_idx0,
corres_idx1,
feats0,
feats1,
2 * self.voxel_size,
num_iterations=80000)
safeguard_time = self.reg_timer.toc()
print(f'=> Safeguard takes {safeguard_time:.2} s')
if self.use_icp:
T = o3d.registration.registration_icp(
make_open3d_point_cloud(xyz0),
make_open3d_point_cloud(xyz1), self.voxel_size * 2, T,
o3d.registration.TransformationEstimationPointToPoint()).transformation
e_time = time.time()
eps = np.finfo(float).eps
inlier_xyz0 = np.array(xyz0[corres_idx0])
inlier_xyz1 = np.array(xyz1[corres_idx1])
pcd0 = o3d.geometry.PointCloud()
pcd0.points = o3d.utility.Vector3dVector(inlier_xyz0)
pcd0.transform(T_gt)
inlier_xyz0_warped = np.array(pcd0.points)
dis = np.sqrt( np.sum( (inlier_xyz0_warped- inlier_xyz1)**2, -1) + eps)
is_correct = dis < 2 * self.voxel_size
target = torch.from_numpy(is_correct).squeeze()
neg_target = (~target).to(torch.bool)
pred = weights > self.clip_weight_thresh
pred_on_pos, pred_on_neg = pred[target], pred[neg_target]
tp = pred_on_pos.sum().item()
fp = pred_on_neg.sum().item()
tn = (~pred_on_neg).sum().item()
fn = (~pred_on_pos).sum().item()
precision = tp / (tp + fp + eps)
recall = tp / (tp + fn + eps)
return T, precision, recall, e_time - s_time, safe_guard
def __getitem__(self, idx):
drive = self.files[idx][0]
t0, t1 = self.files[idx][1], self.files[idx][2]
all_odometry = self.get_video_odometry(drive, [t0, t1])
positions = [self.odometry_to_positions(odometry) for odometry in all_odometry]
fname0 = self._get_velodyne_fn(drive, t0)
fname1 = self._get_velodyne_fn(drive, t1)
# XYZ and reflectance
xyzr0 = np.fromfile(fname0, dtype=np.float32).reshape(-1, 4)
xyzr1 = np.fromfile(fname1, dtype=np.float32).reshape(-1, 4)
xyz0 = xyzr0[:, :3]
xyz1 = xyzr1[:, :3]
key = '%d_%d_%d' % (drive, t0, t1)
filename = self.icp_path + '/' + key + '.npy'
if key not in kitti_icp_cache:
if not os.path.exists(filename):
# work on the downsampled xyzs, 0.05m == 5cm
_, sel0 = ME.utils.sparse_quantize(xyz0 / 0.05, return_index=True)
_, sel1 = ME.utils.sparse_quantize(xyz1 / 0.05, return_index=True)
M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T)
@ np.linalg.inv(self.velo2cam)).T
xyz0_t = self.apply_transform(xyz0[sel0], M)
pcd0 = make_open3d_point_cloud(xyz0_t)
pcd1 = make_open3d_point_cloud(xyz1[sel1])
reg = o3d.registration.registration_icp(
pcd0, pcd1, 0.2, np.eye(4),
o3d.registration.TransformationEstimationPointToPoint(),
o3d.registration.ICPConvergenceCriteria(max_iteration=200))
pcd0.transform(reg.transformation)
# pcd0.transform(M2) or self.apply_transform(xyz0, M2)
M2 = M @ reg.transformation
# o3d.draw_geometries([pcd0, pcd1])
# write to a file
np.save(filename, M2)
else:
M2 = np.load(filename)
kitti_icp_cache[key] = M2
else:
M2 = kitti_icp_cache[key]
if self.random_rotation:
T0 = sample_random_trans(xyz0, self.randg, np.pi / 4)
T1 = sample_random_trans(xyz1, self.randg, np.pi / 4)
trans = T1 @ M2 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = M2
matching_search_voxel_size = self.matching_search_voxel_size
if self.random_scale and random.random() < 0.95:
scale = self.min_scale + \
(self.max_scale - self.min_scale) * random.random()
matching_search_voxel_size *= scale
xyz0 = scale * xyz0
xyz1 = scale * xyz1
# Voxelization
xyz0_th = torch.from_numpy(xyz0)
xyz1_th = torch.from_numpy(xyz1)
_, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True)
_, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True)
# Make point clouds using voxelized points
pcd0 = make_open3d_point_cloud(xyz0[sel0])
pcd1 = make_open3d_point_cloud(xyz1[sel1])
# Get matches
matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size)
if len(matches) < 1000:
raise ValueError(f"{drive}, {t0}, {t1}")
# Get features
npts0 = len(sel0)
npts1 = len(sel1)
feats_train0, feats_train1 = [], []
unique_xyz0_th = xyz0_th[sel0]
unique_xyz1_th = xyz1_th[sel1]
feats_train0.append(torch.ones((npts0, 1)))
feats_train1.append(torch.ones((npts1, 1)))
feats0 = torch.cat(feats_train0, 1)
feats1 = torch.cat(feats_train1, 1)
coords0 = torch.floor(unique_xyz0_th / self.voxel_size)
coords1 = torch.floor(unique_xyz1_th / self.voxel_size)
if self.transform:
coords0, feats0 = self.transform(coords0, feats0)
coords1, feats1 = self.transform(coords1, feats1)
return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(),
coords1.int(), feats0.float(), feats1.float(), matches, trans)
def __getitem__(self, idx):
file0 = os.path.join(self.root, self.files[idx][0])
file1 = os.path.join(self.root, self.files[idx][1])
data0 = np.load(file0)
data1 = np.load(file1)
xyz0 = data0["pcd"]
xyz1 = data1["pcd"]
color0 = data0["color"]
color1 = data1["color"]
matching_search_voxel_size = self.matching_search_voxel_size
if self.random_scale and random.random() < 0.95:
scale = self.min_scale + \
(self.max_scale - self.min_scale) * random.random()
matching_search_voxel_size *= scale
xyz0 = scale * xyz0
xyz1 = scale * xyz1
if self.random_rotation:
T0 = sample_random_trans(xyz0, self.randg, self.rotation_range)
T1 = sample_random_trans(xyz1, self.randg, self.rotation_range)
trans = T1 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = np.identity(4)
# Voxelization
_, sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size, return_index=True)
_, sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size, return_index=True)
# Make point clouds using voxelized points
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
# Select features and points using the returned voxelized indices
pcd0.colors = o3d.utility.Vector3dVector(color0[sel0])
pcd1.colors = o3d.utility.Vector3dVector(color1[sel1])
pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0])
pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1])
# Get matches
matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size)
# Get features
npts0 = len(pcd0.colors)
npts1 = len(pcd1.colors)
feats_train0, feats_train1 = [], []
feats_train0.append(np.ones((npts0, 1)))
feats_train1.append(np.ones((npts1, 1)))
feats0 = np.hstack(feats_train0)
feats1 = np.hstack(feats_train1)
# Get coords
xyz0 = np.array(pcd0.points)
xyz1 = np.array(pcd1.points)
coords0 = np.floor(xyz0 / self.voxel_size)
coords1 = np.floor(xyz1 / self.voxel_size)
if self.transform:
coords0, feats0 = self.transform(coords0, feats0)
coords1, feats1 = self.transform(coords1, feats1)
return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def read_data(feature_path, name):
data = np.load(os.path.join(feature_path, name + ".npz"))
xyz = make_open3d_point_cloud(data['xyz'])
feat = make_open3d_feature_from_numpy(data['feature'])
return data['points'], xyz, feat
def __getitem__(self, idx):
file0 = os.path.join(self.root, self.files[idx][0])
file1 = os.path.join(self.root, self.files[idx][1])
data0 = np.load(file0)
data1 = np.load(file1)
# 这俩点云读进来是可以拼在一起的，只不过都只是完整点云的一部分
xyz0 = data0["pcd"]#50万左右个点
xyz1 = data1["pcd"]
color0 = data0["color"]#每个点都有它自己的颜色，且不一样
color1 = data1["color"]
matching_search_voxel_size = self.matching_search_voxel_size #0.0375
if self.random_scale and random.random() < 0.95:
scale = self.min_scale + \
(self.max_scale - self.min_scale) * random.random()
matching_search_voxel_size *= scale
xyz0 = scale * xyz0
xyz1 = scale * xyz1
if self.random_rotation:
T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) #rotation_range:360 且旋转平移之后的点云中心坐标为0
T1 = sample_random_trans(xyz1, self.randg, self.rotation_range)
# xyz0到xyz1的相对位姿
trans = T1 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = np.identity(4)
# Voxelization voxel_size 0.025
_, sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size, return_index=True)
_, sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size, return_index=True)
# 用稀疏点云生成pcd
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
# Select features and points using the returned voxelized indices
pcd0.colors = o3d.utility.Vector3dVector(color0[sel0])
pcd1.colors = o3d.utility.Vector3dVector(color1[sel1])
pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0])
pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1])
# Get matches
# matches包含了点云1的每一点在点云二的（一定范围内的 #0.0375）匹配点
matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size)
# Get features,特征都为1
npts0 = len(pcd0.colors)
npts1 = len(pcd1.colors)
feats_train0, feats_train1 = [], []
feats_train0.append(np.ones((npts0, 1)))
feats_train1.append(np.ones((npts1, 1)))
feats0 = np.hstack(feats_train0)
feats1 = np.hstack(feats_train1)
# Get coords
xyz0 = np.array(pcd0.points)
xyz1 = np.array(pcd1.points)
coords0 = np.floor(xyz0 / self.voxel_size)
coords1 = np.floor(xyz1 / self.voxel_size)
if self.transform:
coords0, feats0 = self.transform(coords0, feats0)
coords1, feats1 = self.transform(coords1, feats1)
# 返回体素化之后的np点云，稀疏点云坐标、特征，匹配关系，真实旋转平移
return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def main(config):
test_loader = make_data_loader(
config, config.test_phase, 1, num_threads=config.test_num_workers, shuffle=True)
num_feats = 1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Model = load_model(config.model)
model = Model(
num_feats,
config.model_n_out,
bn_momentum=config.bn_momentum,
conv1_kernel_size=config.conv1_kernel_size,
normalize_feature=config.normalize_feature)
checkpoint = torch.load(config.save_dir + '/checkpoint.pth')
model.load_state_dict(checkpoint['state_dict'])
model = model.to(device)
model.eval()
success_meter, rte_meter, rre_meter = AverageMeter(), AverageMeter(), AverageMeter()
data_timer, feat_timer, reg_timer = Timer(), Timer(), Timer()
test_iter = test_loader.__iter__()
N = len(test_iter)
n_gpu_failures = 0
# downsample_voxel_size = 2 * config.voxel_size
for i in range(len(test_iter)):
data_timer.tic()
try:
data_dict = test_iter.next()
except ValueError:
n_gpu_failures += 1
logging.info(f"# Erroneous GPU Pair {n_gpu_failures}")
continue
data_timer.toc()
xyz0, xyz1 = data_dict['pcd0'], data_dict['pcd1']
T_gth = data_dict['T_gt']
xyz0np, xyz1np = xyz0.numpy(), xyz1.numpy()
pcd0 = make_open3d_point_cloud(xyz0np)
pcd1 = make_open3d_point_cloud(xyz1np)
with torch.no_grad():
feat_timer.tic()
sinput0 = ME.SparseTensor(
data_dict['sinput0_F'].to(device), coordinates=data_dict['sinput0_C'].to(device))
F0 = model(sinput0).F.detach()
sinput1 = ME.SparseTensor(
data_dict['sinput1_F'].to(device), coordinates=data_dict['sinput1_C'].to(device))
F1 = model(sinput1).F.detach()
feat_timer.toc()
feat0 = make_open3d_feature(F0, 32, F0.shape[0])
feat1 = make_open3d_feature(F1, 32, F1.shape[0])
reg_timer.tic()
distance_threshold = config.voxel_size * 1.0
ransac_result = o3d.registration.registration_ransac_based_on_feature_matching(
pcd0, pcd1, feat0, feat1, distance_threshold,
o3d.registration.TransformationEstimationPointToPoint(False), 4, [
o3d.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
o3d.registration.CorrespondenceCheckerBasedOnDistance(
distance_threshold)
], o3d.registration.RANSACConvergenceCriteria(4000000, 10000))
T_ransac = torch.from_numpy(
ransac_result.transformation.astype(np.float32))
reg_timer.toc()
# Translation error
rte = np.linalg.norm(T_ransac[:3, 3] - T_gth[:3, 3])
rre = np.arccos((np.trace(T_ransac[:3, :3].t() @ T_gth[:3, :3]) - 1) / 2)
# Check if the ransac was successful. successful if rte < 2m and rre < 5◦
# http://openaccess.thecvf.com/content_ECCV_2018/papers/Zi_Jian_Yew_3DFeat-Net_Weakly_Supervised_ECCV_2018_paper.pdf
if rte < 2:
rte_meter.update(rte)
if not np.isnan(rre) and rre < np.pi / 180 * 5:
rre_meter.update(rre)
if rte < 2 and not np.isnan(rre) and rre < np.pi / 180 * 5:
success_meter.update(1)
else:
success_meter.update(0)
logging.info(f"Failed with RTE: {rte}, RRE: {rre}")
if i % 10 == 0:
logging.info(
f"{i} / {N}: Data time: {data_timer.avg}, Feat time: {feat_timer.avg}," +
f" Reg time: {reg_timer.avg}, RTE: {rte_meter.avg}," +
f" RRE: {rre_meter.avg}, Success: {success_meter.sum} / {success_meter.count}"
+ f" ({success_meter.avg * 100} %)")
data_timer.reset()
feat_timer.reset()
reg_timer.reset()
logging.info(
f"RTE: {rte_meter.avg}, var: {rte_meter.var}," +
f" RRE: {rre_meter.avg}, var: {rre_meter.var}, Success: {success_meter.sum} " +
f"/ {success_meter.count} ({success_meter.avg * 100} %)")
def __getitem__(self, idx):
drive = self.files[idx][0]
t0, t1 = self.files[idx][1], self.files[idx][2]
all_odometry = self.get_video_odometry(drive, [t0, t1])
positions = [self.odometry_to_positions(odometry) for odometry in all_odometry]
fname0 = self._get_velodyne_fn(drive, t0)
fname1 = self._get_velodyne_fn(drive, t1)
# XYZ and reflectance
xyzr0 = np.fromfile(fname0, dtype=np.float32).reshape(-1, 4)
xyzr1 = np.fromfile(fname1, dtype=np.float32).reshape(-1, 4)
xyz0 = xyzr0[:, :3]
xyz1 = xyzr1[:, :3]
coords0 = (xyz0 - xyz0.min(0)) / 0.05
coords1 = (xyz1 - xyz1.min(0)) / 0.05
sel0 = ME.utils.sparse_quantize(coords0, return_index=True)
sel1 = ME.utils.sparse_quantize(coords1, return_index=True)
xyz0 = xyz0[sel0]
xyz1 = xyz1[sel1]
# r0 = xyzr0[:, -1].reshape(-1, 1)
# r1 = xyzr1[:, -1].reshape(-1, 1)
# pcd0 = make_open3d_point_cloud(xyz0_t, 0.7 * np.ones((len(xyz0), 3)))
# pcd1 = make_open3d_point_cloud(xyz1, 0.3 * np.ones((len(xyz1), 3)))
key = '%d_%d_%d' % (drive, t0, t1)
filename = self.icp_path + '/' + key + '.npy'
if key not in kitti_icp_cache:
if not os.path.exists(filename):
if self.IS_ODOMETRY:
M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T)
@ np.linalg.inv(self.velo2cam)).T
else:
M = self.get_position_transform(positions[0], positions[1], invert=True).T
xyz0_t = self.apply_transform(xyz0, M)
pcd0 = make_open3d_point_cloud(xyz0_t)
pcd1 = make_open3d_point_cloud(xyz1)
reg = o3d.registration_icp(pcd0, pcd1, 0.2, np.eye(4),
o3d.TransformationEstimationPointToPoint(),
o3d.ICPConvergenceCriteria(max_iteration=200))
pcd0.transform(reg.transformation)
# pcd0.transform(M2) or self.apply_transform(xyz0, M2)
M2 = M @ reg.transformation
# o3d.draw_geometries([pcd0, pcd1])
# write to a file
np.save(filename, M2)
else:
M2 = np.load(filename)
kitti_icp_cache[key] = M2
else:
M2 = kitti_icp_cache[key]
if self.random_rotation:
T0 = sample_random_trans(xyz0, self.randg, np.pi / 4)
T1 = sample_random_trans(xyz1, self.randg, np.pi / 4)
trans = T1 @ M2 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = M2
matching_search_voxel_size = self.matching_search_voxel_size
if self.random_scale and random.random() < 0.95:
scale = self.min_scale + \
(self.max_scale - self.min_scale) * random.random()
matching_search_voxel_size *= scale
xyz0 = scale * xyz0
xyz1 = scale * xyz1
# Voxelization
coords0 = np.floor(xyz0 / self.voxel_size)
coords1 = np.floor(xyz1 / self.voxel_size)
sel0 = ME.utils.sparse_quantize(coords0, return_index=True)
sel1 = ME.utils.sparse_quantize(coords1, return_index=True)
coords0, coords1 = coords0[sel0], coords1[sel1]
# r0, r1 = r0[sel0], r1[sel1]
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
# Select features and points using the returned voxelized indices
pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0])
pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1])
# Get matches
matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size)
if len(matches) < 1000:
raise ValueError(f"{drive}, {t0}, {t1}")
feats_train0, feats_train1 = [], []
feats_train0.append(np.ones((len(sel0), 1)))
feats_train1.append(np.ones((len(sel1), 1)))
feats0 = np.hstack(feats_train0)
feats1 = np.hstack(feats_train1)
# Get coords
xyz0 = np.array(pcd0.points)
xyz1 = np.array(pcd1.points)
if self.transform:
coords0, feats0 = self.transform(coords0, feats0)
coords1, feats1 = self.transform(coords1, feats1)
return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def __getitem__(self, idx):
file0 = os.path.join(self.root, self.files[idx][0])
file1 = os.path.join(self.root, self.files[idx][1])
data0 = np.load(file0)
data1 = np.load(file1)
xyz0 = data0["pcd"]
xyz1 = data1["pcd"]
matching_search_voxel_size = self.matching_search_voxel_size
if self.random_scale and random.random() < 0.95:
scale = self.min_scale + \
(self.max_scale - self.min_scale) * random.random()
matching_search_voxel_size *= scale
xyz0 = scale * xyz0
xyz1 = scale * xyz1
if self.random_rotation:
T0 = sample_random_trans(xyz0, self.randg, self.rotation_range)
T1 = sample_random_trans(xyz1, self.randg, self.rotation_range)
trans = T1 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = np.identity(4)
# Voxelization
xyz0_th = torch.from_numpy(xyz0)
xyz1_th = torch.from_numpy(xyz1)
_, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size,
return_index=True)
_, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size,
return_index=True)
# Make point clouds using voxelized points
pcd0 = make_open3d_point_cloud(xyz0[sel0])
pcd1 = make_open3d_point_cloud(xyz1[sel1])
# Select features and points using the returned voxelized indices
# 3DMatch color is not helpful
# pcd0.colors = o3d.utility.Vector3dVector(color0[sel0])
# pcd1.colors = o3d.utility.Vector3dVector(color1[sel1])
# Get matches
matches = get_matching_indices(pcd0, pcd1, trans,
matching_search_voxel_size)
# Get features
npts0 = len(sel0)
npts1 = len(sel1)
feats_train0, feats_train1 = [], []
unique_xyz0_th = xyz0_th[sel0]
unique_xyz1_th = xyz1_th[sel1]
# xyz as feats
if self.use_xyz_feature:
feats_train0.append(unique_xyz0_th - unique_xyz0_th.mean(0))
feats_train1.append(unique_xyz1_th - unique_xyz1_th.mean(0))
else:
feats_train0.append(torch.ones((npts0, 1)))
feats_train1.append(torch.ones((npts1, 1)))
feats0 = torch.cat(feats_train0, 1)
feats1 = torch.cat(feats_train1, 1)
coords0 = torch.floor(unique_xyz0_th / self.voxel_size)
coords1 = torch.floor(unique_xyz1_th / self.voxel_size)
if self.transform:
coords0, feats0 = self.transform(coords0, feats0)
coords1, feats1 = self.transform(coords1, feats1)
extra_package = {'idx': idx, 'file0': file0, 'file1': file1}
return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(),
coords1.int(), feats0.float(), feats1.float(), matches, trans,
extra_package)
def register(self, xyz0, xyz1, inlier_thr=0.00):
'''
Main algorithm of DeepGlobalRegistration
'''
self.reg_timer.tic()
with torch.no_grad():
# Step 0: voxelize and generate sparse input
xyz0, coords0, feats0 = self.preprocess(xyz0)
xyz1, coords1, feats1 = self.preprocess(xyz1)
# Step 1: Feature extraction
self.feat_timer.tic()
fcgf_feats0 = self.fcgf_feature_extraction(feats0, coords0)
fcgf_feats1 = self.fcgf_feature_extraction(feats1, coords1)
self.feat_timer.toc()
# Step 2: Coarse correspondences
corres_idx0, corres_idx1 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1)
# Step 3: Inlier feature generation
# coords[corres_idx0]: 1D temporal + 3D spatial coord
# coords[corres_idx1, 1:]: 3D spatial coord
# => 1D temporal + 6D spatial coord
inlier_coords = torch.cat((coords0[corres_idx0], coords1[corres_idx1, 1:]),
dim=1).int()
inlier_feats = self.inlier_feature_generation(xyz0, xyz1, coords0, coords1,
fcgf_feats0, fcgf_feats1,
corres_idx0, corres_idx1)
# Step 4: Inlier likelihood estimation and truncation
logit = self.inlier_prediction(inlier_feats.contiguous(), coords=inlier_coords)
weights = logit.sigmoid()
if self.clip_weight_thresh > 0:
weights[weights < self.clip_weight_thresh] = 0
wsum = weights.sum().item()
# Step 5: Registration. Note: torch's gradient may be required at this stage
# > Case 0: Weighted Procrustes + Robust Refinement
wsum_threshold = max(200, len(weights) * 0.05)
sign = '>=' if wsum >= wsum_threshold else '<'
print(f'=> Weighted sum {wsum:.2f} {sign} threshold {wsum_threshold}')
T = np.identity(4)
if wsum >= wsum_threshold:
try:
rot, trans, opt_output = GlobalRegistration(xyz0[corres_idx0],
xyz1[corres_idx1],
weights=weights.detach().cpu(),
break_threshold_ratio=1e-4,
quantization_size=2 *
self.voxel_size,
verbose=False)
T[0:3, 0:3] = rot.detach().cpu().numpy()
T[0:3, 3] = trans.detach().cpu().numpy()
dgr_time = self.reg_timer.toc()
print(f'=> DGR takes {dgr_time:.2} s')
except RuntimeError:
# Will directly go to Safeguard
print('###############################################')
print('# WARNING: SVD failed, weights sum: ', wsum)
print('# Falling back to Safeguard')
print('###############################################')
else:
# > Case 1: Safeguard RANSAC + (Optional) ICP
pcd0 = make_open3d_point_cloud(xyz0)
pcd1 = make_open3d_point_cloud(xyz1)
T = self.safeguard_registration(pcd0,
pcd1,
corres_idx0,
corres_idx1,
feats0,
feats1,
2 * self.voxel_size,
num_iterations=80000)
safeguard_time = self.reg_timer.toc()
print(f'=> Safeguard takes {safeguard_time:.2} s')
if self.use_icp:
T = o3d.registration.registration_icp(
make_open3d_point_cloud(xyz0),
make_open3d_point_cloud(xyz1), self.voxel_size * 2, T,
o3d.registration.TransformationEstimationPointToPoint()).transformation
return T
def __getitem__(self, idx):
drive = self.U.pairs[idx, 0]
t0 = self.U.pairs[idx, 1]
t1 = self.U.pairs[idx, 2]
M2, xyz0, xyz1 = self.U.get_pair(idx)
if self.random_rotation:
T0 = sample_almost_planar_rotation(self.randg)
T1 = sample_almost_planar_rotation(self.randg)
trans = T1 @ M2 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = M2
matching_search_voxel_size = self.matching_search_voxel_size
if self.random_scale and random.random() < 0.95:
scale = self.min_scale + \
(self.max_scale - self.min_scale) * random.random()
matching_search_voxel_size *= scale
xyz0 = scale * xyz0
xyz1 = scale * xyz1
M2[:3, 3] *= scale
# Voxelization
xyz0_th = torch.from_numpy(xyz0)
xyz1_th = torch.from_numpy(xyz1)
_, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size,
return_index=True)
_, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size,
return_index=True)
# Make point clouds using voxelized points
pcd0 = make_open3d_point_cloud(xyz0[sel0])
pcd1 = make_open3d_point_cloud(xyz1[sel1])
# Get matches
matches = get_matching_indices(pcd0, pcd1, trans,
matching_search_voxel_size)
# Get features
npts0 = len(sel0)
npts1 = len(sel1)
feats_train0, feats_train1 = [], []
unique_xyz0_th = xyz0_th[sel0]
unique_xyz1_th = xyz1_th[sel1]
feats_train0.append(torch.ones((npts0, 1)))
feats_train1.append(torch.ones((npts1, 1)))
feats0 = torch.cat(feats_train0, 1)
feats1 = torch.cat(feats_train1, 1)
coords0 = torch.floor(unique_xyz0_th / self.voxel_size)
coords1 = torch.floor(unique_xyz1_th / self.voxel_size)
if self.transform:
coords0, feats0 = self.transform(coords0, feats0)
coords1, feats1 = self.transform(coords1, feats1)
extra_package = {'drive': drive, 't0': t0, 't1': t1}
return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(),
coords1.int(), feats0.float(), feats1.float(), matches, trans,
extra_package)
def main(config):
test_loader = make_data_loader(config,
config.test_phase,
1,
num_threads=config.test_num_thread,
shuffle=False)
num_feats = 1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Model = load_model(config.model)
model = Model(num_feats,
config.model_n_out,
bn_momentum=config.bn_momentum,
conv1_kernel_size=config.conv1_kernel_size,
normalize_feature=config.normalize_feature)
checkpoint = torch.load(config.save_dir + '/checkpoint.pth')
model.load_state_dict(checkpoint['state_dict'])
model = model.to(device)
model.eval()
success_meter, rte_meter, rre_meter = AverageMeter(), AverageMeter(
), AverageMeter()
data_timer, feat_timer, reg_timer = Timer(), Timer(), Timer()
test_iter = test_loader.__iter__()
N = len(test_iter)
n_gpu_failures = 0
# downsample_voxel_size = 2 * config.voxel_size
list_results_to_save = []
for i in range(len(test_iter)):
data_timer.tic()
try:
data_dict = test_iter.next()
except ValueError:
n_gpu_failures += 1
logging.info(f"# Erroneous GPU Pair {n_gpu_failures}")
continue
data_timer.toc()
xyz0, xyz1 = data_dict['pcd0'], data_dict['pcd1']
T_gth = data_dict['T_gt']
xyz0np, xyz1np = xyz0.numpy(), xyz1.numpy()
#import pdb
# pdb.set_trace()
pcd0 = make_open3d_point_cloud(xyz0np)
pcd1 = make_open3d_point_cloud(xyz1np)
with torch.no_grad():
feat_timer.tic()
sinput0 = ME.SparseTensor(
data_dict['sinput0_F'].to(device),
coordinates=data_dict['sinput0_C'].to(device))
F0 = model(sinput0).F.detach()
sinput1 = ME.SparseTensor(
data_dict['sinput1_F'].to(device),
coordinates=data_dict['sinput1_C'].to(device))
F1 = model(sinput1).F.detach()
feat_timer.toc()
# saving files to pkl
print(i)
dict_sample = {
"pts_source": xyz0np,
"feat_source": F0.cpu().detach().numpy(),
"pts_target": xyz1np,
"feat_target": F1.cpu().detach().numpy()
}
list_results_to_save.append(dict_sample)
import pickle
path_results_to_save = "fcgf.results.pkl"
print('Saving results to ', path_results_to_save)
pickle.dump(list_results_to_save, open(path_results_to_save, 'wb'))
print('Saved!')
import pdb
pdb.set_trace()
def __getitem__(self, idx): # split, idx):
if self.split == "test":
sample = self.list_test_sample[idx]
else:
T_noise = self.generate_noise_T()
sample = self.get_sample(idx, T_noise)
xyz0 = sample["source"].points
xyz1 = sample["target"].points
trans = torch.Tensor(self.data["T_map"][idx])
matching_search_voxel_size = self.matching_search_voxel_size
import copy
#
sel0 = ME.utils.sparse_quantize(
xyz0.contiguous() / self.voxel_size, return_index=True)[1]
sel1 = ME.utils.sparse_quantize(
xyz1.contiguous() / self.voxel_size, return_index=True)[1]
unique_xyz0_th = xyz0[sel0] # [ind_0]
unique_xyz1_th = xyz1[sel1] # [ind_1]
pcd0 = make_open3d_point_cloud(unique_xyz0_th)
pcd1 = make_open3d_point_cloud(unique_xyz1_th)
# Get matches
matches = get_matching_indices(
pcd0, pcd1, trans, matching_search_voxel_size)
# Get features
feats_train0, feats_train1 = [], []
npts0 = unique_xyz0_th.shape[0]
npts1 = unique_xyz1_th.shape[0]
feats_train0.append(torch.ones((npts0, 1)))
feats_train1.append(torch.ones((npts1, 1)))
feats0 = torch.cat(feats_train0, 1)
feats1 = torch.cat(feats_train1, 1)
coords0 = np.floor(unique_xyz0_th / self.voxel_size)
coords1 = np.floor(unique_xyz1_th / self.voxel_size)
#coords0_mean = coords0.min(axis=0).int()
#coords0 -=coords0_mean
if False:#len(matches) < 300: # idx == 113:#len(matches) <
print(self.split, "num matches = ", len(matches))
print(coords0)
pcd_target = o3d.geometry.PointCloud()
pcd_target.points = o3d.utility.Vector3dVector(coords0)
o3d.io.write_point_cloud("coords0_before_%d.ply" % idx, pcd_target)
#
pcd_target = o3d.geometry.PointCloud()
pcd_target.points = o3d.utility.Vector3dVector(coords1)
o3d.io.write_point_cloud("coords1_before_%d.ply" % idx, pcd_target)
#
pcd_target = o3d.geometry.PointCloud()
pcd_target.points = o3d.utility.Vector3dVector(unique_xyz1_th)
o3d.io.write_point_cloud("unique_xyz1_th_%d.ply" % idx, pcd_target)
#
pcd_target = o3d.geometry.PointCloud()
pcd_target.points = o3d.utility.Vector3dVector(unique_xyz0_th)
o3d.io.write_point_cloud("unique_xyz0_th_%d.ply" % idx, pcd_target)
#
pcd_target = o3d.geometry.PointCloud()
pcd_target.points = o3d.utility.Vector3dVector(unique_xyz0_th)
pcd_target.transform(trans)
o3d.io.write_point_cloud(
"unique_xyz0_th_trans_%d.ply" % idx, pcd_target)
#
import pdb
pdb.set_trace()
#if self.transform: # add noises to the point clouds
# coords0, feats0 = self.transform(coords0, feats0)
# coords1, feats1 = self.transform(coords1, feats1)
return (unique_xyz0_th, unique_xyz1_th, coords0,
coords1, feats0.float(), feats1.float(), matches, trans)
