[论文解读] Chest X-ray lung and heart segmentation based on minimal training sets

Balázs Maga|arXiv (Cornell University)|Jan 20, 2021

COVID-19 diagnosis using AI参考文献 16被引用 4

一句话总结

该论文提出了一种基于注意力机制的BCDU-Net，一种结合双向ConvLSTM与注意力门的U-Net变体架构，旨在使用极少的训练数据实现胸部X光片中肺部和心脏的精确分割。该方法在JSRT数据集上实现了最先进性能——Dice分数达98.1%，IoU达95.2%；在仅使用10张训练样本时，Dice分数达到97.0%，IoU达92.2%，性能显著提升，主要得益于mixup数据增强技术。

ABSTRACT

As the COVID-19 pandemic aggravated the excessive workload of doctors globally, the demand for computer aided methods in medical imaging analysis increased even further. Such tools can result in more robust diagnostic pipelines which are less prone to human errors. In our paper, we present a deep neural network to which we refer to as Attention BCDU-Net, and apply it to the task of lung and heart segmentation from chest X-ray (CXR) images, a basic but ardous step in the diagnostic pipeline, for instance for the detection of cardiomegaly. We show that the fine-tuned model exceeds previous state-of-the-art results, reaching $98.1\pm 0.1\%$ Dice score and $95.2\pm 0.1\%$ IoU score on the dataset of Japanese Society of Radiological Technology (JSRT). Besides that, we demonstrate the relative simplicity of the task by attaining surprisingly strong results with training sets of size 10 and 20: in terms of Dice score, $97.0\pm 0.8\%$ and $97.3\pm 0.5$, respectively, while in terms of IoU score, $92.2\pm 1.2\%$ and $93.3\pm 0.4\%$, respectively. To achieve these scores, we capitalize on the mixup augmentation technique, which yields a remarkable gain above $4\%$ IoU score in the size 10 setup.

研究动机与目标

• 开发一种深度学习模型，用于在训练数据有限的情况下实现胸部X光片中肺部和心脏的自动分割。
• 通过结合注意力机制与双向特征融合的新型网络架构，超越现有最先进模型的分割精度。
• 探究mixup数据增强在极小数据场景下对医学图像分割性能的影响。
• 分析在小样本设置下观察到的特殊训练动态，特别是尽管训练指标稳定，验证指标却出现延迟提升的现象。

提出的方法

• 提出Attention BCDU-Net，一种改进的U-Net变体，通过引入双向ConvLSTM层增强跳跃连接的特征融合。
• 引入注意力门机制，使网络在上采样过程中聚焦于相关特征，提升定位精度。
• 采用mixup数据增强技术，通过超参数δ对输入图像与掩码进行线性插值，生成合成训练样本。
• 通过调节网络深度与损失函数的超参数优化模型性能，使用Dice与IoU作为评估指标。
• 基于验证损失应用早停策略，但观察到验证指标在训练指标稳定后出现反常的延迟提升。
• 在JSRT数据集上，使用10、20和209张训练样本对模型进行训练与评估。

实验结果

研究问题

• RQ1 在仅使用10–20张训练样本的情况下，深度学习模型能否在胸部X光片分割任务中达到最先进性能？
• RQ2 在训练数据极度有限时，mixup增强技术在提升分割性能方面的有效性如何？
• RQ3 为何在小样本设置下，验证Dice分数在训练指标稳定后出现延迟提升？
• RQ4 双向ConvLSTM与注意力门的结合如何增强医学图像分割中的特征融合？
• RQ5 随着训练集规模增大，由于医学图像掩码本身具有较高相似性，是否降低了对数据增强的需求？

主要发现

• Attention BCDU-Net在完整JSRT数据集上达到98.1±0.1%的Dice分数与95.2±0.1%的IoU，超越了先前最先进结果。
• 仅使用10张训练样本时，模型Dice分数达97.0±0.8%，IoU达92.2±1.2%，展现出在极小数据下的强大性能。
• 使用20张训练样本时，模型Dice分数达97.3±0.5%，IoU达93.3±0.4%，相较于更小数据集持续提升。
• 在10样本设置下，δ=0.2的mixup增强使IoU分数相比无mixup提升4%以上，凸显其在小样本泛化中的关键作用。
• 在10样本设置中，验证Dice分数在第400轮后出现突然上升，尽管训练分数保持稳定，表明存在出人意料的泛化动态。
• 随着训练集规模增大，mixup的增益逐渐减弱，表明大规模数据本身因掩码相似性而降低了对增强的需求。