[Paper Review] 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 AIMedicine12 references128 citations

TL;DR

The paper proposes an Attention U-Net with an adversarial critic (adv. ATTN U-Net) for chest X-ray lung segmentation, achieving high Dice scores and good generalization across datasets. It uses Focal Tversky Loss and a structure-correcting adversarial framework to refine segmentation shapes.

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.

Motivation & Objective

• Improve lung segmentation accuracy on chest X-rays using an Attention U-Net architecture.
• Incorporate adversarial guidance to produce more realistic segmentation masks without excessive preselection of ROIs.
• Evaluate generalization of the method across multiple public chest X-ray datasets with ground-truth masks.
• Explore effective loss functions tailored to medical image segmentation, particularly for imbalanced or small ROIs.

Proposed method

• Adopt an Attention U-Net architecture to focus on relevant local features without external ROI preselection.
• Integrate an adversarial critic model inspired by SCAN to encourage realistic segmentation masks.
• Replace the standard loss in the adversarial framework with Focal Tversky Loss to address class imbalance and small lung regions.
• Train with a min-max objective: S (segmentor) and D (critic) with a balance parameter lambda (set to 0.1) controlling pixel-wise vs adversarial losses.
• Pretrain the generator for 50 epochs and then train adversarially with 1 critic update per 5 generator updates.
• Preprocess data with CLAHE, resize to 512x512, and normalize to [-1,1] for robust training.

Experimental results

Research questions

• RQ1 Can Attention Gates within U-Net improve lung segmentation without ROI preselection?
• RQ2 Does an adversarial critic improve the realism of segmentation masks beyond a strong Attention U-Net baseline?
• RQ3 How does the model generalize across JSRT, Montgomery, and Shenzhen datasets with different patient profiles and imaging conditions?
• RQ4 What loss functions best handle imbalanced and small target regions in chest X-ray segmentation?

Key findings

• The proposed method achieved a Dice score of 0.975 on the JSRT dataset (final DSC 97.5% reported in abstract).
• On JSRT, SCAN, and Ours (Adv. ATTN) comparisons show JSRT: 97.3±0.8% (SCAN), 96.3±0.7% (ATTN U-Net), 97.6±0.5% (Adv. ATTN).
• When considering all datasets, Ours (Adv. ATTN) achieved 96.2±0.4% (All) vs 95.8±0.4% (SC AN).
• All datasets combined show All / JSRT: 97.8±0.6% for Adv. ATTN vs 96.6±0.6% for ATTN U-Net.
• The adversarial scheme provided shape refinements and modest Dice improvements, with preprocessing contributing to strong cross-dataset generalization.
• Training is faster than some prior work, with 40 epochs overall and a stable SGD optimization regime.