[Paper Review] Lung Segmentation from Chest X-rays using Variational Data Imputation

Raghavendra Selvan, Erik B. Dam+5 more|arXiv (Cornell University)|May 20, 2020

COVID-19 diagnosis using AI32 references46 citations

TL;DR

The paper proposes a U-net style lung segmentation model augmented with a variational encoder to impute missing data in high-opacity chest X-rays, trained only on normal CXRs, and shows improved segmentation using specialized data augmentations that simulate opacification. The best results combine block masking and diffused noise augmentations with variational data imputation.

ABSTRACT

Pulmonary opacification is the inflammation in the lungs caused by many respiratory ailments, including the novel corona virus disease 2019 (COVID-19). Chest X-rays (CXRs) with such opacifications render regions of lungs imperceptible, making it difficult to perform automated image analysis on them. In this work, we focus on segmenting lungs from such abnormal CXRs as part of a pipeline aimed at automated risk scoring of COVID-19 from CXRs. We treat the high opacity regions as missing data and present a modified CNN-based image segmentation network that utilizes a deep generative model for data imputation. We train this model on normal CXRs with extensive data augmentation and demonstrate the usefulness of this model to extend to cases with extreme abnormalities.

Motivation & Objective

• Motivate robust lung segmentation in CXRs with extreme opacification where manual segmentation is difficult.
• Develop a segmentation model that can generalize from normal CXRs to abnormal, highly opaque CXRs by imputing missing data.
• Leverage a variational encoder to augment a U-net segmentation network for improved performance on challenging CXRs.
• Explore data augmentation strategies that simulate missing regions and increased opacity to train the imputation model.

Proposed method

• Use a U-net-like segmentation network for lung masks with an added variational encoder to perform data imputation.
• Introduce a latent variable z sampled from q_phi(z|x) and concatenate z with the U-net encoder output before decoding.
• Train with a segmentation loss plus KL divergence between q_phi(z|x) and p(z).
• Share the VAE decoder with the segmentation decoder to jointly predict the segmentation and impute data.
• Experiment with data augmentations (standard, block masking, diffused noise) to simulate missing data and opacification.
• Evaluate on normal-training-set-derived models and test on high-opacity CXRs with expert-provided masks.

Experimental results

Research questions

• RQ1 Can a variational data imputation module trained on normal CXRs improve segmentation on high-opacity CXRs lacking labeled abnormal examples?
• RQ2 Do advanced data augmentations that simulate occluded regions and increased opacity boost performance of a lung segmentation model under severe opacification?
• RQ3 Does sharing the decoder between the VAE imputation and segmentation networks yield better lung masks than a baseline U-net?
• RQ4 What is the impact of combining block masking with diffused noise augmentations on Dice and accuracy for lung segmentation?
• RQ5 How well does the proposed method generalize from normal to highly abnormal CXRs for COVID-19 related analysis?

Key findings

• The proposed model with variational data imputation and block masking plus diffused noise achieves the best Dice of 0.8503 (±0.07) and Accuracy of 0.8815 (±0.11).
• Baseline with standard augmentation yields Dice 0.7335 (±0.17) and Accuracy 0.8449 (±0.09).
• Baseline with block augmentation yields Dice 0.7563 (±0.15) and Accuracy 0.8522 (±0.09).
• Baseline with diffused augmentation yields Dice 0.7757 (±0.15) and Accuracy 0.8654 (±0.10).
• Proposed with standard augmentation yields Dice 0.7204 (±0.18) and Accuracy 0.8392 (±0.10).
• Proposed with block augmentation yields Dice 0.7688 (±0.17) and Accuracy 0.8552 (±0.10).