Data Augmentation and Transfer Learning to Improve Generalizability of an Automated Prostate Segmentation Model

Area of Science:

• Radiology
• Artificial Intelligence
• Medical Imaging

Background:

• Deep learning models in radiology can overfit, limiting their performance on external datasets.
• Developing generalizable prostate segmentation models is crucial for consistent clinical application.

Purpose of the Study:

• To develop a high-quality prostate segmentation model using transfer learning and data augmentation.
• To evaluate the model's performance across multiple independent datasets, including external centers.

Main Methods:

• A deep learning model combining 2D and 3D architectures was trained using transfer learning on 648 prostate MRI scans.
• A specialized data augmentation strategy addressed deformations, intensity variations, and image quality changes.
• Model performance was tested on five independent datasets, with and without fine-tuning.

Main Results:

• The model achieved high performance on internal validation (Dice scores: 93.1% whole prostate, 89.0% transition zone).
• Data augmentation alone improved performance by 2.2% (whole prostate) and 3.0% (transition zone) on test sets.
• Fine-tuning on test center data yielded the best results (Dice scores: 91.5% whole prostate, 89.7% transition zone).

Conclusions:

• Transfer learning is effective for creating high-performing prostate segmentation models.
• Data augmentation and fine-tuning enhance model generalizability to external datasets, improving diagnostic consistency.