def demo(config):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
checkpoint = torch.load(config.model)
model = ResUNetBN2C(1,
16,
normalize_feature=True,
conv1_kernel_size=3,
D=3)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
model = model.to(device)
pcd = o3d.io.read_point_cloud(config.input)
xyz_down, feature = extract_features(model,
xyz=np.array(pcd.points),
voxel_size=config.voxel_size,
device=device,
skip_check=True)
vis_pcd = o3d.geometry.PointCloud()
vis_pcd.points = o3d.utility.Vector3dVector(xyz_down)
vis_pcd = get_colored_point_cloud_feature(vis_pcd,
feature.detach().cpu().numpy(),
config.voxel_size)
o3d.visualization.draw_geometries([vis_pcd])
def extract_features_batch(model, config, source_path, target_path, voxel_size,
device):
folders = get_folder_list(source_path)
assert len(
folders) > 0, f"Could not find 3DMatch folders under {source_path}"
logging.info(folders)
list_file = os.path.join(target_path, "list.txt")
f = open(list_file, "w")
timer, tmeter = Timer(), AverageMeter()
num_feat = 0
model.eval()
for fo in folders:
if 'evaluation' in fo:
continue
files = get_file_list(fo, ".ply")
fo_base = os.path.basename(fo)
f.write("%s %d\n" % (fo_base, len(files)))
for i, fi in enumerate(files):
# Extract features from a file
pcd = o3d.io.read_point_cloud(fi)
save_fn = "%s_%03d" % (fo_base, i)
if i % 100 == 0:
logging.info(f"{i} / {len(files)}: {save_fn}")
timer.tic()
xyz_down, feature = extract_features(model,
xyz=np.array(pcd.points),
rgb=None,
normal=None,
voxel_size=voxel_size,
device=device,
skip_check=True)
t = timer.toc()
if i > 0:
tmeter.update(t)
num_feat += len(xyz_down)
np.savez_compressed(os.path.join(target_path, save_fn),
points=np.array(pcd.points),
xyz=xyz_down,
feature=feature.detach().cpu().numpy())
if i % 20 == 0 and i > 0:
# 最后一项算的是每个点的特征提取时间
logging.info(
f'Average time: {tmeter.avg}, FPS: {num_feat / tmeter.sum}, time / feat: {tmeter.sum / num_feat}, '
)
f.close()