【论文阅读】R2-Gaussian: Rectifying Radiative Gaussian Splatting for Tomographic Reconstruction

Abstract

3DGS 在图像渲染和表面重建方面显示出可喜的结果。然而它在体积重建任务（如X射线计算机断层扫描）中的潜力仍未得到充分探索。本文介绍了$R^2$-Gaussian，这是第一个基于3DGS的稀疏视图断层重建框架。通过仔细推导X射线光栅化函数，我们发现了标准 3DGS 公式中以前未知的积分偏差，这阻碍了准确的体积检索。

为了解决这个问题，本文提出了一种新的校正技术，通过将投影从3D重构为2D高斯。新方法提出了三个关键创新：（1）引入定制的高斯核，（2）将光栅化拓展到 X 射线成像，以及（3）开发基于 CUDA 的可微体素。大量的实验表明，本文方法在 PSNR 和 SSIM 中分别优于最先进的方法。重要的是可以在3分钟内提供高质量的结果，比基于 NeRF 的方法快12×与传统算法相当。该方法的优越性能和快速收敛性凸显了实用价值。

Introduction

本文揭示了 3DGS 中的固有的积分偏差（integration bias）。尽管这种偏差对图像渲染的影响可以忽略，但严重阻碍了体积重建。标准3DGS 在将 3D 高斯核投影到 2D 图像平面上时忽略了与协方差相关的缩放因子。这种公式导致从不同视图查询的体积属性不一致。除了积分偏差之外，将 3DGS 应用于断层扫描还面临其他挑战，例如自然光和X射线成像之间的差异以及缺乏从核查询体积的有效技术。

$R^2$-Gaussian 通过三项重要改进实现了无偏差的训练流程：

1. 引入了新型辐射高斯核，充当由中心密度、位置和协方差参数化的局部密度场；使用分析方法 FDK 初始化高斯参数，并使用光度损失对其进行优化；
2. 修改3DGS光栅化器以支持X射线成像。这是通过推导新的X射线渲染函数并校正积分偏差以实现精确的密度检索来实现的。
3. 开发了基于CUDA的可微分体素化器，不仅可以从高斯中提取体积，还可以在训练期间实现基于体素的正则化。

本文在多个模态上评估了该方法，包括人体器官、动物和植物以及人造物体。

Related work

Tomographic reconstruction

3DGS

Preliminary

3.1 X-ray imaging

投影 $I\in {\mathbb{R}}^{H\times W}$ 测量了穿过材料的射线衰减。对于具有初始强度 $I_0$ 和路径边界 $t_n$ 和 $t_f$ 的X射线，相应的原始像素值通过 Beer-Lambert 定律给出：$I^{\prime }(r)=I_{o}exp(-{\int }_{t_n}^{t_f}\sigma (\mathbf{r}(t))dt)$，这里 $\sigma (\mathbf{x})$ 是$\mathbf{x}\in {\mathbb{R}}^3$处的密度（或物理学中的衰减系数）。为了计算简单，断层扫描通常将原始数据转换成对数空间（公式1）：
$$I(r)=log{I}_o-log{I}^{\prime }(\mathbf{r})={\int }_{t_n}^{t_f}\sigma (\mathbf{r}(t))dt$$
每个像素值 $I(r)$ 表示沿光线路径的密度积分。除非另有说明，我们使用对数投影作为输入。断层扫描重建的目标是估计 $\sigma (X)$ 的3D分布，输出为离散体积，并使用从N个不同角度捕获的噪声X射线投影 { I i } i = 1 , . . . , N \{I_i\}_{i=1,...,N} {Ii​}i=1,...,N​。

3.2 3DGS

3DGS 利用一系列 3D高斯核 G 3 = { G i 3 } i = 1 , . . . , M G_3=\{G_i^3\}_{i=1,...,M} G3​={Gi3​}i=1,...,M​对场景进行建模。每个高斯核由位置、协方差、颜色和不透明度参数化。光栅化器 R 从这些高斯函数渲染 RGB 图像 $I_{rgb}\in {\mathbb{R}}^{H\times W\times 3}$（公式2）：
$${\mathbf{I}}_{rgb}=\mathcal{R}\left({\mathbb{G}}^3\right)=\mathcal{C}\circ \mathcal{P}\circ \mathcal{T}\left({\mathbb{G}}^3\right)$$

其中 $\mathcal{C},\mathcal{P},\mathcal{T}$ 分别是变换、投影和合成模块。首先T将3D高斯变换到ray space，将观察光线和坐标轴对齐以提高计算效率，然后将变换后的3D高斯投影到图像平面上。投影的2D高斯和3D的保持相同的透明度和颜色，但省略了位置和协方差的第三行和第三列。随后使用alpha混合合成这些2D高斯函数来渲染一个RGB图像。光栅化器是可微分的，允许使用光度损失来优化内核参数。3DGS 利用SfM点初始化稀疏高斯函数。训练时用自适应控制策略动态地致密化高斯来改善场景表示。

Methods

4.1 Representing objects with radiative gaussians

我们用一组可学习的3D高斯核表示目标对象，将其称为辐射高斯。每个核定义一个局部高斯形密度场（由中心密度值、位置和协方差参数化）（公式3）：
$$G_i^3(x|{\rho }_{\mathbf{i}},{\mathbf{p}}_{\mathbf{i}},{\mathbf{\Sigma }}_{\mathbf{i}})={\rho }_i\cdot exp(-\frac{1}{2}(\mathbf{x}-{\mathbf{p}}_i{)}^T{\Sigma }_i^{-1}(\mathbf{x}-{\mathbf{p}}_i))$$
优化时将协方差矩阵 $\Sigma$ 进一步分解成旋转矩阵 $R_i$ 和尺度矩阵 $S_i$ : ${\Sigma }_i=R_i{S}_i{S}_i^T{R}_i^T$。然后通过对核的密度贡献和来计算位置$\mathbf{x}\in {\mathbb{R}}^3$ 处的总密度（公式4）：
$$\sigma (\mathbf{x})=\sum _{i=1}^M{G}_i^3(\mathbf{x}|{\rho }_i,{\mathbf{p}}_i,{\Sigma }_i)$$

和标准3DGS相比，我们的kernel消除了和视图相关的颜色，因为X射线衰减仅取决于各向同性密度。更重要的是，我们为辐射高斯定义了密度查询函数 $\sigma (\mathbf{x})$ ，使其对于2D图像渲染和3D体积重建都很有用。

Initialization

使用从分析方法获得的初步结果来初始化辐射高斯。具体地，使用FDK在不到1秒的时间内低质量重建体积。然后排除具有低于密度阈值的空白区域，并随机采样M个点作为kernel的位置。和3DGS一样，将高斯的大小设置为最近邻距离并且假设物旋转。中心密度从FDK体积中查询得到，凭经验用k缩小查询的密度，来补偿kernel之间的overlap。

4.2 训练

辐射高斯先从FDK体积初始化，然后对投影进行光栅化来计算光度损失，并对微小密度体积进行体素化来进行3D正则化。使用改进的自适应控制来进行高斯致密化，以获得更好的表示。

流程如图：

4.2.1 X-ray 光栅化

本节重点介绍X射线光栅化的理论推导。投影的像素值是沿着相应光线路径的密度积分。
把公式（4）代入（1）得到公式（5）：
$$I_r(\mathbf{r})=\int \sum _{i=1}^M{G}_i^3(\mathbf{r}(t)|{\rho }_i,{\mathbf{p}}_i,{\Sigma }_i)dt=\sum _{i=1}^M\int G_i^3(\mathbf{r}(t)|{\rho }_i,{\mathbf{p}}_i,{\Sigma }_i)dt$$

其中$I_r(\mathbf{r})$ 就是渲染的像素值。这意味着我们可以单独积分每个3D高斯以光栅化X射线投影。之前公式1里的边界被忽略，因为假设所有高斯都在目标空间内有界。

变换

由于锥束X射线扫描仪可以类似于针孔相机进行建模，因此我们按照[64]将高斯从世界空间转移到射线空间。在射线空间中，射线平行于第三坐标轴，有利于解析积分。由于光线空间的非笛卡尔性质，我们对方程采用局部放射变换，得到公式（6）：
$$I_r(\mathbf{r})\approx \sum _{i=1}^M\int G_i^3(\tilde{\mathbf{x}}\mid {\rho }_i,\underset{{\tilde{\mathbf{p}}}_i}{\underbrace{\varphi (\mathbf{p})}},\underset{{\tilde{\mathbf{\Sigma }}}_i}{\underbrace{{\mathbf{J}}_i\mathbf{W}{\mathbf{\Sigma }}_i{\mathbf{W}}^{\top }{\mathbf{J}}_i^{\top }}})d{x}_2,$$
其中 $\tilde{\mathbf{x}}=[x_0,x_1,x_2{]}^T$是光线空间中的点，${\tilde{\mathbf{p}}}_i\in {\mathbb{R}}^3$是通过投影映射得到的新高斯位置，${\tilde{\mathbf{\Sigma }}}_i\in {\mathbb{R}}^{3\times 3}$是由局部逼近矩阵 $J_i$ 和观察变换矩阵 W 控制的新高斯协方差矩阵。$\varphi ,J_i,W$根据扫描仪参数确定。

投影和合成

归一化3D高斯分布的一个良好特性是它沿一个坐标轴的积分会产生归一化2D高斯分布。
将公式（3）代入公式（6），可以得到公式（7）：

其中 $\hat{\mathbf{x}}\in {\mathbb{R}}^2,\hat{\mathbf{p}}\in {\mathbb{R}}^2,\hat{\mathbf{\Sigma }}\in {\mathbb{R}}^{2\times 2}$ 分别是通过删除对应的第三行和第三列获得的。公式（7）表明X射线投影可以通过简单地对2D高斯进行求和来渲染，而不是进行alpha合成。

Intergration bias

在投影过程中我们的2D高斯与3DGS中的原始高斯之间的一个关键区别是中心密度 ${\hat{\rho }}_i$。如公式（7），我们使用协方差相关因子 ${\mu }_i$ 来缩放密度，$\widehat{{\rho }_i}={\mu }_i{\rho }_i$，而3DGS中没有。这意味着 3DGS 实际上学习的是2D图像平面中的集成密度，而不是3D空间中的实际密度。这种集成偏差虽然对成像渲染的影响可以忽略不计，但会导致密度检索出现严重的不一致（图5展示了简化的2D-1D投影的不一致性）。当尝试在3D空间中用 ${\rho }_i=\widehat{{\rho }_i}/{\mu }_i$ 恢复中心密度 $\rho$ 时，我们发现不同的视角下的 ${\mu }_i$ 导致不同的结果。这违反了 ${\rho }_i$ 的各向同性性质，使我们无法确定正确的值。相反，我们的方法将实际的3D密度分配给kernel并向前计算2D投影，从而从根本上解决了这个问题。虽然概念上简单，但实现该想法需要大量的工程工作，包括重新变成CUDA中的所有反向传播例程。

4.2.2 密度体素化

我们开发了一个体素化器 $\mathcal{V}$ 来有效地从辐射高斯查询密度体积 $\mathbf{V}\in {\mathbf{R}}^{X\times Y\times Z}$。受到tile-based光栅化器的启发，我们的体素化器首先将目标空间划分为多个8×8×8大小的3D图块。然后，它会剔除高斯分布，保留那些有99%置信度与图块相交的分布。在每个3D图块中，通过公式（4）求和附近kernel的贡献来并行计算体素值。我们在CUDA中实现体素化器和反向传播，使其可微以进行优化。这种设计不仅加速了查询过程，还允许我们使用3D先验对辐射高斯进行正则化。

4.2.3 优化

使用随机梯度下降来优化辐射高斯。除了光度 L1损失和 D-SSIM 损失之外，我们还进一步结合了3D TV正则化损失作为同质性先验。在每次迭代中，随机查询一个小密度体积（和目标输出的间距相同）并最小化总变分。总体的训练损失定义如下：
$${\mathcal{L}}_{total}={\lambda }_1(I_r,I_m)+{\lambda }_{ssim}{\mathcal{L}}_{ssim}(I_r,I_m)+{\lambda }_{tv}{\mathcal{L}}_{tv}(V_{tv})$$
r代表render projection，m表示measured projection。

高斯自适应调整：我们删除空的高斯并致密化（clone或者split）那些具有大的损失梯度的高斯。考虑到人体器官等物体具有广泛的同质区域，我们不会修剪大高斯。至于致密化，我们将原始高斯函数和复制高斯函数的密度减半。该策略降低了新高斯导致的性能突然下降，从而稳定训练过程。

讨论和结论

R2-Gaussian继承了3DGS的一些局限性，比如不同模态的训练时间不同，极其稀疏视图条件下的针状伪影以及其他断层扫描任务的suboptimal extrapolation。此外，我们不考虑有关扫描几何和辐射特性的校准误差。尽管有这些限制，本文方法的性能和速度使其对于现实世界的应用很有价值。