[Paper Review] Transfer Learning U-Net Deep Learning for Lung Ultrasound Segmentation

Dorothy J. Cheng, Edmund Y. Lam|arXiv (Cornell University)|Oct 5, 2021

COVID-19 diagnosis using AI19 references20 citations

TL;DR

This study proposes a transfer learning U-Net framework for precise lung ultrasound segmentation, comparing VGG16-based (V-Unet) and natural image pre-trained (X-Unet) approaches. X-Unet achieved superior visual accuracy and fewer artifacts despite a slightly lower Dice score, and partial-fine-tuning further improved performance, highlighting the synergy between transfer learning and data augmentation for limited medical datasets.

ABSTRACT

Transfer learning (TL) for medical image segmentation helps deep learning models achieve more accurate performances when there are scarce medical images. This study focuses on completing segmentation of the ribs from lung ultrasound images and finding the best TL technique with U-Net, a convolutional neural network for precise and fast image segmentation. Two approaches of TL were used, using a pre-trained VGG16 model to build the U-Net (V-Unet) and pre-training U-Net network with grayscale natural salient object dataset (X-Unet). Visual results and dice coefficients (DICE) of the models were compared. X-Unet showed more accurate and artifact-free visual performances on the actual mask prediction, despite its lower DICE than V-Unet. A partial-frozen network fine-tuning (FT) technique was also applied to X-Unet to compare results between different FT strategies, which FT all layers slightly outperformed freezing part of the network. The effect of dataset sizes was also evaluated, showing the importance of the combination between TL and data augmentation.

Motivation & Objective

• To improve segmentation accuracy of ribs in lung ultrasound images using transfer learning with U-Net.
• To evaluate the effectiveness of different pre-training strategies—VGG16 and natural salient object datasets—on medical image segmentation.
• To investigate fine-tuning strategies in transfer learning for U-Net in low-data regimes.
• To assess the impact of dataset size and data augmentation when combined with transfer learning.
• To determine the optimal balance between pre-training source and fine-tuning strategy for robust segmentation.

Proposed method

• Adapted U-Net architecture with encoder pre-trained on ImageNet via VGG16 (V-Unet).
• Constructed U-Net with encoder pre-trained on grayscale natural salient object dataset (X-Unet).
• Applied partial-fine-tuning by freezing varying numbers of encoder layers during training on lung ultrasound data.
• Used dice coefficient (DICE) and visual inspection to evaluate segmentation performance.
• Employed data augmentation techniques to enhance model generalization under limited training data.
• Compared model performance across different transfer learning techniques and fine-tuning strategies.

Experimental results

Research questions

• RQ1 Which pre-training source—VGG16 or natural salient object dataset—yields better segmentation performance for lung ultrasound images?
• RQ2 How does partial-fine-tuning compare to full fine-tuning in terms of segmentation accuracy and artifact reduction?
• RQ3 What is the impact of dataset size and data augmentation when combined with transfer learning in medical image segmentation?
• RQ4 Why does X-Unet produce more artifact-free predictions despite a lower Dice score than V-Unet?
• RQ5 Can transfer learning with a non-medical pre-trained model outperform medical-specific pre-training in low-data settings?

Key findings

• X-Unet produced visually more accurate and artifact-free segmentation masks compared to V-Unet, despite a lower Dice score.
• The partial-frozen fine-tuning strategy slightly outperformed freezing only a subset of layers, indicating broader fine-tuning benefits.
• V-Unet achieved a higher Dice coefficient than X-Unet, but visual quality was inferior due to increased artifacts.
• The combination of transfer learning and data augmentation significantly improved model performance, especially with limited training data.
• Pre-training on natural salient object datasets yielded better generalization and visual fidelity than ImageNet-based pre-training in this context.
• The study demonstrates that non-medical pre-training can be competitive with medical-specific pre-training when combined with effective fine-tuning and data augmentation.