def main(configs):
# train and validation dataloaders
train_loader = make_data_loader(configs,
"train",
configs.train_batch_size,
num_threads=configs.train_num_thread,
shuffle=True)
val_loader = make_data_loader(configs,
"valid",
1,
num_threads=configs.val_num_thread,
shuffle=False)
trainer = BlindPnPTrainer(configs, train_loader, val_loader)
trainer.train()
def main(config, resume=False):
train_loader = make_data_loader(
config,
config.train_phase,
config.batch_size,
num_threads=config.train_num_thread)
# default: False
if config.test_valid:
val_loader = make_data_loader(
config,
config.val_phase,
config.val_batch_size,
num_threads=config.val_num_thread)
else:
val_loader = None
Trainer = get_trainer(config.trainer)
trainer = Trainer(
config=config,
data_loader=train_loader,
val_data_loader=val_loader,
)
trainer.train()
def main(config, resume=False):
train_loader = make_data_loader(config,
config.train_phase,
config.batch_size,
shuffle=True,
repeat=True,
num_workers=config.train_num_workers)
if config.test_valid:
val_loader = make_data_loader(config,
config.val_phase,
config.val_batch_size,
shuffle=True,
repeat=True,
num_workers=config.val_num_workers)
else:
val_loader = None
Trainer = get_trainer(config.trainer)
trainer = Trainer(
config=config,
data_loader=train_loader,
val_data_loader=val_loader,
)
trainer.train()
if config.final_test:
test_loader = make_data_loader(config,
"test",
config.val_batch_size,
num_workers=config.val_num_workers)
trainer.val_data_loader = test_loader
test_dict = trainer._valid_epoch()
test_loss = test_dict['loss']
trainer.writer.add_scalar('test/loss', test_loss, config.max_epoch)
logging.info(f" Test loss: {test_loss}")
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 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 main(config):
# loading the test data
val_data_loader = make_data_loader(configs, "valid", 1, num_threads=configs.val_num_thread, shuffle=False)
# no gradients
with torch.no_grad():
# Model initialization
Model = load_model(config.model)
model = Model(config)
# limited GPU
if config.gpu_inds > -1:
torch.cuda.set_device(config.gpu_inds)
device = torch.device('cuda', config.gpu_inds)
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
# load the weights
if config.weights:
checkpoint = torch.load(config.weights)
model.load_state_dict(checkpoint['state_dict'])
logging.info(model)
# evaluation model
model.eval()
num_data = 0
data_timer, matching_timer = Timer(), Timer()
tot_num_data = len(val_data_loader.dataset)
data_loader_iter = val_data_loader.__iter__()
# collecting the errors in rotation, errors in tranlsation, num of inliers, inlier ratios
measure_list = ["err_q", "err_t", "inlier_ratio"]
nb_outlier_ratios = config.outlier_ratios.shape[0]
eval_res = {}
for measure in measure_list:
eval_res[measure] = np.zeros((nb_outlier_ratios, tot_num_data))
# for gpu memory consideration
max_nb_points = 20000
for batch_idx in range(tot_num_data):
data_timer.tic()
matches, numInliers, input_3d_xyz, input_2d_xy, R_gt, t_gt, _ = data_loader_iter.next()
data_timer.toc()
p3d_cur = input_3d_xyz
p2d_cur = input_2d_xy
matches_cur = matches
# for each outlier_ratio level
for oind in range(nb_outlier_ratios):
outlier_ratio = config.outlier_ratios[oind]
if outlier_ratio > 0:
# adding outliers
nb_added_outlier_2d = min(int(round(numInliers.item() * outlier_ratio)), max_nb_points - numInliers.item())
nb_added_outlier_3d = min(int(round(numInliers.item() * outlier_ratio)), max_nb_points - numInliers.item())
if config.outlier_type == "synthetic":
# add synthetic outliers
p2d_outlier, p3d_outlier = generate_outliers_synthetic(input_2d_xy, input_3d_xyz, nb_added_outlier_2d, nb_added_outlier_3d)
elif config.outlier_type == "real":
# add real outliers, you can load real-outlier from the non-matchable-2d/3d.pkl
raise NameError('Please modify dataloader')
# p2d_outlier, p3d_outlier = generate_outliers_real(p2d_outlier_all, p3d_outlier_all, nb_added_outlier_2d, nb_added_outlier_3d)
else:
raise NameError('Invalid outlier type')
# padding the outliers
p2d_cur = torch.cat((input_2d_xy, p2d_outlier), -2)
p3d_cur = torch.cat((input_3d_xyz, p3d_outlier), -2)
# Fill ground-truth matching indexes with outliers
b = p2d_cur.size(0)
m = p3d_cur.size(-2)
n = p2d_cur.size(-2)
matches_cur = matches.new_full((b,m,n),0.0)
matches_cur[:,:numInliers.item(),:numInliers.item()] = matches
# print([p3d_cur.size(1), p2d_cur.size(1)])
p2d_cur, p3d_cur, R_gt, t_gt, matches_cur = p2d_cur.to(device), p3d_cur.to(device), R_gt.to(
device), t_gt.to(device), matches_cur.to(device)
# Compute output
matching_timer.tic()
prob_matrix = model(p3d_cur, p2d_cur)
matching_timer.toc()
# compute the topK correspondences
# note cv2.solvePnPRansac is not stable, sometimes wrong!
# please use https://github.com/k88joshi/posest, and cite the original paper
k = min(2000, round(p3d_cur.size(1) * p2d_cur.size(1))) # Choose at most 2000 points in the testing stage
_, P_topk_i = torch.topk(prob_matrix.flatten(start_dim=-2), k=k, dim=-1, largest=True, sorted=True)
p3d_indices = P_topk_i / prob_matrix.size(-1) # bxk (integer division)
p2d_indices = P_topk_i % prob_matrix.size(-1) # bxk
# let's check the inliner ratios within the topK matches
# retrieve the inlier/outlier 1/0 logit
inlier_inds = matches_cur[:, p3d_indices, p2d_indices].cpu().numpy()
# inlier ratio
inlier_ratio = np.sum(inlier_inds) / k * 100.0
# in case cannot be estimated
err_q = np.pi
err_t = np.inf
# more than 5 2D-3D matches
if k > 5:
# compute the rotation and translation error
K = np.float32(np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]))
dist_coeff = np.float32(np.array([0.0, 0.0, 0.0, 0.0]))
# RANSAC p3p
p2d_np = p2d_cur[0, p2d_indices[0, :k], :].cpu().numpy()
p3d_np = p3d_cur[0, p3d_indices[0, :k], :].cpu().numpy()
retval, rvec, tvec, inliers = cv2.solvePnPRansac(
p3d_np, p2d_np, K, dist_coeff,
iterationsCount=1000,
reprojectionError=0.01,
flags=cv2.SOLVEPNP_P3P)
if rvec is not None and tvec is not None:
R_est, _ = cv2.Rodrigues(rvec)
t_est = tvec
err_q, err_t = evaluate_R_t(R_est, t_est, R_gt[0, :, :].cpu().numpy(), t_gt.cpu().numpy())
torch.cuda.empty_cache()
eval_res["err_q"][oind, batch_idx] = err_q
eval_res["err_t"][oind, batch_idx] = err_t
eval_res["inlier_ratio"][oind, batch_idx] = inlier_ratio
logging.info(' '.join([
f"Validation iter {num_data} / {tot_num_data} : Data Loading Time: {data_timer.avg:.3f},",
f"outlier ratio: {outlier_ratio:.3f},",
f"Matching Time: {matching_timer.avg:.3f},",
f"err_rot: {err_q:.3f}, err_t: {err_t:.3f}, inlier_ratio: {inlier_ratio:.3f},",
]))
data_timer.reset()
num_data += 1
# after checking all the validation samples, let's calculate statistics
# for each outlier_ratio level
for oind in range(nb_outlier_ratios):
outlier_ratio = config.outlier_ratios[oind]
eval_res_cur = {}
for measure in measure_list:
eval_res_cur[measure] = eval_res[measure][oind,:]
recall = recalls(eval_res_cur)
logging.info(' '.join([
f" outlier_ratio: {outlier_ratio:.3f}, recall_rot: {recall[0]:.3f}, med. rot. : {recall[1]:.3f}, med. trans. : {recall[2]:.3f}, avg. inlier ratio: {recall[3]:.3f},",
]))