[Paper Review] 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 AI16 references4 citations

TL;DR

This paper proposes Attention BCDU-Net, a U-Net-based architecture with bidirectional ConvLSTM and attention gates, for accurate lung and heart segmentation in chest X-rays using minimal training data. It achieves state-of-the-art performance—98.1% Dice score and 95.2% IoU on the JSRT dataset—while attaining 97.0% Dice and 92.2% IoU with only 10 training samples, significantly boosted by mixup data augmentation.

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.

Motivation & Objective

• To develop a deep learning model for automated lung and heart segmentation in chest X-rays that performs well even with limited training data.
• To improve upon existing state-of-the-art models in segmentation accuracy using a novel network architecture combining attention mechanisms and bidirectional feature fusion.
• To investigate the effectiveness of mixup data augmentation in low-data regimes for medical image segmentation.
• To analyze the peculiar training dynamics observed in small-data settings, particularly the delayed improvement in validation metrics despite stable training metrics.

Proposed method

• Proposes Attention BCDU-Net, a modified U-Net variant integrating bidirectional ConvLSTM layers for enhanced skip-connection feature fusion.
• Incorporates attention gates to focus the network on relevant features during upsampling, improving localization accuracy.
• Employs mixup data augmentation with hyperparameter δ to generate synthetic training samples by linear interpolation of input images and masks.
• Optimizes model performance through hyperparameter tuning of network depth and loss function, using Dice and IoU as evaluation metrics.
• Applies early stopping based on validation loss, though observes counterintuitive late improvements in validation metrics.
• Trains and evaluates the model on the JSRT dataset using training set sizes of 10, 20, and 209 samples.

Experimental results

Research questions

• RQ1 Can a deep learning model achieve state-of-the-art segmentation performance on chest X-rays with only 10–20 training samples?
• RQ2 How effective is mixup augmentation in improving segmentation performance when training data is extremely limited?
• RQ3 Why do validation Dice scores show a delayed improvement after training metrics have stabilized in small-data setups?
• RQ4 How does the integration of bidirectional ConvLSTM and attention gates enhance feature fusion in medical image segmentation?
• RQ5 Does the inherent similarity of medical image masks reduce the need for augmentation as training set size increases?

Key findings

• Attention BCDU-Net achieves a Dice score of 98.1±0.1% and IoU of 95.2±0.1% on the full JSRT dataset, surpassing previous state-of-the-art results.
• With only 10 training samples, the model reaches 97.0±0.8% Dice and 92.2±1.2% IoU, demonstrating strong performance on minimal data.
• With 20 training samples, the model achieves 97.3±0.5% Dice and 93.3±0.4% IoU, showing consistent improvement over smaller sets.
• Mixup augmentation with δ=0.2 yields a 4%+ gain in IoU score compared to no mixup in the 10-sample setup, highlighting its critical role in small-data generalization.
• The validation Dice score exhibits a sudden rise after epoch 400 in the 10-sample setup, despite stable training scores, indicating unexpected generalization dynamics.
• As training set size increases, the benefit of mixup diminishes, suggesting that larger datasets inherently reduce the need for augmentation due to mask similarity.