[论文解读] Attention U-Net Based Adversarial Architectures for Chest X-ray Lung Segmentation

Gusztáv Gaál, Balázs Maga+1 more|arXiv (Cornell University)|Mar 23, 2020

COVID-19 diagnosis using AI参考文献 12被引用 128

一句话总结

本文提出一种带对抗评判器的 Attention U-Net（adv. ATTN U-Net），用于胸部 X 射线肺部分割，在 Dice 分数方面表现出色并且在跨数据集有较好泛化。它使用 Focal Tversky Loss 和结构校正的对抗框架来细化分割形状。

ABSTRACT

Chest X-ray is the most common test among medical imaging modalities. It is applied for detection and differentiation of, among others, lung cancer, tuberculosis, and pneumonia, the last with importance due to the COVID-19 disease. Integrating computer-aided detection methods into the radiologist diagnostic pipeline, greatly reduces the doctors' workload, increasing reliability and quantitative analysis. Here we present a novel deep learning approach for lung segmentation, a basic, but arduous task in the diagnostic pipeline. Our method uses state-of-the-art fully convolutional neural networks in conjunction with an adversarial critic model. It generalized well to CXR images of unseen datasets with different patient profiles, achieving a final DSC of 97.5% on the JSRT dataset.

研究动机与目标

• 在胸部 X 光影像上使用 Attention U-Net 架构提升肺部分割的准确性。
• 结合对抗引导，在不大量预选 ROI 的情况下生成更逼真的分割掩膜。
• 评估该方法在多个公开胸部 X 光数据集（带有真实掩膜）上的泛化能力。
• 探索针对医学影像分割的有效损失函数，特别是针对不平衡或小 ROI 的情况。

提出的方法

• 采用 Attention U-Net 架构，在不进行外部 ROI 预选的情况下聚焦相关局部特征。
• 集成一个受 SCAN 启发的对抗评判器模型，鼓励生成更真实的分割掩膜。
• 在对抗框架中用 Focal Tversky Loss 替代标准损失，以应对类别不平衡和小的肺区。
• 采用最小-最大目标训练：S（分割器）和 D（对抗评判器），平衡参数 lambda（设为 0.1）控制像素级损失与对抗损失。
• 对生成器预训练 50 个 epoch，然后进行对抗训练：每 5 次生成器更新更新 1 次判别器。
• 数据预处理使用 CLAHE，调整为 512x512，并标准化到 [-1,1] 以实现鲁棒训练。

实验结果

研究问题

• RQ1 在 U-Net 中的注意力门能否在不进行 ROI 预选的情况下提升肺部分割？
• RQ2 对抗评判器是否比强大的 Attention U-Net 基线进一步提升分割掩膜的真实度？
• RQ3 在拥有不同患者特征和成像条件的 JSRT、Montgomery 和 Shenzhen 数据集上，该模型的泛化能力如何？
• RQ4 哪些损失函数最适合处理胸部 X 光分割中不平衡和小目标区域？

主要发现

• 所提出的方法在 JSRT 数据集上达到 Dice 分数 0.975（摘要中报告的最终 DSC 为 97.5%）
• 在 JSRT、SCAN 与我们的方法（Adv. ATTN）的比较中 JSRT：97.3±0.8%（SCAN），96.3±0.7%（ATTN U-Net），97.6±0.5%（Adv. ATTN）
• 综合所有数据集时，我们的方法（Adv. ATTN）达到 96.2±0.4%（All），而 SC AN 为 95.8±0.4%
• All 数据集汇总显示 All / JSRT：Adv. ATTN 为 97.8±0.6%，ATTN U-Net 为 96.6±0.6%
• 对抗方案提供了形状细化和适度的 Dice 提升，且预处理有助于强跨数据集泛化。
• 训练比某些前期工作更快，总共 40 个 epoch，且采用稳定的 SGD 优化策略。