【论文阅读|cryoET】Learning Adaptive Tensorial Density Fields for Clean Cryo-ET Reconstruction-CSDN博客

本文链接：https://blog.csdn.net/m0_38068229/article/details/136380703

该工作发表在 NeurIPS2023(Poster)
有GitHub链接但是未开源。

摘要

提出了一种基于学习的框架，用于从cryoET倾斜序列数据进行3D重建。
cryoET面临的挑战主要有：缺失楔，数据量大，高噪声等。该工作通过对扫描样本的3D密度场使用基于自适应张量表示来解决以上挑战。
主要工作：
1. 优化四叉树结构来划分感兴趣的volume；
2. 学习表示每个节点中密度场的张量vector-matrix分解；
3. 使用的损失函数将可微断层扫描形成模型和三个正则化项结合：总变分（total variation），边界一致性约束（boundary consistency constraint）和各向同性傅里叶先验（isotropic Fourier prior）。
结果：该框架可以使用学习到的表示，去查询任何位置的密度，并获得高质量的3D tomogram。在合成数据和真实数据上证明了方法的优越性。在提高了重建质量的同时减少了计算时间和内存占用。

简介

cryoET面对的主要挑战有：

缺失楔现象；
采集期间发生的样本运动和变形引起的投影序列之间的未对准；
计算时间和内存方面的资源密集型任务；
使用低强度电子束扫描导致的高噪声。

已有一些方法来克服其中一些挑战：

可以使用标记跟踪（marker tracking）或通过跟踪投影之间的特定特征来校正未对准和运动；
最初通过分别在投影或者断层图像上应用基于图像或者基于体积的去噪方法来降低高水平噪声；最近有人提出可以使用奇数/偶数投影来重建断层图像的两个副本，然后使用Noise2noise方法学习断层图的去噪。
将扫描透射电子显微镜（STEM）的噪声模型和可微分框架中的隐式三维形状相结合，对重建数据进行去噪。

本文工作：

提出了一种cryoET重建框架，将重建和去噪联合起来。该架构基于四叉树结构，用它来定义自适应张量密度场表示。四叉树的每个节点都具有张量密度场，以表示其相关区域的断层图像的密度。
损失函数有四项：可微分断层扫描图像形成模型，促进重建平滑的总变分项，强制相邻节点之间一致的边界一致性约束，以及各向同性傅里叶先验惩罚方向性伪影并帮助断层图降噪。

方法

1. 概览

Adaptive TensorDF，利用四叉树结构创建多尺度且有效的基于张量的表示，用于重建噪声倾斜序列冷冻电子断层扫描图。框架将扫描样本的密度场表示为3D连续神经场。由于cryoET沿z轴的延伸范围较小，因此本文使用自适应四叉树结构而不是八叉树来划分体积。在四叉树的每个节点内，密度场使用张量表示。框架由三个步骤组成：四叉树更新、张量表示优化和3D体积重建。

四叉树更新步骤包括优化四叉树结构以实现重建体积的最佳划分。这是通过鼓励统一节点同时限制节点总数来实现的。在第二步中，优化损失函数以构建每个节点内密度场的基于张量的表示。损失函数包含一个数据保真度项和三个正则项。重建步骤包括对感兴趣区域（ROI）中的体积进行均匀采样并查询采样位置处的密度值。

在这里插入图片描述

3.2 图像形成模型（Image formation model）

cryoET的投影图像对应于对数空间中沿源和探测器之间的射线的密度积分。离散化后，给定射线的成像模型可以写成如下：
$b_i=A_i\mathbf{x}+n_i$

$A_i$ 表示沿射线i的Radon变换操作。 $b_i$ 和 $n_i$ 分别对应于探测器i测量的强度及其相关噪声。 $x$ 是我们想要重建的3D密度向量。实际上采集图像中的噪声非常复杂且难以建模，尤其是应用了不同的预处理步骤之后，为了简化，这里假设预处理投影中存在高斯白噪声。

通过重新组合所有捕获的光线，从公式1中推导出数据保真度损失（data-fidelity loss）：
在这里插入图片描述
这里K是一个二元掩码，用于禁用和基准标记相交的光线。引入此mask是为了抑制由这些标记引起的投影伪影。

3.3 基于坐标的表示（Coordinate-based representation, CBR）

在断层扫描应用中，基于坐标的网络已经被提议将感兴趣体积内的3D空间坐标映射到密度场。该映射由下式给出：
在这里插入图片描述
这里 $p_i$ 是要重建的体积中的3D坐标， $x_i$ 是对应的密度。 $f$ 是要优化的映射函数。使用CBR，捕获的投影是用过沿射线采样位置、将映射应用于每个样本，然后使用Radon系数求和它们的贡献来估计的。在常见的隐式表示中，映射函数被选为全连接的MLP（比如NeRF、ImplicitET），如图2.a。
然而这种表示法不适用于cryoET中遇到的大体积。为了解决这个限制，一些工作包括 KiloNeRF、ACORN、NeAT等建议使用统一块或多尺度基于八叉树的结构对感兴趣的体积进行分区，并且为每个分区分配更小的MLP或者解码器（块或八叉树节点）用于局部密度场表示。

如图2b，本文使用四叉树结构对ROI进行分区。因为cryoET在z轴上的拓展有限，所以这里的表示基于在z轴上拓展的xy平面的自适应四叉树结构。对于每个四叉树节点，使用基于张量的表示。因此对于感兴趣区域中给定的3D点 $p_i$ ，使用向量矩阵外积之和表示密度场：
在这里插入图片描述
其中D是单层解码器，将编码特征转换为输出密度。 $V_r^X$ 、 $V_r^Y$ 、 $V_r^Z$ 对应于张量分解的第r个分量的三种模式的因式分解向量。M表示矩阵因子。R是表示的等级，是根据重建样本的复杂性进行调整的超参数。选择较小的R有助于对断层图像进行降噪，但可能会丢失扫描样本的一些详细特征。
在这里插入图片描述
张量表示的主要优点是大大较少了表示体积所需要的参数数量，所以很适合cryoET数据中的大尺寸体积。
此外，在投影过程中考虑ROI之外区域的贡献也非常重要，以确保更准确的输出。然而这些区域不需要有高质量的表达，这里建议使用相同的表示方法但是用较小尺寸的张量。

3.4 正则化

Total Variation (TV loss)

总变分损失（TV loss） 在传统断层扫描重建中常用作空间正则化器来提高重建质量。在神经场表示方法中TV并不总是直接的，还可能会增加重建的计算复杂性。NeAT 是在应用解码器获得密度之前计算特征空间中的TV损失。TensoRF 使用向量和矩阵因子上的这种损失来处理观测值较少的区域中的噪声和异常值问题。在实现中，也将TV先验应用于ROI中启用的四叉树节点的向量和矩阵因子。然后对这些节点的贡献进行平均来计算TV loss：
在这里插入图片描述
其中α是缩放因子（经验值取0.1）。后面求的是向量和矩阵因子的梯度，然后 $N_{EN}$ 是四叉树的启用节点数。

边缘一致性约束（BCC loss）

在框架中，四叉树的每个节点都有子集的特征表示并单独优化，这会不可避免地在最终重建中引入不连续伪影。为了解决这个问题，提出了一种边界一致性约束（BCC），来惩罚从共享边缘上的采样点的相邻节点的张量表示中获得的特征之间的差异。

BCC loss：
在这里插入图片描述
其中 $O_b$ 指所有相邻四叉树节点对，∩n,m是指节点n和m之间的边界表面上的采样点集合，fφ§m 和fφ§n 是在点p处评估的密度，分别使用节点m、n的张量表示。

各向同性傅里叶先验（IFP loss）

通过将四叉树结构与张量密度场表示相结合，减少了表示3D体积所需的参数数量以及重建时间。但是低秩张量表示的一个缺点是它倾向于以轴线排列的条纹形式进行结构化伪影的重建。

在傅里叶空间中，这些伪影表现为沿垂直和水平方向的高峰。为了解决这个问题，在损失中引入了各向同性傅里叶先验（IFP），来惩罚傅里叶空间中的此类峰值。该约束旨在限制傅里叶域中水平频率和垂直频率与其他频率之间的差异。 为此，首先计算傅里叶域中每个环的平均振幅，它代表所有可能方向的给定空间频率。然后对远高于平均振幅的水平和垂直频率加以惩罚。此外，应用加权系数来惩罚更多的高频，主要对应于噪声。在实现中，对ROI进行粗采样，并查询这些样本处的体积。然后逐个slice计算傅里叶变换。

IFP loss：
在这里插入图片描述
ζ0 是频率参数，根据频率控制损失的权重。

3.5 自适应张量密度场优化

损失函数定义如下：
在这里插入图片描述
下面描述框架训练步骤的关键要素：更新四叉树和沿射线采样。

四叉树更新：

如图3，使用四叉树结构构建张量表示。
在这里插入图片描述

首先，从ROI初始化四叉树结构，并且禁用该区域之外的节点。
然后，对每个节点的体积进行均匀采样，并计算采样位置的密度值的标准差STD，来定义四叉树更新损失。给定节点中的高STD表明它包含精细的细节，因此这个节点更有可能被分裂成四个子节点。具有低STD的节点可能包含较少的特征并代表扫描样本的均匀区域，这样的节点更有可能被合并或者保持不变。这里使用 mixed-integer program(MIP) 来优化这个过程。其中，在每次迭代中根据更新损失，节点要么合并、分裂，要么保持相同。优化期间节点总数低于固定限制。

使用 coarse-to-fine 的策略来加速四叉树更新并处理高噪声水平：

先使用下采样投影来更新四叉树，在此步骤中还更新每个节点内的张量表示。该策略减少了噪声对四叉树更新的影响。
经过一些epochs后，保持四叉树结构固定，并继续使用原始投影优化每个节点的张量表示。

光线采样：

在优化过程中对每条光线进行采样，以获得用于密度积分和损失评估的3D位置列表。所应用的采样不是均匀的，而是遵循考虑四叉树结构的分层随机采样策略。对于光线经过的每个节点，在光线上选取 $N_q$ 个3D位置，从光线的 $N_q$ 个均匀段中随机采样。 $N_q$ 定义如下：
在这里插入图片描述
其中 $N_{max}$ 是对应于每个节点的最大样本数的超参数， $l_q$ 是四叉树节点q内部的射线长度， $d_q$ 是q的对角线长度。

4. 实验

本文代码用c++实现来加速重建过程。
设计了一些列实验来展示本文框架在合成数据集和真实数据集上的有效性。然后将本文方法和不同基线方法的性能进行比较。

SART+TV: 一种成熟的迭代重建技术SART和全变分先验结合；
ImplicitCT: 用于cryoET的隐式神经表示；
I-NGP：在cryoET重建任务上对Instant-NGP的重新实现;
TensoRF：在cryoET重建任务上对TensoRF的重新实现。

为了公平对比，所有输出密度都归一化到[0,1]之间。

4.1 合成数据集实验

实验中使用了ImplicitCT中介绍的合成数据集，该数据集由随机分布的椭圆体壳组成，其随机密度模仿了15埃的ZIKV病毒粒子的密度范围。仿真tomogram的角度范围是 [-70°, 69.5°]，角度步距1.5°。然后向投影中添加高斯噪声，并测试了几个噪声水平。当噪声的标准偏差超过0.05时，模拟投影就噪声而言在视觉上和真实投影很接近。该合成数据集用于本文的噪声水平方法的参数调整和鲁棒性评估。

参数调整
实验发现最重要的参数是张量维度（矩阵向量因子的维度）和特征大小。张量维度的选择影响训练和重建的速度和质量。较小的维度会导致更快的训练和更平滑的重建，但是可能会损失一些细节。较大的维度捕获更多细节，但也可能引入过拟合或噪声，因为网络在学习结构后往往会学习噪声。特征大小也会影响训练速度和重建质量。较小的特征可以加快训练速度，但会丢失一些细节；较大的特征可以保留更多细节，但增加计算成本。选择时需要一些权衡。
实验针对这些参数的不同值测量了重建体积的3D峰值信噪比（PSNR）和3D结构相似性指数测量（SSIM），结果如下：
对噪声等级的鲁棒性
模拟在不同的噪声水平下进行实验，和基线方法进行对比。图4说明了根据标准差范围[0.01, 0.12]内不同噪声水平的每种方法的结果计算出的PSNR和SSIM。无论噪声水平如何，本文方法在这两方面都是最优。基于神经表示的方法对噪声非常鲁棒，因为对于较大的噪声水平，性能不会因此下降。然而implicitET方法不适应高噪声的情况，因为它依赖于从中学到的噪声水平。

图5中显示了可视化的切片，表示使用不同方法从模拟噪声投影中得到的重建结果。噪声水平选用0.08，来说明每种方法对噪声水平的鲁棒性。数据集中的两个区域被放大以说明细节恢复（红框）和均匀区域（蓝框）的去噪效果。
SART+TV 的细节恢复较差，且保持较高的噪声水平。
ImplicitET 相对来说在均匀区域去噪效果良好，但是细节无法正确恢复。
I-NGP和TensoRF方法在保留细节和降低噪声方面表现相对较好，但TensoRF引入了一些伪影。可以在常规视图的某些椭圆中注意到这些伪影。
本文方法在两项任务中都表现最好。通过比较Ours和Ours W/O $L_{IFP}$ ，可以看到可向同性傅里叶先验在减少伪影和提高降噪效果方面的影响。

4.2 真实数据集实验

在EMPIAR 10643 和 EMPIAR 10761上进行评估。

EMPIAR 10643：HIV-1 角度范围 [-60°,60°]，增量为3。重建了series 40和51。
EMPIAR 10761：HEK cell 角度范围 [-60°,60°]。

使用IMOD对数据集进行预处理进行投影对齐和CTF校正。为了公平，将所有输出密度标准化为范围[0,1]。

图6展示了重建结果的比较。
在这里插入图片描述

去噪能力评估

结果显示本文方法在均匀区域具有最佳去噪效果。

SART+TV：尽管使用了TV正则化，仍然在均匀区域内表现出噪声。增加TV约束的权重来增强SART+TV的去噪能力会导致主要特征的损失。
implicitET 会产生过度平滑的重建，可能是由于用于学习噪声统计的数据集和我们使用的数据集之间的差异引起的。
I-NGP和TensoRF具有相对较好的去噪效果。一些方向性伪影会较低TensoRF的质量。它们是由于ROI之外的区域的影响造成的。该方法没有很好的对这些区域进行建模。
本文方法除了最佳的去噪能力之外，在提高特征与背景之间的对比度方面也最有效。

为了定量评估每种方法的去噪效果，使用两个统计指标进行了比较：对比度和噪声比（CNR）和等效外观数（ENL）。

CNR指标测量感兴趣区域和均匀背景之间的对比度。同时，ENL指标评估均匀区域的平滑度。

这里没有比较implicitCT，因为该方法在重建真实数据时表现不佳。表3中的结果证明本文方法的CNR和ENL更高，增强了特征区域和背景之间的对比度，同时产生更平滑的均匀区域。该表还表明SART+TV的结果噪声最大。
在这里插入图片描述

四叉树正则化

四叉树结构的好处不仅仅是加快计算速度。和TensRF相比，它允许在XZ和YZ平面上拥有更多的局部矩阵表示，而TensoRF只能具有整个场景的全局平面。这可以减少重建中的噪声和误表示，但参数大小仅仅减少50%。

表4和图8显示了本文方法和TensoRF W/O $L_{TV}$ 、TensoRF、TensoRF W $L_{IFP}$ 。Ours 具有更好的细节保留和去噪效果，实现了质量的显著提高。
在这里插入图片描述

详细特征分析

图6显示本文方法具有最好的特征恢复。对EMPIAR 10643-40，重建可以更好地解析刺突蛋白。对EMPIAR 10751 的重建使得 HEK细胞结构更清楚。图7中的intensity剖面证实了这些观察结果，其中本文方法产生了更规则的剖面，峰值（背景区域）和谷值（特征）之间存在重要差异。所有其他方法由于重建中的残留噪声都显示出更多的中间峰谷。
在这里插入图片描述