Generalized Diffusion MRI Denoising and Super-Resolution using Swin Transformers

基于Swin Transformers的广义扩散MRI去噪与超分辨率

论文地址：[2303.05686] Generalized Diffusion MRI Denoising and Super-Resolution using Swin Transformers (arxiv.org)

项目地址： https://github.com/ucsfncl/dmri-swin

Abstract

扩散 MRI 是一种非侵入性体内医学成像方法，能够绘制人脑的组织微观结构和结构连接图，并检测其他临床神经成像技术无法看到的变化，例如大脑发育和损伤。然而，获取具有高角度和空间采样的高信噪比 (SNR) 数据集需要非常长的扫描时间，限制了在许多重要临床环境中的使用，尤其是儿童、老年人和患有急性神经系统疾病的急诊患者可能不会能够在没有清醒镇静或全身麻醉的情况下配合 MRI 扫描。在这里，我们建议使用 Swin UNET Transformers (Swin UNETR) 模型，在增强的人类连接组计划 (HCP) 数据上训练并以注册的 T1 扫描为条件，对不随采集参数、患者群体、扫描仪和站点变化的扩散 MRI 执行广义去噪和超分辨率。我们在正常成人志愿者中使用人工下采样的 HCP 数据定性地展示了超分辨率。我们对另外两个不相关的数据集（一个患有神经发育障碍的儿童和一个创伤性脑损伤患者）进行的实验表明，尽管数据分布变化很大，但我们的方法表现出出色的去噪效果。可以通过仅对一个额外的主题进行微调来实现进一步的改进。我们将我们的模型应用于扩散张量（二阶球谐函数）和高阶球谐函数估计，并显示优于当前最先进方法的结果。我们的方法可以开箱即用或进行微调，以对各种扩散 MRI 数据集进行去噪和超分辨。

1 Introduction

扩散MRI（dMRI）可以提供有价值的临床信息并评估组织微观结构；然而，其低信噪比（SNR）导致较差的诊断和定量准确性。为了提高信噪比，大多数dMRI协议需要低角度和空间分辨率和/或长扫描时间，这限制了在许多重要的临床环境中使用。因此，人们对在不影响SNR或分辨率的情况下具有短的患者扫描时间非常感兴趣。

已经提出了几种有监督的方法来对大脑dMRI扫描进行去噪；然而，它们由于缺乏普遍性而受到限制。通常，它们只在一个b值和一组预先指定的扩散编码方向上工作，只用于预测一组微观结构参数，或者在同一数据集上进行训练和验证，如人类连接体项目[16] [10]。由于不同的采集参数、扫描仪和患者群体，扩散数据可能会有很大的差异，因此通常首选无监督或自监督的去噪方法[5]。

在这里，我们建议使用Swin UNET Transformers（Swin UNETR）模型[7]来对以注册T1扫描为条件的dMRI数据进行去噪。与其他监督方法不同，其他监督方法利用HCP数据集的一小部分（通常为40名受试者）进行训练[16] [10]，我们使用全套HCP数据，使用所有b值和扩散编码方向进行训练，并应用简单的数据增强，如随机翻转、旋转、缩放和k空间下采样。在大型数据集（约27万次3D扫描）上进行训练，使Swin UNETR模型能够学习在许多设置中表现良好的去噪函数。我们在一个保留的HCP数据集以及两个不相关的数据集上验证了我们的方法，并显示出与当前最先进的无监督方法相比的显著性能优势。此外，我们还表明，即使只对一名受试者进行微调，也能提高性能，而且我们的方法还可以通过对HCP受试者的定性评估来超级解析dMRI数据。

2 Methods

2.1 Data

在我们的实验中，我们使用了来自三个不同数据集的1155名独特受试者的数据。

• HCP：1065名受试者，来自人类连接体项目（HCP）青年数据集[18]，使用90个扩散编码方向采集，b值为b=1000、2000、3000 s/mm²，分辨率为1.25 mm。我们随机选择985名受试者进行训练，40名受试人进行验证，40名接受测试。
• TBI：20年前，45名成人轻度创伤性脑损伤（mTBI）患者在GE扫描仪上使用与[20]相同的方案获得，55个扩散编码方向，b=1000 s/mm²，标称分辨率为1.8 mm（在k空间中进行零插值后，xy中为0.9 mm）。我们随机选择了5名受试者进行微调，40名受试人进行测试。
• SPIN：45名8-12岁患有神经发育障碍的儿童在西门子扫描仪上获得，具有64个和96个扩散编码方向，b值分别为b=1000、2500 s/mm²（TE=72.20 ms，TR=2420 ms，翻转角度=85^◦ ) 分辨率为2.00mm。我们随机选择了5名受试者进行微调，40名受试人进行测试。

DMRI数据用Synthstrip[8]进行颅骨剥离，用Eddy[1]校正涡流引起的畸变和受试者的运动，并用基于边界的配准（BBR）[6]与结构T1扫描对齐。T1扫描也使用Synthstrip剥离头骨，并使用SynthSeg[2]进行分割。在1.25mm处对配准的T1和dMRI扫描进行采样，并将其用作模型的输入。

2.2 Evaluation Strategy

为了进行评估，将模型预测重新采样到dMRI分辨率。为了评估扩散张量估计，发现了主特征向量（V1）、分数各向异性（FA）、轴向扩散率（AD）、径向扩散率（RD）和平均扩散率（MD）的基本事实和模型预测之间的平均绝对误差（MAE）。为了评估高阶球面谐波，使用了地面实况和模型预测之间的Jensen–Shannon距离（JSD），投影到均匀分布的362方向半球[3]。使用[17]中描述的程序，选择了唯一拟合所需的最小扩散梯度数：扩散张量为6，4阶球谐为15，6阶球谐为28。我们将我们的模型的性能与用于扩散张量拟合的块匹配和4D滤波（BM4D）[14]以及用于4阶和6阶球面谐波估计的Marchenko Pastur主成分分析（MPPCA）[19]进行了比较，因为我们发现BM4D和MPPCA在这两个领域表现最好。在每种情况下，通过使用所有获得的扩散梯度方向拟合模型来找到基本事实。为了评估超分辨率，将来自HCP受试者的dMRI数据在k空间下采样2倍，然后线性上采样以模拟低分辨率采集。

2.3 Training and Implementation

Swin UNETR模型[7]使用PyTorch和MONAI实现，其架构如图1所示。该模型在1个NVIDIA V100 GPU上使用模型输出和地面实况之间的均方误差损失进行训练。为了获得用于训练的地面实况dMRI数据，对每个壳体拟合六阶球面谐波，并将其投影到所获取的方向上。我们选择AdamW[12]作为学习率为1e-4的优化器，并执行最大范数为1.0和随机权重平均的梯度裁剪[9]。为了减少内存使用，我们使用模型检查点并以16位精度进行训练。由于计算的限制，只完成了五个时期的训练。在训练过程中，我们首先使用类似于[13]的程序，在频率空间中以0.5的概率对dMRI扫描进行下采样，使其各向异性分辨率在1.25和3 mm之间，并线性上采样回到1.25 mm分辨率。从扫描中裁剪128 x 128 x 128的随机补丁，并以0.5的概率沿所有轴随机翻转，并以相等的概率沿全部轴随机旋转0、90、180或270度。将输入dMRI贴片归一化为具有零平均值和单位方差，并且将输入T1贴片归一化为在0.25和4.0之间。为了进行推理，我们使用了一种重叠为0.875的滑动窗口方法。微调是通过对五个受试者中的一个的外部数据进行额外训练来实现的，三个时期的学习率为1e-6，并报告了平均结果。

图1：Swin UNETR体系结构的概述。输入是级联的3D T1和dMRI扫描，由Swin transformer[11]以多分辨率编码，并输入到残差卷积神经网络（CNN）解码器中，以重建地面实况dMRI扫描。有关更多详细信息，请参见[7]。

3 Results and Discussion

3.1 Comparison with Other Methods

对于扩散张量拟合，表1显示，在HCP和TBI测试数据集中，在白质（WM）、灰质（GM）和脑脊液（CSF）的所有指标中，Swin模型在没有微调的情况下实现了比BM4D更低的MAE。对于SPIN数据集，除了GM和CSF中的FA估计外，Swin模型优于BM4D；然而，对于CSF中的MD估计，不应用去噪优于BM4D和Swin模型，这可能是由于Eddy之后未校正的运动引起的噪声。

表1：FA、MD、RD、AD和V1的MAE估计，使用白质（WM）、灰质（GM）和脑脊液（CSF）的六方向HCP、SPIN和TBI测试数据，通过无去噪（RAW）、BM4D、无微调的Swin（Swin）和对一个受试者进行微调的Swin（Swin-F1）。最佳结果以粗体显示。

图2中的定性比较与这些定量结果一致，并表明Swin模型能够在没有过度平滑BM4D的情况下捕捉WM和GM微观结构中更多的精细特征。

图2：对来自HCP（顶部）、TBI（中间）和SPIN（底部）的测试数据进行去噪的基本事实（GT）、无去噪（RAW）、BM4D、无微调的Swin（Swin）之间的视觉比较。MAE图像显示了平均绝对误差。

对于4阶和6阶球面谐波拟合，表2显示，Swin模型在所有数据集的GM和CSF中实现了最低的JSD。对于六阶球面谐波拟合，即使在HCP数据中，Swin模型在WM中的表现也比MPPCA差，这可能是因为模型不足（由于计算限制，只能完成五个训练时期），也可能是因为Swin模型一次对一个方向进行去噪，而MPPCA能够对所有方向进行集体去噪。

表2。使用白质（WM）、灰质（GM）和脑脊液（CSF）中的15个方向（4阶）和28个方向（6阶）HCP、SPIN和TBI测试数据，通过无去噪（RAW）、MPPCA、无微调的Swin（Swin）和对一个受试者进行微调的Swin（Swin-F1），在基本事实和模型估计之间进行JSD。最佳结果以粗体显示。

我们的结果是通过简单的数据扩充实现的，尽管使用更广泛的模拟，如[15]中的模拟，可能会导致更大的可推广性。此外，通过分别处理每个dMRI体积，我们能够考虑完整的大脑体积，但不要利用体积之间的相关性，例如[10]中基于Transformers 的patch方法。由于GPU内存的限制，在利用空间相关性和角度相关性之间存在权衡。最后，与细胞神经网络相比，Swin转换器可以有效地捕获长程依赖性，但需要更广泛的预训练，这是我们方法成功的关键[4]。

3.2 Finetuning

即使对一个项目进行微调，也始终提高了扩散张量估计的去噪性能，但结果在高阶球面谐波拟合中是混合的。这可能是因为微调减少了模型偏差，但增加了方差，与扩散张量估计中使用的六个方向相比，方差可以在用于计算球面谐波的15或28个方向上累积误差。此外，通过对多个主题进行微调，我们没有体验到显著的性能优势。据我们所知，在dMRI去噪的背景下，微调从未被报道过，值得进一步研究。

3.3 Super-Resolution

尽管我们的模型没有针对超分辨率进行训练，但它可以用于将dMRI数据集重新采样到1.25 mm的分辨率。定性比较表明，在心室后周WM中，Swin能够比BM4D捕捉更多的精细微观结构，并避免过度模糊。

图3：HCP受试者心室后周WM超分辨率的基本事实（GT）、无后处理（RAW）、BM4D和无微调的Swin（Swin）之间的视觉比较。数据被k空间下采样2倍，然后线性上采样回到1.25mm。