项目地址: https://github.com/zhulf0804/PCReg.PyTorch/tree/main
网络简介: 网络是基于PointNet + Concat + FC的,它没有其它复杂的结构,易于复现。因其简洁性,这里暂且把其称作点云配准的Benchmark。因作者源码中复杂的(四元数, 旋转矩阵, 欧拉角之间)的变换操作和冗余性,且其PyTorch版本的不完整性(缺少评估模型等,最近又更新了),
项目详细介绍: 基于深度学习的点云配准Benchmark
本文方法与常见的图像配准逻辑类似,基于采样与transfrom操作从源点云生成目标点云,然后进行训练与评测。总体看来效果不如open3d自带的fgr方法,可以作为入门级项目进行使用。
1、运行环境安装
1.1 项目下载
打开https://github.com/zhulf0804/PCReg.PyTorch/tree/main,点Download ZIP然后将代码解压到指定目录下即可。
1.2 依赖项安装
在装有pytorch的环境终端,进入PCReg.PyTorch-main目录,执行以下安装命令:
pip install -r requirements.txt
python -m pip install open3d>=0.9
emd loss编译
如果不做训练使用,可以不用进行emd loss编译。
cd loss/cuda/emd_torch & python setup.py install
在编译过程中,很有可能碰到报错
File "C:\Users\Administrator\miniconda3\lib\site-packages\torch\utils\cpp_extension.py", line 499, in build_extensions_check_cuda_version(compiler_name, compiler_version)File "C:\Users\Administrator\miniconda3\lib\site-packages\torch\utils\cpp_extension.py", line 387, in _check_cuda_versionraise RuntimeError(CUDA_MISMATCH_MESSAGE.format(cuda_str_version, torch.version.cuda))
RuntimeError:
The detected CUDA version (12.1) mismatches the version that was used to compile
PyTorch (11.7). Please make sure to use the same CUDA versions.
这是由于PyTorch 的cuda版本与系统自带的cuda版本不同所导致的,可以先使用一下命令卸载过往的torch版本(慎重操作),然后重新安装torch;也可以在conda环境中重新创建一个符合系统cuda版本的torch环境。
pip uninstall torch torchvision torchaudio
在上述信息输出中,博主的cuda版本为12.1,我们可以打开pytorch官网找打符合自己电脑cuda版本的pytorch安装命令.
如果cuda版本较早,可以在 https://pytorch.org/get-started/previous-versions/ 中找到安装命令。
在正确的安装cuda版本后,重新执行命令即可实现emd-loss的安装
1.3 模型与数据下载
-
modelnet40数据集 [here,
435M
] -
可用的预训练模型 [Complete, pwd:
c4z7
,16.09 M
] or [Paritial, pwd:pcno
,16.09
] first) 模型下载好后将其放置到PCReg.PyTorch项目根路径下即可。
2. 关键代码说明
2.1 数据加载器
在data目录下有CustomData.py和ModelNet40.py两个文件,其中ModelNet40文件对应modelnet40数据集的加载,CustomData文件对应自己个人数据集的加载。从两个文件的__getitem__函数中可以发现,模型不是基于数据对进行训练的。其依据ref_cloud随机采样生成ref_cloud,然后对ref_cloud进行transform操作。具体实例如下所示:
def __getitem__(self, item):file = self.files[item]ref_cloud = readpcd(file, rtype='npy')ref_cloud = random_select_points(ref_cloud, m=self.npts)ref_cloud = pc_normalize(ref_cloud)R, t = generate_random_rotation_matrix(-20, 20), \generate_random_tranlation_vector(-0.5, 0.5)src_cloud = transform(ref_cloud, R, t)if self.train:ref_cloud = jitter_point_cloud(ref_cloud)src_cloud = jitter_point_cloud(src_cloud)return ref_cloud, src_cloud, R, t
在以上代码中需要注意的是,所有的点云都进行了坐标值的归一化处理
2.2 模型结构
在model目录下有benchmark.py、fgr.py、icp.py,分别为模型配准,fgr配准,icp配准方法。其中fgr配准与icp配准方法是使用open3d库实现。
benchmark
benchmark为本文模型,其是基于PointNet所实现的一个孪生网络,核心代码如Benchmark类所示。其基于encoder提取2个点云的特征,然后简单的使用全连接层将两个点云的特征进行交互,然后再输出两个点云的特征。在这里最为重要的是loss的设计,即如何设计优化目标,使模型参数以点云配准为优化方向
class Benchmark(nn.Module):def __init__(self, gn, in_dim1, in_dim2=2048, fcs=[1024, 1024, 512, 512, 256, 7]):super(Benchmark, self).__init__()self.in_dim1 = in_dim1self.encoder = PointNet(in_dim=in_dim1, gn=gn)self.decoder = nn.Sequential()for i, out_dim in enumerate(fcs):self.decoder.add_module(f'fc_{i}', nn.Linear(in_dim2, out_dim))if out_dim != 7:if gn:self.decoder.add_module(f'gn_{i}',nn.GroupNorm(8, out_dim))self.decoder.add_module(f'relu_{i}', nn.ReLU(inplace=True))in_dim2 = out_dimdef forward(self, x, y):x_f, y_f = self.encoder(x), self.encoder(y)concat = torch.cat((x_f, y_f), dim=1)out = self.decoder(concat)batch_t, batch_quat = out[:, :3], out[:, 3:] / torch.norm(out[:, 3:], dim=1, keepdim=True)batch_R = batch_quat2mat(batch_quat)if self.in_dim1 == 3:transformed_x = batch_transform(x.permute(0, 2, 1).contiguous(),batch_R, batch_t)elif self.in_dim1 == 6:transformed_pts = batch_transform(x.permute(0, 2, 1)[:, :, :3].contiguous(),batch_R, batch_t)transformed_nls = batch_transform(x.permute(0, 2, 1)[:, :, 3:].contiguous(),batch_R)transformed_x = torch.cat([transformed_pts, transformed_nls], dim=-1)else:raise ValueErrorreturn batch_R, batch_t, transformed_xclass IterativeBenchmark(nn.Module):def __init__(self, in_dim, niters, gn):super(IterativeBenchmark, self).__init__()self.benckmark = Benchmark(gn=gn, in_dim1=in_dim)self.niters = nitersdef forward(self, x, y):transformed_xs = []device = x.deviceB = x.size()[0]transformed_x = torch.clone(x)batch_R_res = torch.eye(3).to(device).unsqueeze(0).repeat(B, 1, 1)batch_t_res = torch.zeros(3, 1).to(device).unsqueeze(0).repeat(B, 1, 1)for i in range(self.niters):batch_R, batch_t, transformed_x = self.benckmark(transformed_x, y)transformed_xs.append(transformed_x)batch_R_res = torch.matmul(batch_R, batch_R_res)batch_t_res = torch.matmul(batch_R, batch_t_res) \+ torch.unsqueeze(batch_t, -1)transformed_x = transformed_x.permute(0, 2, 1).contiguous()batch_t_res = torch.squeeze(batch_t_res, dim=-1)#transformed_x = transformed_x.permute(0, 2, 1).contiguous()return batch_R_res, batch_t_res, transformed_xs
fgr配准方法
import copy
import open3d as o3ddef fpfh(pcd, normals):pcd.normals = o3d.utility.Vector3dVector(normals)pcd_fpfh = o3d.registration.compute_fpfh_feature(pcd,o3d.geometry.KDTreeSearchParamHybrid(radius=0.3, max_nn=64))return pcd_fpfhdef execute_fast_global_registration(source, target, source_fpfh, target_fpfh):distance_threshold = 0.01result = o3d.registration.registration_fast_based_on_feature_matching(source, target, source_fpfh, target_fpfh,o3d.registration.FastGlobalRegistrationOption(maximum_correspondence_distance=distance_threshold))transformation = result.transformationestimate = copy.deepcopy(source)estimate.transform(transformation)R, t = transformation[:3, :3], transformation[:3, 3]return R, t, estimatedef fgr(source, target, src_normals, tgt_normals):source_fpfh = fpfh(source, src_normals)target_fpfh = fpfh(target, tgt_normals)R, t, estimate = execute_fast_global_registration(source=source,target=target,source_fpfh=source_fpfh,target_fpfh=target_fpfh)return R, t, estimate
ICP配准方法
import copy
import numpy as np
import open3d as o3ddef icp(source, target):max_correspondence_distance = 2 # 0.5 in RPM-Netinit = np.eye(4, dtype=np.float32)estimation_method = o3d.pipelines.registration.TransformationEstimationPointToPoint()reg_p2p = o3d.pipelines.registration.registration_icp(source=source,target=target,init=init,max_correspondence_distance=max_correspondence_distance,estimation_method=estimation_method)transformation = reg_p2p.transformationestimate = copy.deepcopy(source)estimate.transform(transformation)R, t = transformation[:3, :3], transformation[:3, 3]return R, t, estimate
3.基本使用
modelnet40数据集的评测及训练可以使用一下代码实现
# Iterative Benchmarkpython modelnet40_evaluate.py --root your_data_path/modelnet40_ply_hdf5_2048 --checkpoint your_ckpt_path/test_min_loss.pth --cuda# Visualization# python modelnet40_evaluate.py --root your_data_path/modelnet40_ply_hdf5_2048 --checkpoint your_ckpt_path/test_min_loss.pth --show# ICP# python modelnet40_evaluate.py --root your_data_path/modelnet40_ply_hdf5_2048 --method icp# FGR# python modelnet40_evaluate.py --root your_data_path/modelnet40_ply_hdf5_2048 --method fgr --normal
-
train
CUDA_VISIBLE_DEVICES=0 python modelnet40_train.py --root your_data_path/modelnet40_ply_hdf5_2048
这里注意讲述训练与评测自己的数据集,其中自己数据集的路径如下所示,里面都是处理好的pcd点云数据。
具体格式为:
|- CustomData(dir)|- train_data(dir)- train1.pcd- train2.pcd- ...|- val_data(dir)- val1.pcd- val2.pcd- ...
3.1 ICP方法性能评测
可以加上 --show 参数来查看每一个配准的数据
python custom_evaluate.py --root cumstom_data --infer_npts 2048 --method icp --normal
如果出现以下报错,则用open3d.pipelines.registration替换open3d.registration,具体可以用本博文的ICP配准方法替换掉models\icp.py中的内容
Traceback (most recent call last):File "custom_evaluate.py", line 142, in <module>evaluate_icp(args, test_loader)File "custom_evaluate.py", line 98, in evaluate_icpR, t, pred_ref_cloud = icp(npy2pcd(src_cloud), npy2pcd(ref_cloud))File "D:\点云AI配准\PCReg.PyTorch-main\models\icp.py", line 9, in icpestimation_method = o3d.registration.TransformationEstimationPointToPoint()
AttributeError: module 'open3d' has no attribute 'registration'
具体执行输出如下所示:
3.2 模型训练
训练命令:
python custom_train.py --root cumstom_data --train_npts 2048
训练好的模型保存在work_dirs\models\checkpoints目录中
评测命令:
python custom_evaluate.py --infer_npts 2048 --root cumstom_data --checkpoint work_dirs\models\checkpoints\test_min_loss.pth --show
其中绿色点云为源点云,红色点云为参考点云,蓝色点云为配准后的源点云,可以看到蓝色点云与红色点云完全没有对齐,这表明训练效果极其不佳。这或许是训练数据太少所导致的,毕竟本次实验只有18个点云数据。
4、原文效果
下图是作者论文中的配准效果图
在modelnet40数据集上相关精度信息如下所示,可以确定,本文方法与FGR方法相比没有显著性优势。
- Point-to-Point Correspondences(R error is large due to EMDLoss, see here)
Method | isotropic R | isotropic t | anisotropic R(mse, mae) | anisotropic t(mse, mae) | time(s) |
---|---|---|---|---|---|
ICP | 11.44 | 0.16 | 17.64(5.48) | 0.22(0.07) | 0.07 |
FGR | 0.01 | 0.00 | 0.07(0.00) | 0.00(0.00) | 0.19 |
IBenchmark | 5.68 | 0.07 | 9.77(2.69) | 0.12(0.03) | 0.022 |
IBenchmark + ICP | 3.65 | 0.04 | 9.22(1.66) | 0.11(0.02) |
- Noise Data(infer_npts = 1024)
Method | isotropic R | isotropic t | anisotropic R(mse, mae) | anisotropic t(mse, mae) |
---|---|---|---|---|
ICP | 12.14 | 0.17 | 18.32(5.86) | 0.23(0.08) |
FGR | 4.27 | 0.06 | 11.55(2.43) | 0.09(0.03) |
IBenchmark | 6.25 | 0.08 | 9.28(2.94) | 0.12(0.04) |
IBenchmark + ICP | 5.10 | 0.07 | 10.51(2.39) | 0.13(0.03) |
- Partial-to-Complete Registration(infer_npts = 1024)
Method | isotropic R | isotropic t | anisotropic R(mse, mae) | anisotropic t(mse, mae) |
---|---|---|---|---|
ICP | 21.33 | 0.32 | 22.83(10.51) | 0.31(0.15) |
FGR | 9.49 | 0.12 | 19.51(5.58) | 0.17(0.06) |
IBenchmark | 15.02 | 0.22 | 15.78(7.45) | 0.21(0.10) |
IBenchmark + ICP | 9.21 | 0.13 | 14.73(4.43) | 0.18(0.06) |