使用PCReg.PyTorch项目训练自己的数据集&进行点云配准

项目地址: 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目录,执行以下安装命令:

bash 复制代码
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

在编译过程中,很有可能碰到报错

bash 复制代码
  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_version
    raise 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操作。具体实例如下所示:

python 复制代码
    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.pyfgr.pyicp.py,分别为模型配准,fgr配准,icp配准方法。其中fgr配准与icp配准方法是使用open3d库实现。

benchmark

benchmark为本文模型,其是基于PointNet所实现的一个孪生网络,核心代码如Benchmark类所示。其基于encoder提取2个点云的特征,然后简单的使用全连接层将两个点云的特征进行交互,然后再输出两个点云的特征。在这里最为重要的是loss的设计,即如何设计优化目标,使模型参数以点云配准为优化方向

python 复制代码
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_dim1
        self.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_dim

    def 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 ValueError
        return batch_R, batch_t, transformed_x


class 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 = niters

    def forward(self, x, y):
        transformed_xs = []
        device = x.device
        B = 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配准方法
python 复制代码
import copy
import open3d as o3d


def 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_fpfh


def execute_fast_global_registration(source, target, source_fpfh, target_fpfh):
    distance_threshold = 0.01
    result = o3d.registration.registration_fast_based_on_feature_matching(
        source, target, source_fpfh, target_fpfh,
        o3d.registration.FastGlobalRegistrationOption(
            maximum_correspondence_distance=distance_threshold))
    transformation = result.transformation
    estimate = copy.deepcopy(source)
    estimate.transform(transformation)
    R, t = transformation[:3, :3], transformation[:3, 3]
    return R, t, estimate


def 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配准方法
python 复制代码
import copy
import numpy as np
import open3d as o3d


def icp(source, target):
    max_correspondence_distance = 2 # 0.5 in RPM-Net
    init = 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.transformation
    estimate = copy.deepcopy(source)
    estimate.transform(transformation)
    R, t = transformation[:3, :3], transformation[:3, 3]
    return R, t, estimate

3.基本使用

modelnet40数据集的评测及训练可以使用一下代码实现

    # Iterative Benchmark
    python 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点云数据。

具体格式为:

txt 复制代码
    |- CustomData(dir)
        |- train_data(dir)
            - train1.pcd
            - train2.pcd
            - ...
        |- val_data(dir)
            - val1.pcd
            - val2.pcd
            - ...

3.1 ICP方法性能评测

可以加上 --show 参数来查看每一个配准的数据

python 复制代码
python custom_evaluate.py --root cumstom_data --infer_npts 2048  --method icp --normal

如果出现以下报错,则用open3d.pipelines.registration替换open3d.registration,具体可以用本博文的ICP配准方法替换掉models\icp.py中的内容

bash 复制代码
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_icp
    R, t, pred_ref_cloud = icp(npy2pcd(src_cloud), npy2pcd(ref_cloud))
  File "D:\点云AI配准\PCReg.PyTorch-main\models\icp.py", line 9, in icp
    estimation_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)
相关推荐
C7211BA1 小时前
基于网页的大语言模型聊天机器人
人工智能·语言模型·机器人
铭瑾熙2 小时前
深度学习之人脸检测
人工智能·深度学习
白光白光3 小时前
量子卷积神经网络
人工智能·神经网络·cnn
陈苏同学4 小时前
机器翻译 & 数据集 (NLP基础 - 预处理 → tokenize → 词表 → 截断/填充 → 迭代器) + 代码实现 —— 笔记3.9《动手学深度学习》
人工智能·pytorch·笔记·python·深度学习·自然语言处理·机器翻译
狂放不羁霸4 小时前
组会 | 大语言模型 + LoRA
人工智能·语言模型·自然语言处理
sp_fyf_20244 小时前
【大语言模型】ACL2024论文-20 SCIMON:面向新颖性的科学启示机器优化
人工智能·深度学习·机器学习·语言模型·自然语言处理·数据挖掘
宋138102797204 小时前
SouVR Feedback force7 力反馈设备
人工智能·机器人·vr
我叫白小猿5 小时前
【大模型-智能体】AutoGen Studio测试和导出工作流程
人工智能·python·workflow·工作流·智能体·autogen
CopyLower5 小时前
AI赋能电商:智能购物推荐、会员分类与商品定价的创新探索
人工智能·分类·数据挖掘
界面开发小八哥5 小时前
界面控件DevExpress WinForms v24.2新功能预览 - 人工智能(AI)
人工智能·.net·界面控件·devexpress·ui开发