[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features

零、概要

  • 论文: RPM-Net: Robust Point Matching using Learned Features
  • tag: CVPR 2020; registration
  • 代码: https://github.com/yewzijian/RPMNet/
  • 作者: Zi Jian Yew, Gim Hee Lee
  • 机构: Department of Computer Science, National University of Singapore
  • 笔者整理了一个最近几年150多篇点云的论文列表,欢迎大家一块学习交流。

0.0 摘要

ICP算法通过迭代的方式解决点云的刚性配准: (1) 基于最近空间距离寻找hard的点对应关系; (2) 求解最小二乘变换。其中第(1)步对待配准点云的初始位置和噪声点/离群点敏感,会导致迭代收敛到错误的局部最小值。在本文中,作者提出了RPM-Net,一种对初始化不太敏感且鲁棒的、基于学习的用于刚性点云配准方法。为此,RMP-Net使用可微的Sinkhorn层和退火(Annealing)技术,利用从空间坐标和局部几何信息中学习到的混合特征,实现点对应(point correspondences)的软(soft)分配。为了提高配准性能,作者引入了第二个网络来预测最佳退火参数。与现有的一些方法不同,RPM-Net可以处理部分可见和丢失correspondences的点云。实验结果表明,与现有的非深度学习和最近的基于深度学习的方法相比,RPM-Net达到了SOTA。

读到这里,我们会产生如下的疑问:

  • RPM-Net怎么解决部分可见和丢失correspondences的点云配准 ?
  • Sinkhorn和退火是什么,它们如何产生点对应的软分配的 ?

上述问题将会在二、论文的方法中介绍。

一、论文的出发点和贡献

现有点云配准的方案主要包括以下三种:

  • 基于距离的匹配算法: ICP, RPM等。这类算法对待配准点云的初始位置或噪声敏感。
  • 基于特征的匹配算法: 基于特征(LRF、FPH、FPFH等)建立point correspondences。这类算法要求点云具有独特的几何结构。
  • 基于学习的匹配算法: PointNetLK, DCP, PRNet等。上述算法或者不能解决局部-局部点云的匹配,或者设计太复杂。

针对上述点云配准中的现状,作者提出了RPM-Net,它是一个end-to-end的可微分深度网络,它保留了RPM对噪声/离群点的鲁棒性,同时从学习的特征距离(而不是空间距离)中预测点对应(point correspondences)关系,降低了配准对初始值的依赖。

它和ICP算法类似,也是一个迭代配准算法。由于使用了混合特征(空间特征和集合特征),算法可以在少量迭代中收敛。

在我看来,作者的主要贡献在于: (1)引进了match matrix,以及如何生成match matrix; (2)提出了一种改进的Chamfer distance度量,以提高点云局部可见情况下的配准质量。

二、论文的方法

RPM-Net的架构如Figure 2所示, 它是一个迭代网络。在每次迭代中,Source点云X和Reference点云Y会经过特征提取模块,参数预测模块,计算Match Matrix模块和SVD分解模块。当前迭代求解的R, t,会作用于当前的source点云X,并将得到的点云与Reference点云Y作为输入传递到下一次迭代中。

[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features_第1张图片

2.1 特征提取模块

对于点云中的一个点 x c \text{x}_c xc,定义它的邻域 N ( x c ) N(\text{x}_c) N(xc),则

F x c = f θ ( x c , { Δ x c , i } , { P P F ( x c , x i ) } ) F_{\text{x}_c} = f_\theta(\text{x}_c, \lbrace \Delta \text{x}_{c, i} \rbrace, \lbrace PPF(\text{x}_c, \text{x}_i) \rbrace) Fxc=fθ(xc,{Δxc,i},{PPF(xc,xi)})

f θ f_\theta fθ是一个神经网络,待学习的参数是 θ \theta θ x i ∈ N ( x c ) \text{x}_i \in N(\text{x}_c) xiN(xc) Δ x c , i = x i − x c \Delta \text{x}_{c, i} = \text{x}_i - \text{x}_c Δxc,i=xixc P P F ( x c , x i ) PPF(\text{x}_c, \text{x}_i) PPF(xc,xi)是一个4D的特征,其组成如下:

P P F ( x c , x i ) = ( ∠ ( n c , Δ x c , i ) , ∠ ( n i , Δ x c , i ) , ∠ ( n c , n i ) , ∣ ∣ Δ x c , i ∣ ∣ 2 ) PPF(\text{x}_c, \text{x}_i) = (\angle (\text{n}_c, \Delta \text{x}_{c, i}), \angle(\text{n}_i, \Delta \text{x}_{c, i}), \angle (\text{n}_c, \text{n}_i), ||\Delta \text{x}_{c, i}||_2) PPF(xc,xi)=((nc,Δxc,i),(ni,Δxc,i),(nc,ni),Δxc,i2)

∠ \angle 表示角度angle(v1, v2) = atan2(cross(v1, v2), dot(v1, v2))。PPF特征是一个四维向量,因此对于特定的 i i i, ( x c , { Δ x c , i } , { P P F ( x c , x i ) } ) (\text{x}_c, \lbrace \Delta \text{x}_{c, i} \rbrace, \lbrace PPF(\text{x}_c, \text{x}_i) \rbrace) (xc,{Δxc,i},{PPF(xc,xi)})是一个10维向量。

特征提取模块如Figure 2(b)所示,对于 x c \text{x}_c xc的每个邻域点 x i \text{x}_i xi,构成一个10维向量,batch的数据格式为(B, 10, N), 送入到一个给定共享参数的MLP,得到维度为(B, C1, N)的特征, 而后经过Max Pooling,得到维度为(B, C1)的特征,接着经过几个MLP,得到维度为(B, C2)的特征。最后对特征进行归一化,产生维度为(B, C2)的特征 F x c F_{\text{x}_c} Fxc

2.2 参数预测模块

在原始的RPM算法中,参数 α , β \alpha, \beta α,β依赖于数据集,需要人工调制。在RMP-Net中,这些参数由于依赖于学习到的特征,很难调制。因此,作者设计了一个网络来预测这两个参数。如Figure 2©所示,网络的输入来自source点云X(J, 3)和reference点云Y(K, 3),组织形式为concat, 构成维度为(J + K, 3)的输入。另外为了表征当前点来自于哪个数据集,增加了第4列特征,0表示当前点来自于X点云, 1表示当前点来自于Y点云。所以,参数预测模块的输入数据是(J+K, 4)维的。对于(B, 4, J + K)的数据首先经过共享参数的MLP,产生维度为(B, C3, J+K)的数据,经过Max Pooling,得到维度为(B, C3)的数据,接下来经过MLP得到(B, 2)的数据。为了确保预测的 α , β \alpha, \beta α,β是正的,作者在最后使用了softplux激活函数。

2.3 计算Match Matrix

经过上述两个模块,得到了两个特征矩阵 F X ∈ R J × C , F Y ∈ R K × C F_X \in \mathbb R^{J \times C}, F_Y \in \mathbb R^{K \times C} FXRJ×C,FYRK×C和退火参数 α , β \alpha, \beta α,β。接下来如何找点对应(point correspondences)关系呢 ? 这里是通过Match Matrix M ∈ R J × K M \in \mathbb R^{J \times K} MRJ×K实现的。 M M M的初始值 m j k m_{jk} mjk通过如下方式获得:

m j k = e − β ( ∣ ∣ F x j − F y k ∣ ∣ 2 − α ) m_{jk} = e^{-\beta(||F_{\text{x}_j} - F_{\text{y}_k}||_2 - \alpha)} mjk=eβ(FxjFyk2α)

接下来是论文的一个重点,如Figure 2(d)所示, 通过Sinkhorn迭代来获取一个双随机限制的矩阵(doubly stochastic matrix)。相关的理论证明在A relationship between arbitrary positive matrices and doubly stochastic matrices。这里把Sinkhorn相关代码搬过来吧。它是一个迭代算法,大致做法交替进行行归一化和列归一化。

def sinkhorn(log_alpha, n_iters: int = 5, slack: bool = True, eps: float = -1) -> torch.Tensor:
    """ Run sinkhorn iterations to generate a near doubly stochastic matrix, where each row or column sum to <=1

    Args:
        log_alpha: log of positive matrix to apply sinkhorn normalization (B, J, K)
        n_iters (int): Number of normalization iterations
        slack (bool): Whether to include slack row and column
        eps: eps for early termination (Used only for handcrafted RPM). Set to negative to disable.

    Returns:
        log(perm_matrix): Doubly stochastic matrix (B, J, K)

    Modified from original source taken from:
        Learning Latent Permutations with Gumbel-Sinkhorn Networks
        https://github.com/HeddaCohenIndelman/Learning-Gumbel-Sinkhorn-Permutations-w-Pytorch
    """

    # Sinkhorn iterations
    prev_alpha = None
    if slack:
        zero_pad = nn.ZeroPad2d((0, 1, 0, 1))
        log_alpha_padded = zero_pad(log_alpha[:, None, :, :])

        log_alpha_padded = torch.squeeze(log_alpha_padded, dim=1)

        for i in range(n_iters):
            # Row normalization
            log_alpha_padded = torch.cat((
                    log_alpha_padded[:, :-1, :] - (torch.logsumexp(log_alpha_padded[:, :-1, :], dim=2, keepdim=True)),
                    log_alpha_padded[:, -1, None, :]),  # Don't normalize last row
                dim=1)

            # Column normalization
            log_alpha_padded = torch.cat((
                    log_alpha_padded[:, :, :-1] - (torch.logsumexp(log_alpha_padded[:, :, :-1], dim=1, keepdim=True)),
                    log_alpha_padded[:, :, -1, None]),  # Don't normalize last column
                dim=2)

            if eps > 0:
                if prev_alpha is not None:
                    abs_dev = torch.abs(torch.exp(log_alpha_padded[:, :-1, :-1]) - prev_alpha)
                    if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
                        break
                prev_alpha = torch.exp(log_alpha_padded[:, :-1, :-1]).clone()

        log_alpha = log_alpha_padded[:, :-1, :-1]
    else:
        for i in range(n_iters):
            # Row normalization (i.e. each row sum to 1)
            log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=2, keepdim=True))

            # Column normalization (i.e. each column sum to 1)
            log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=1, keepdim=True))

            if eps > 0:
                if prev_alpha is not None:
                    abs_dev = torch.abs(torch.exp(log_alpha) - prev_alpha)
                    if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
                        break
                prev_alpha = torch.exp(log_alpha).clone()

    return log_alpha

这里就回答了Sinkhorn和退火是什么, 如何产生point soft correspondence将在2.4小节介绍

先解释一下RPM-Net怎么解决部分可见和丢失correspondences的点云配准的吧。

因为X中的每一个点 x j \text{x}_j xj并不总是在Y中存在一个对应点 y k \text{y}_k yk,因此 M M M矩阵的第 j j j行之和就是0,不满足doubly stochastic 矩阵的性质。因此,作者在 M M M矩阵最右边增加了一列,这样就可以满足doubly stochastic的性质了;同理在 M M M矩阵的最下面增加了一行。这样就产生了 M ∈ R ( J + 1 ) × ( K + 1 ) M \in \mathbb R^{(J+1) \times (K + 1)} MR(J+1)×(K+1)的矩阵了。在行归一化时,不归一最后一行,在列归一化时不归一化最后一列。结束后,取前 J J J行, 前 M M M列,带入下面的计算。

2.4 计算R, t

对于点云X中每一个点 x j \text{x}_j xj,构造响应的对应点:

y j = 1 Σ k = 1 K m j , k Σ k = 1 K m j , k y k \text{y}_j = \frac{1}{\Sigma_{k=1}^Km_{j,k}}\Sigma_{k=1}^Km_{j, k}\text{y}_k yj=Σk=1Kmj,k1Σk=1Kmj,kyk

另外,因为并不是每一个 x j \text{x}_j xj都有一个对应点,因此在每一个对应关系 ( x j , y i ) (\text{x}_j, \text{y}_i) (xj,yi)设置一个权重 w j = Σ k = 1 K m j , k w_j = \Sigma_{k=1}^Km_{j,k} wj=Σk=1Kmj,k。最后通过SVD求解R, t。

对每一个对应关系 ( x j , y i ) (\text{x}_j, \text{y}_i) (xj,yi)设置一个权重 w j w_j wj是否由必要呢 ?

应该是必要的,论文中对每一个 x j \text{x}_j xj构造了soft correspondence, 虽然soft correspondence是根据权重生成的,但它仍然是不准确的(比如, 权重为0, 生成了(0, 0, 0),此时的对应关系就是错误的),因此在最后的对应关系有必要增加一个权重,来惩罚不同置信度下correspondence所带来的损失。

2.5 损失函数

L reg = 1 J Σ j J ∣ ( R gt x j + t gt ) − ( R pred x j + t pred ) ∣ L_{\text{reg}} = \frac{1}{J}\Sigma_j^J |(R_{\text{gt}}\text{x}_j + t_{\text{gt}}) - (R_{\text{pred}}\text{x}_j + t_{\text{pred}})| Lreg=J1ΣjJ(Rgtxj+tgt)(Rpredxj+tpred)

这个损失函数感觉还是挺有创意的,既考虑了R, t,还考虑点云X中的所有点。

作者说,只使用上述loss,网络倾向于把大部分点标记为outliers。为此,作者增加了第2个loss,鼓励 M M M得到更大的inliers:

L inlier = − 1 J Σ j J Σ k K m j k − 1 K Σ k K Σ j J m j k L_{\text{inlier}} = -\frac{1}{J}\Sigma_j^J\Sigma_k^Km_{jk} - \frac{1}{K}\Sigma_k^K\Sigma_j^Jm_{jk} Linlier=J1ΣjJΣkKmjkK1ΣkKΣjJmjk

整体的loss为两者的和:

L total = L reg + λ L inlier L_{\text{total}} = L_{\text{reg}} + \lambda L_{\text{inlier}} Ltotal=Lreg+λLinlier

因为RPM-Net是一个迭代的网络,因此考虑每次迭代的损失,作者这里把损失设置为每次迭代损失的加权和:

L = Σ i N ( 1 2 ) ( N − i ) L total i L = \Sigma_i^N(\frac{1}{2})^{(N - i)}L_{\text{total}}^i L=ΣiN(21)(Ni)Ltotali

还有一个需要注意的地方: 每次迭代都是独立的,都是学习GroundTruth的R,t。因此,作者在训练时采用了2次迭代,在测试时采用了5次迭代。这个和以往的一些算法,如PRNet是不同的。

三、论文的实验

作者在ModelNet40上进行了实验,使用前20类进行训练和验证(调参),使用其余20类进行测试。训练集有5112个,验证集有1202个,测试集有1266个。在训练和测试时,旋转角度在[0, 45°]进行采样,平移从[-0.5, 0.5]进行采样。

3.1评价指标

  • anisotropic

    旋转矩阵和平移向量的MSE和MAE值,在DCP和PRNet中都有使用。

  • isotropic
    Error ( R ) = arc cos ( tr ( R GT − 1 R ) − 1 2 ) , Error ( t ) = ∣ ∣ t GT − t ∣ ∣ 2 \text{Error}(R) = \text{arc cos}(\frac{\text{tr}(R_{\text{GT}}^{-1}R) - 1}{2}), \text{Error}(t) = ||t_{\text{GT}} - t||_2 Error(R)=arc cos(2tr(RGT1R)1),Error(t)=tGTt2

    此方式在3DRegNet论文中有使用。

  • Chamfer distance

    anisotropic和isotropic的评价都是在R,t上面, Chamfer distance有些不同,它考虑了样本。在计算最近距离时,考虑了原始点云,并不是只比较ref和src之间的距离。

    C D ( X , Y ) = 1 ∣ X ∣ Σ x ∈ X min ⁡ y ∈ Y clean ∣ ∣ x − y ∣ ∣ 2 + 1 ∣ Y ∣ Σ y ∈ Y min ⁡ x ∈ X clean ∣ ∣ x − y ∣ ∣ 2 CD(X, Y) = \frac{1}{|X|}\Sigma_{x \in X} \min_{y \in Y_{\text{clean}}}||x - y||_2 + \frac{1}{|Y|}\Sigma_{y \in Y} \min_{x \in X_{\text{clean}}}||x - y||_2 CD(X,Y)=X1ΣxXyYcleanminxy2+Y1ΣyYxXcleanminxy2

作者在Clean, Gaussian Noise, Partial Visibility数据集上进行了实验,并与ICP, RPM, FGR, PointNetLK, DCP-v2算法进行了比较,可视化结果如Figure 3所示。我对这部分的实验细节比较感兴趣,所以记录一下作者是怎么处理这些实验的(大部分是结合作者开源的代码进行描述)

3.2 Clean Data

  • 训练集: 对于单个点云S(2048, ), 随机采样1024个点, 生成src和ref点云, 这里src和ref是相同的点云。随机生成R, t, 将R, t作用于src得到新的点云src_transformed。随机打乱src_transformed和ref点云中点的顺序,并分别记这两个点云为src, ref。此时训练的数据和Gt就是src, ref和R, t。
  • 测试集: 对于单个点云S(2048, ), 固定采样1024个点, 生成src和ref点云, 这里src和ref是相同的点云。接下来的过程和训练集一致。

实验结果如Table 2所示,RPM-Net和FGR处于第一或者第二的位置。在接下来的章节中,可以看懂FGRd对噪声非常敏感。

[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features_第2张图片

3.3 Gaussian Noise

  • 训练集: 对于单个点云S(2048, ),首先生成src和ref点云, 这两个点云是一致的且都包含2048个点。随机生成R, t,将R, t作用于src得到新的点云src_transformed。从src_transformed和ref点云随机采样1024个点生成新的点云, 得到src_transformed_sampled和ref_sampled, 此时得到的src_transformed_sampled和ref_sampled不具有一一对应关系的。接下来对src_transformed_sampled和ref_sampled进行随机抖动和随机打乱顺序,生成新的点云,分别即为src, ref。
  • 测试集: 生成过程和训练集一样,只是设定了随机种子。

实验结果如Table 2所示,RPM-Net优于其他方法。

[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features_第3张图片

3.4 Partial Visibility

  • 训练集: 对于单个点云S(2048, ),首先生成src和ref点云, 这两个点云是一致的且都包含2048个点。对src和ref进行partial visibility(这里通过随机一个方向的半空间,移动确保剩余的点在原始点云的70%左右)点云生成新的点云src_crop和ref_crop,这两个点云是不一样的。接下来就是随机生成R,t等操作,和Gaussian Noise的数据处理方式相同。
  • 测试集: 成过程和训练集一样,只是设定了随机种子。

实验结果如Table 3所示, RPM-Net明显优于其它方法。

[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features_第4张图片

3.5 Ablation Studies

作者在partial visibility实验上进行了ablation studies,来更好的了解各种选择是如何影响算法的性能的,实验结果如Table 4所示。比较1-3行和第5行,可以看到x, Δ x \Delta \text x Δx, PPF都是必须的,且x对网络的性能影响较大。第4行,把使用复杂网络预测 α , β \alpha, \beta α,β改成了学习 α , β \alpha, \beta α,β,相比于第5行,性能也降低了。

[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features_第5张图片

四、对论文的想法

  • 缺点
    • 只在ModelNet40上进行了实验,不具有说服力。
    • 只比较了基于学习方法的PointNetLK和DCP(局部点云匹配不佳),比较不充分。
  • 优点
    • RPM-Net是一个迭代网络,在设计损失函数和累加每次迭代损失时具有创新性;而且,每次迭代都是独立的,都是学习GroundTruth的R,t,这一点与以往的算法是不同的。
    • 引入传统的RPM和SkinHorn的思想,设计基于深度学习的配准方法。

五、RPM-Net预测可视化代码

在src文件夹中新建vis.py,内容如下:

"""Evaluate RPMNet. Also contains functionality to compute evaluation metrics given transforms

Example Usages:
    1. Visualize RPMNet
        python vis.py --noise_type crop --resume [path-to-model.pth] --dataset_path [your_path]/modelnet40_ply_hdf5_2048

"""
import os
import open3d as o3d
import random
from tqdm import tqdm
import torch

from arguments import rpmnet_eval_arguments
from common.misc import prepare_logger
from common.torch import dict_all_to_device, CheckPointManager, to_numpy
from common.math_torch import se3
from data_loader.datasets import get_test_datasets
import models.rpmnet


def vis(npys):
    pcds = []
    colors = [[1.0, 0, 0],
              [0, 1.0, 0],
              [0, 0, 1.0]]
    for ind, npy in enumerate(npys):
        color = colors[ind] if ind < 3 else [random.random() for _ in range(3)]
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(npy)
        pcd.paint_uniform_color(color)
        pcds.append(pcd)
    return pcds


def inference_vis(data_loader, model: torch.nn.Module):
    _logger.info('Starting inference...')
    model.eval()

    with torch.no_grad():
        for data in tqdm(data_loader):

            #gt_transforms = data['transform_gt']
            points_src = data['points_src'][..., :3]
            points_ref = data['points_ref'][..., :3]
            #points_raw = data['points_raw'][..., :3]
            dict_all_to_device(data, _device)
            pred_transforms, endpoints = model(data, _args.num_reg_iter)
            src_transformed = se3.transform(pred_transforms[-1], points_src)

            src_np = torch.squeeze(points_src).cpu().detach()
            src_transformed_np = torch.squeeze(src_transformed).cpu().detach()
            ref_np = torch.squeeze(points_ref).cpu().detach()

            pcds = vis([src_np, src_transformed_np, ref_np])
            o3d.visualization.draw_geometries(pcds)


def get_model():
    _logger.info('Computing transforms using {}'.format(_args.method))
    assert _args.resume is not None
    model = models.rpmnet.get_model(_args)
    model.to(_device)
    if _device == torch.device('cpu'):
        model.load_state_dict(
            torch.load(_args.resume, map_location=torch.device('cpu'))['state_dict'])
    else:
        model.load_state_dict(torch.load(_args.resume)['state_dict'])
    return model


def main():
    # Load data_loader
    test_dataset = get_test_datasets(_args)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=1, shuffle=False)

    model = get_model()
    inference_vis(test_loader, model)  # Feedforward transforms

    _logger.info('Finished')


if __name__ == '__main__':
    # Arguments and logging
    parser = rpmnet_eval_arguments()
    _args = parser.parse_args()
    _logger, _log_path = prepare_logger(_args, log_path=_args.eval_save_path)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu)
    if _args.gpu >= 0 and (_args.method == 'rpm' or _args.method == 'rpmnet'):
        os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu)
        _device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
    else:
        _device = torch.device('cpu')

    main()

部分-部分匹配的可视化结果如下:

[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features_第6张图片
红色为src点云,蓝色为ref点云,绿色为RPM-Net输出R,t作用到src的点云。

使用方法:

python vis.py --noise_type crop --resume [path-to-model.pth] --dataset_path [your_path]/modelnet40_ply_hdf5_2048

如,
# python vis.py --noise_type crop --resume partial-trained.pth --dataset_path /Users/zhulf/data/modelnet40_ply_hdf5_2048

你可能感兴趣的:(点云,点云配准,CVPR2020)