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()