[论文解读] Neuro-TransUNet: Segmentation of stroke lesion in MRI using transformers

Muhammad Nouman, Mohamed A. Mabrok+1 more|arXiv (Cornell University)|Jun 10, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

Neuro-TransUNet 提出了一种混合深度学习框架，将 U-Net 的空间特征提取能力与 SwinUNETR 的全局上下文建模能力相结合，以提升 3D MRI 缺血性卒中病灶的分割性能。通过集成先进的特征融合技术与全面的预处理流程，该模型在 ATLAS v2.0 基准测试中取得了 0.730 的测试 Dice 分数，优于现有方法，在缺血性卒中病灶分割任务中达到了新的最先进水平。

ABSTRACT

Accurate segmentation of the stroke lesions using magnetic resonance imaging (MRI) is associated with difficulties due to the complicated anatomy of the brain and the different properties of the lesions. This study introduces the Neuro-TransUNet framework, which synergizes the U-Net's spatial feature extraction with SwinUNETR's global contextual processing ability, further enhanced by advanced feature fusion and segmentation synthesis techniques. The comprehensive data pre-processing pipeline improves the framework's efficiency, which involves resampling, bias correction, and data standardization, enhancing data quality and consistency. Ablation studies confirm the significant impact of the advanced integration of U-Net with SwinUNETR and data pre-processing pipelines on performance and demonstrate the model's effectiveness. The proposed Neuro-TransUNet model, trained with the ATLAS v2.0 \emph{training} dataset, outperforms existing deep learning algorithms and establishes a new benchmark in stroke lesion segmentation.

研究动机与目标

• 解决在 3D MRI 中准确分割缺血性卒中病灶的挑战，这些病灶通常较小、形态多变，且因解剖结构复杂而难以检测。
• 克服现有模型在病灶大小差异性和数据不平衡方面存在的局限，特别是对小病灶或多病灶的处理能力不足。
• 通过开发全面的预处理流程，提升模型的鲁棒性与泛化能力，改善数据质量与一致性，同时避免头骨去除（skull stripping）操作。
• 通过先进的特征融合与分割合成技术，整合 U-Net（局部空间特征）与 SwinUNETR（全局上下文理解）的优势。
• 通过在具有挑战性的 ATLAS v2.0 数据集上验证模型，建立卒中病灶分割的新基准，证明其在性能上优于先前的深度学习方法。

提出的方法

• 将 U-Net 的编码器-解码器架构与基于 Transformer 的 SwinUNETR 解码器相结合，实现局部空间特征学习与全局长距离依赖建模的协同优化。
• 实施多阶段特征融合机制，在多个分辨率层级上聚合来自 U-Net 和 SwinUNETR 分支的特征，以增强表征学习能力。
• 应用全面的预处理流程，包括重采样、偏置场校正与数据标准化，以提升跨受试者输入数据的质量与一致性。
• 采用体素级推理方法，绕过传统的 2D/3D 后处理步骤，实现端到端的 3D 分割，提升空间一致性。
• 利用 ATLAS v2.0 数据集进行训练与评估，该数据集包含多种缺血性卒中病例，涵盖不同大小与解剖位置的病灶。
• 采用考虑类别不平衡的损失函数进行模型优化，提升对小病灶及低频病灶的检测敏感性。

实验结果

研究问题

• RQ1 结合 U-Net 与 SwinUNETR 的混合架构是否能提升 3D MRI 中所有大小缺血性卒中病灶的分割准确性？
• RQ2 全面的预处理流程（不包含头骨去除）在真实世界 MRI 数据上对模型性能与鲁棒性的提升程度如何？
• RQ3 先进的特征融合与分割合成技术在整合局部与全局特征以改善病灶边界勾画方面的有效性如何？
• RQ4 Neuro-TransUNet 框架是否在 ATLAS v2.0 基准测试中优于现有的最先进模型？
• RQ5 该模型在处理复杂、多样化或罕见病灶模式时存在哪些局限性？未来工作如何进一步提升其泛化能力？

主要发现

• Neuro-TransUNet 在 ATLAS v2.0 基准数据集上采用体素级推理方法，取得了 0.730 的测试 Dice 分数，优于现有深度学习模型。
• 消融实验证实，通过先进特征融合整合 U-Net 与 SwinUNETR 后，模型分割性能显著优于单一架构。
• 全面的预处理流程（包括重采样、偏置场校正与标准化）对模型性能至关重要，若移除该流程，Dice 分数将下降至 0.65。
• 该模型在不同大小的病灶上均表现出强泛化能力，可在同一受试者中成功检测到小、中、大及多发性病灶。
• 尽管性能优异，该模型在处理高度复杂或罕见病灶模式时仍存在困难，表明需要更大且更多样化的训练数据。
• 该模型的参数量极高（超过 1 亿参数），反映了其复杂架构，使其能够实现精细的特征学习，但也引发了在临床环境中计算可行性的担忧。