def test_detector(args):
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
model = Detector(args)
model = model.cuda()
model.load_state_dict(torch.load(args.model_path))
model.eval()
velodyne_dir = os.path.join(args.data_dir, 'sequences', args.test_seq,
'velodyne_txt')
velodyne_names = os.listdir(velodyne_dir)
velodyne_names = sorted(velodyne_names)
for filename in velodyne_names:
filepath = os.path.join(velodyne_dir, filename)
kp_path = os.path.join(args.save_dir, "keypoints", filename)
pc, sn = get_pointcloud(filepath, args.npoints)
feature = torch.cat((pc, sn), dim=-1)
feature = feature.unsqueeze(0)
feature = feature.cuda()
startT = datetime.datetime.now()
kp, sigmas, _, _ = model(feature)
endT = datetime.datetime.now()
computation_time = (endT - startT).microseconds
kp_sigmas = torch.cat((kp, sigmas.unsqueeze(1)), dim=1)
kp_sigmas = kp_sigmas.squeeze().cpu().detach().numpy().transpose()
np.savetxt(kp_path, kp_sigmas)
def test_descriptor(args):
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
model = RSKDD(args)
model = model.cuda()
model.load_state_dict(torch.load(args.model_path))
model.eval()
velodyne_dir = os.path.join(args.data_dir, 'sequences', args.test_seq,
'velodyne_txt')
velodyne_names = glob.glob(os.path.join(velodyne_dir, '*.txt'))
velodyne_names = sorted(velodyne_names)
kp_save_dir = os.path.join(args.save_dir, args.test_seq, "keypoints")
desc_save_dir = os.path.join(args.save_dir, args.test_seq, "desc")
if not os.path.exists(kp_save_dir):
os.makedirs(kp_save_dir)
if not os.path.exists(desc_save_dir):
os.makedirs(desc_save_dir)
for filename in velodyne_names:
filepath = os.path.join(velodyne_dir, filename)
basename = filename.split('/')[-1]
kp_path = os.path.join(kp_save_dir, "keypoints", basename)
desc_path = os.path.join(desc_save_dir, "desc", basename)
pc, sn = get_pointcloud(filepath, args.npoints)
feature = torch.cat((pc, sn), dim=-1)
feature = feature.unsqueeze(0)
feature = feature.cuda()
startT = datetime.datetime.now()
kp, sigmas, desc = model(feature)
endT = datetime.datetime.now()
computation_time = (endT - startT).microseconds
kp_sigmas = torch.cat((kp, sigmas.unsqueeze(1)), dim=1)
kp_sigmas = kp_sigmas.squeeze().cpu().detach().numpy().transpose()
desc = desc.squeeze().cpu().detach().numpy().transpose()
print(filename, computation_time)
np.savetxt(kp_path, kp_sigmas)
np.savetxt(desc_path, desc)
def demo(args):
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
model = RSKDD(args)
model = model.cuda()
model.load_state_dict(torch.load(args.model_path))
model.eval()
file_names = os.listdir(args.data_dir)
kp_save_dir = os.path.join(args.save_dir, "keypoints")
desc_save_dir = os.path.join(args.save_dir, "desc")
if not os.path.exists(kp_save_dir):
os.makedirs(kp_save_dir)
if not os.path.exists(desc_save_dir):
os.makedirs(desc_save_dir)
for file_name in file_names:
file_path = os.path.join(args.data_dir, file_name)
kp_save_path = os.path.join(kp_save_dir, file_name)
desc_save_path = os.path.join(desc_save_dir, file_name)
pc, sn = get_pointcloud(file_path, args.npoints)
feature = torch.cat((pc, sn), dim=-1)
feature = feature.unsqueeze(0)
feature = feature.cuda()
startT = datetime.datetime.now()
kp, sigmas, desc = model(feature)
endT = datetime.datetime.now()
computation_time = (endT - startT).microseconds
kp_sigmas = torch.cat((kp, sigmas.unsqueeze(1)), dim=1)
kp_sigmas = kp_sigmas.squeeze().cpu().detach().numpy().transpose()
desc = desc.squeeze().cpu().detach().numpy().transpose()
print(file_name, "processed",
' computation time: {} ms'.format(computation_time))
np.savetxt(kp_save_path, kp_sigmas, fmt='%.04f')
np.savetxt(desc_save_path, desc, fmt='%.04f')
print("Done")