ICLR 2023｜场景三维重建新SOTA！基于3D Transformer的单目场景重建

2023-05-11 12:05

©作者 | 原力

单位 | 阿里巴巴达摩院

研究方向 | 三维重建

从单目序列图像进行场景三维重建是一个很具有挑战性的任务，最近有很多 volume-based 方法学习直接预测 TSDF volume 来实现三维重建，展现了非常不错的结果，相比于之前的 depth-based 方法可以重建出更加完整、一致、连贯的场景几何。然而，大多数这类方法都是在研究如何提取 2D 特征、如何融合 2D 特征形成 3D 特征体（feature volume），但它们的 3D 特征的聚合方式还是一直沿用了 3D CNN，没有任何改变。

在这项工作中，我们认为对于三维重建这个任务来说，三维特征的聚合是很重要的事情，所以我们提出了一个新的三维特征聚合方式——3D Transformer。就像 2D CNN 到 3D CNN 的转变一样，我们将 2D Transformer 升格到 3D，用来处理 3D 特征。

这个事情之所以之前没有人做，是因为 3D Attention 巨大的计算量，为了减少计算复杂度，我们提出了一个稀疏窗口注意力模块，从而只在局部窗口内的非空体素之间计算注意力。然后我们构建了一个自上而下-自下而上的 3D Transformer 网络用于 3D 特征聚合，其中的膨胀-注意力结构用来防止几何退化，两个全局模块用来增加全局感受野。

在多个数据集上的实验表明，这个 3D Transformer 网络生成了更准确、更完整的三维几何重建，相比之前的方法有明显的提升。尤其是，在 ScanNet 数据集上，重建精度提高了 41.8%，重建完整度提高了 25.3%。

论文题目：

3D-Former: Monocular Scene Reconstruction with SDF 3D Transformers

论文作者：

Weihao Yuan, Xiaodong Gu, Heng Li, Zilong Dong, Siyu Zhu

项目主页：

https://weihaosky.github.io/former3d/

论文链接：

https://arxiv.org/abs/2301.13510

代码链接：

https://github.com/alibaba-damo-academy/former3d

Introduction

单目三维重建是计算机视觉中的经典任务，对于自主导航、机器人和增强/虚拟现实等众多应用至关重要。这个视觉任务旨在仅利用一系列单目 RGB 图像，从一个不规则场景中重建一个精确和完整的三维几何形状。

目前，摄像机姿态可以通过 SLAM 或 SfM 系统比较准确地估计得到，然而从这些已有位姿的图像中恢复场景三维几何仍然是一个具有挑战性的问题，因为一个大场景环境是很复杂的，如各种物体、多样的光照、反光表面，还有不同焦距、畸变和传感器噪声的多样摄像机。

许多先前的方法采用多视角深度估计（MVS）来重建场景，他们预测每一帧图像的稠密深度图，这种 depth-based 方法可以估计准确的局部几何形状，但需要额外的步骤来融合这些深度图，例如解决不同视角之间的不一致性，这是比较困难的。

最近，一些方法尝试直接从截断有符号距离函数（TSDF）表示中回归完整场景的完整 3D 表面。他们首先使用 2D 卷积神经网络（CNN）提取 2D 特征，然后将这些特征反投影到 3D 空间。之后，3D feature volume 由 3D CNN 网络处理，输出一个 TSDF volume 预测，该预测经过 marching cubes 就可以提取到 mesh。这种重建方式是端到端可训练的，已经证明了能够输出准确、连贯和完整的网格。在本文中，我们沿着这种 volume-based 三维重建路径，直接回归 TSDF volume。

然而，之前的 volume-based 方法都是在研究如何提取 2D 特征、如何融合 2D 特征形成 3D feature volume，但它们的 3D 特征的聚合方式还是一直沿用了 3D CNN，没有任何改变。有些方法使用了 2D Transformer 在进行 2D 特征的处理，但是都没有人使用 3D Transformer 来做 3D 特征的处理，这是因为 3D multi-head attention 的计算量是爆炸的，每个 voxel 直接都需要计算 attention，在现有硬件下这几乎是不可实现。

在本文中，为了解决这个问题，使 3D Transformer 在三维重建中成为可能，我们提出了一个 sparse window multi-head attention 的模块，并在此之上构建了第一个 top-down-bottom-up 的 3D Transformer 网络。

Methodology

2.1 整体结构

我们进行三维重建的整体框架如图 1 所示。给定一个场景的图像序列和相应的相机内参和外参，我们首先在三个尺度上，提取 2D 空间中的图像特征，然后将这些 2D 特征反投影到 3D 空间中，并分别融合到粗、中、细三个特征体中。

然后，这三个特征体按照粗到细的顺序被我们的 SDF 3D Transformer 聚合。在粗和中尺度上，3D Transforme 的输出是两个占用体（occupancy volume）预测 O1、O2，用于将更精细层的 volume 稀疏化，而在细尺度上，输出是预测的 TSDF volume。这样以来，特征体可以被稀疏地处理以减少计算复杂度。最后，从 TSDF volume 中可以使用 Marching Cubes 来提取预测的网格 mesh。

2.2 三维特征构建

与之前的方法类似，我们使用了加权平均的方式融合不同视角投影得到的 3D 特征。与之前不同的是，我们在计算加权权重时加入了每个视角的 variance，并且在最终融合的 feature volume 中加入了 total variance，因为我们认为每个视角的 variance 反映了对最终融合的应有贡献度，比方说与整体平均差很多的值很可能是 outlier。具体细节和计算公式可以查看论文。

2.3 稀疏窗口多头注意力

多头注意力结构已被证明在许多视觉任务中很有效。然而，大多数任务仅限于 2D 特征处理，而非 3D 特征处理。这是因为多头注意力的计算复杂度通常比卷积网络高，而这个问题在 3D 特征中被进一步扩大。要计算一个 3D feature volume 的注意力，需要计算每一个体素和任何其他体素之间的注意力，即一个体素需要计算N_X×N_Y×N_Z 个注意力，而所有体素需要计算 N_X×N_Y×N_Z×N_X×N_Y×N_Z 个注意力，这是非常大的，很难在常规 GPU 中实现。

为了解决这个问题，使得 3D 特征的多头注意力成为可能，我们提出使用稀疏窗口结构来计算注意力。如图 1 所示，在中等和精细级别上，我们使用占用预测 O2、O1 来稀疏化体积，并仅计算非空体素的注意力。此外，考虑到附近的体素对当前体素的形状贡献更大，而远处的体素贡献较少，我们仅计算每个体素局部窗口内的注意力，如图 2 所示。因此，我们只计算小窗口内非空体素的多头注意力，从而显著减少了计算复杂度。

具体而言，对于特征体积 V 中的任何非空体素 vi，我们首先在以该体素为中心的一个 n×n×n 窗口内搜索所有非空体素，并得到相邻的体素 {vj，j∈Ω(i)}。然后计算查询、键和值的嵌入：

对于位置嵌入 P，我们希望对三维世界的坐标尺度具有不变性。因此，我们根据 volume 中相对体素位置而不是实际世界坐标进行计算，如下所示：

然后我们就可以计算注意力如下：

这样一来，计算复杂度就可以从

减少到

这里 O(ij) 是体素 vi 和 vj 之间的一个注意力计算的复杂度，N_occu 是 volume 中占用体素的数量，n_occu 是局部窗口内占用体素的数量。举个例子，假设 volume 的非空率为 10%，窗口大小为 volume 大小的 1/10，那么稀疏窗口注意力的计算复杂度只有稠密 3D 注意力的

2.4 SDF 3D Transformer

▲ 图3. SDF Transformer结构图

由于多头注意力的高资源消耗限制，大多数之前与 3D Transformer 相关的工作仅小心翼翼地进行了省资源的特征处理，基本没有下采样和上采样。然而，在 3D 重建中，自上而下自下而上的结构对于特征提取和预测生成更为合理，就像大多数基于 3D CNN 的结构一样。因此，在本文中，我们设计了第一个基于 3D Transformer 的自上而下自下而上结构，如图 3 所示。

以图 1 中的 fine volume（V0）为例，总共有四个特征尺度层级，即 1/2、1/4、1/8、1/16，如图 3 所示。在编码器部分，每个层级采用下采样和膨胀注意力的组合来下采样特征体积。然后使用两个稀疏窗口多头注意力模块来聚合特 feature volume。在底层，采用全局注意力模块来补偿窗口注意力的局部感受野，并使用全局上下文编码模块提取全局信息。

在解码器部分，使用逆稀疏 3D CNN 来上采样特征体积。因此，上采样后的最终形状应与输入相同。与 FPN 类似，下采样中的特征也添加到相应级别的上采样特征中。为了给网络增加变形能力，在上采样完成之后设置了一个后膨胀注意力块。最后，附加了一个带有 Tanh 激活的子流形 3D CNN 头来输出 TSDF 预测。

对于膨胀-注意力模块、全局注意力模块、全局上下文编码模块的详细介绍可以查看论文。