基于PointNet的迭代方法

文章来源

知乎：https://zhuanlan.zhihu.com/p/289620126

Github：https://github.com/zhulf0804/PCReg.PyTorch

源码分析

源码来自Github

目录结构

PCReg.PyTorch

• data

  • CustomDataset.py # 自定义模型的Dataset定义

  • ModelNet40.py # ModelNet模型的Dataset定义

• loss

  • earth_mover_distance # 推土机距离损失定义

• matrics

  • metrics.py # 计算评价指标

  • helper.py # 计算R和t的误差

• models

  • benchmark.py # 网络结构定义

  • fgr.py # Fast Global Registration (FGR) 算法执行点云之间的配准

  • icp.py # ICP算法执行点云之间的配准

• utils

  • dist.py # 计算的是两组点之间点对点的平方距离

  • format.py # 读取pcd、npy和pcd文件相互转换

  • time.py # 一个装饰器，计算传入函数的执行时间

• custom_evaluate.py # 评估ICP和基于迭代的网络后结果的性能

• custom_train.py # 进行配准任务的训练

• modelnet40_evaluate.py # 同上，只是模型换成了modelnet40

• modelnet40_train.py # 同上

数据处理

from utils import readpcd
from utils import pc_normalize, random_select_points, shift_point_cloud, \
    jitter_point_cloud, generate_random_rotation_matrix, \
    generate_random_tranlation_vector, transform


class CustomData(Dataset):
    def __init__(self, root, npts, train=True):
        super(CustomData, self).__init__()
        # 根据train参数选择训练集或验证集的目录名
        dirname = 'train_data' if train else 'val_data'
        path = os.path.join(root, dirname)
        self.train = train
        # 获取所有文件路径
        self.files = [os.path.join(path, item) for item in sorted(os.listdir(path))]
        self.npts = npts

    def __getitem__(self, item):
        # 获取指定索引的数据项
        file = self.files[item]
        ref_cloud = readpcd(file, rtype='npy')  # 加载pcd文件
        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

    def __len__(self):
        return len(self.files)

1. 初始化：CustomData 类继承自 PyTorch 的 Dataset 类，接受数据集根目录、采样点数和指示是否为训练集的参数。根据参数选择训练集或验证集的目录名，读取该目录下的所有点云文件路径。

2. 获取数据项：__getitem__ 方法用于按索引获取指定数据项。首先从文件列表中获取对应索引的点云文件，然后加载该点云数据。

3. 随机选点：对点云数据执行随机选点操作，将点云数据减少到指定的点数。这有助于减小计算复杂度并提高处理速度。

4. 归一化：将点云数据进行归一化处理，将其坐标范围限制在一个固定的范围内。这有助于提高模型的泛化性能和收敛速度。

5. 生成随机变换矩阵：生成一个随机的旋转矩阵和平移向量，用于模拟真实场景中点云之间的相对位置变化。旋转角度范围为 -20 到 20 度，平移范围为 -0.5 到 0.5。

6. 应用变换：将生成的随机旋转矩阵和平移向量应用到参考点云数据，得到源点云数据。

7. 添加随机噪声：如果是训练集，为参考点云和源点云添加随机噪声。这有助于提高模型的鲁棒性，使其能够在噪声环境下更好地执行点云配准。

8. 返回数据：返回预处理后的参考点云、源点云、旋转矩阵和平移向量。这些数据将作为网络的输入和监督信号进行训练。

9. 数据集大小：__len__ 方法返回数据集中点云文件的数量。

网络定义

Encoder (PointNet)

class PointNet(nn.Module):  # 用于提取点云特征
    def __init__(self, in_dim, gn, mlps=[64, 64, 64, 128, 1024]):
        super(PointNet, self).__init__()
        self.backbone = nn.Sequential()
        # 初始化网络的骨架（backbone），它是一个 Sequential 容器，用于按顺序存放神经网络的层
        # 遍历 mlps 列表中的每个输出维度（out_dim）
        for i, out_dim in enumerate(mlps):
            # 向骨架中添加 1D 卷积层，输入维度为 in_dim，输出维度为 out_dim
            self.backbone.add_module(f'pointnet_conv_{i}',
                                     nn.Conv1d(in_dim, out_dim, 1, 1, 0))
            # 如果 gn（GroupNorm）为 True，则添加 GroupNorm 层
            if gn:
                self.backbone.add_module(f'pointnet_gn_{i}',
                                    nn.GroupNorm(8, out_dim))
                # 向骨架中添加 ReLU 激活函数
            self.backbone.add_module(f'pointnet_relu_{i}',
                                     nn.ReLU(inplace=True))
            in_dim = out_dim

    def forward(self, x):
        x = self.backbone(x)
        x, _ = torch.max(x, dim=2)
        return x

这个网络用于提取点云数据的特征，为一个1024维的特征

1. 输入维度：in_dim (例如 3，表示位置信息)

2. 第一个卷积层输出维度：64

   • ReLU 激活函数

3. 第二个卷积层输出维度：64

   • ReLU 激活函数

4. 第三个卷积层输出维度：64

   • ReLU 激活函数

5. 第四个卷积层输出维度：128

   • ReLU 激活函数

6. 第五个卷积层输出维度：1024

   • ReLU 激活函数

7. 全局最大池化（沿 num_points 维度）：维度变为 (batch_size, 1024)

在每个卷积层之后，我们都添加了一个 ReLU 激活函数，以引入非线性，使得网络能够学习更复杂的特征表示。这种网络结构可以更有效地捕捉点云数据中的信息。在全局最大池化之后，我们得到一个尺寸为 (batch_size, 1024) 的特征向量，可以用于后续任务。

Decoder (Benchmark)

class Benchmark(nn.Module):  # 这个名字其实不是很恰当，PointCloudAlignment
    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)  # 编码：PointNet提取特征
        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:  # 是否添加GroupNorm层
                    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 和 y 进行特征提取
        x_f, y_f = self.encoder(x), self.encoder(y)
        # 沿着维度 1 对特征 x_f 和 y_f 进行拼接
        concat = torch.cat((x_f, y_f), dim=1)
        # 将拼接后的特征传入解码器
        out = self.decoder(concat)
        # 对输出结果进行分割，得到平移向量（batch_t）和四元数（batch_quat）
        batch_t, batch_quat = out[:, :3], out[:, 3:] / torch.norm(out[:, 3:], dim=1, keepdim=True)
        # 将四元数转换为旋转矩阵（batch_R）
        batch_R = batch_quat2mat(batch_quat)
        # 根据输入的维度（in_dim1）选择不同的变换方式
        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

1. 输入点云 x 和 y 的维度：(batch_size, in_dim1, num_points)。

2. 使用 PointNet 编码器对输入 x 和 y 进行特征提取：

   • 输出特征维度：(batch_size, 1024)。

3. 将特征 x_f 和 y_f 进行拼接：

   • 拼接后的特征维度：(batch_size, 2048)。

4. 将拼接后的特征传入解码器。解码器包含一系列全连接层和 ReLU 激活函数：

   • 第一个全连接层输出维度：1024

     • ReLU 激活函数

   • 第二个全连接层输出维度：1024

     • ReLU 激活函数

   • 第三个全连接层输出维度：512

     • ReLU 激活函数

   • 第四个全连接层输出维度：512

     • ReLU 激活函数

   • 第五个全连接层输出维度：256

     • ReLU 激活函数

   • 第六个全连接层输出维度：7（输出旋转和平移参数）

5. 输出：(batch_size, 7)，其中前 3 个元素表示平移向量，后 4 个元素表示旋转四元数。

Benchmark 网络首先使用 PointNet 编码器对输入点云 x 和 y 提取特征，然后将两个特征向量拼接在一起。接下来，通过一个解码器（由一系列全连接层和 ReLU 激活函数组成）将拼接后的特征映射到旋转和平移参数。最后，网络输出一个大小为 (batch_size, 7) 的向量，表示点云之间的相对变换。

IterativeNet (IterativeBenchmark)

class IterativeBenchmark(nn.Module):  # 迭代
    def __init__(self, in_dim, niters, gn):
        super(IterativeBenchmark, self).__init__()
        # 初始化 Benchmark 模型
        self.benckmark = Benchmark(gn=gn, in_dim1=in_dim)
        # 设置迭代次数
        self.niters = niters

    def forward(self, x, y):
        # 用于存储每次迭代后变换的 x
        transformed_xs = []
        device = x.device
        B = x.size()[0]
        # 创建一个与 x 相同的张量 transformed_x
        transformed_x = torch.clone(x)
        # 初始化旋转矩阵为单位矩阵，大小为 Bx3x3
        batch_R_res = torch.eye(3).to(device).unsqueeze(0).repeat(B, 1, 1)
        # 初始化平移向量为零矩阵，大小为 Bx3x1
        batch_t_res = torch.zeros(3, 1).to(device).unsqueeze(0).repeat(B, 1, 1)
        # 进行 self.niters 次迭代
        for i in range(self.niters):
            # 对 transformed_x 和 y 使用 Benchmark 模型
            batch_R, batch_t, transformed_x = self.benckmark(transformed_x, y)
            # 将本次迭代的变换后的 x 添加到列表中
            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 = 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_xsn

一个基于迭代方法的点云对其网络，用Pointnet作为编码器，Benchmark作为解码器，多次迭代来优化点云之间的变换

以下是 IterativeBenchmark 网络的各层和维度变化：

1. 输入点云 x 和 y 的维度：(batch_size, in_dim, num_points)。

2. 初始化 Benchmark 网络。

3. 设置迭代次数 niters。

4. 初始化变换后的点云 transformed_x 为与输入点云 x 相同。

5. 初始化旋转矩阵为单位矩阵，大小为 (batch_size, 3, 3)。

6. 初始化平移向量为零矩阵，大小为 (batch_size, 3, 1)。

7. 对于 i 从 0 到 niters-1 进行迭代：

   • 使用 Benchmark 模型对 transformed_x 和 y 进行特征提取，计算旋转矩阵 batch_R 和平移向量 batch_t，并获得变换后的点云 transformed_x。

   • 更新旋转矩阵和平移向量。

   • 调整变换后的点云 transformed_x 的维度顺序。

8. 输出：最终的旋转矩阵 batch_R_res，最终的平移向量 batch_t_res 和每次迭代后变换的点云 transformed_xs。

IterativeBenchmark 网络首先初始化一个 Benchmark 网络，然后对输入点云 x 和 y 进行多次迭代。在每次迭代中，网络计算 transformed_x 和 y 之间的相对变换，并更新旋转矩阵和平移向量。经过 niters 次迭代后，网络输出最终的旋转矩阵、平移向量以及每次迭代后的变换点云。这种迭代策略有助于提高网络的性能和鲁棒性。