[Paper Review] Generalized Dice Focal Loss trained 3D Residual UNet for Automated Lesion Segmentation in Whole-Body FDG PET/CT Images

Shadab Ahamed, Arman Rahmim|arXiv (Cornell University)|Sep 24, 2023

Medical Imaging Techniques and ApplicationsMedicine3 citations

TL;DR

This paper proposes a 3D Residual UNet model trained with Generalized Dice Focal Loss for automated lesion segmentation in whole-body FDG-PET/CT images. Using 5-fold cross-validation and ensemble inference on the AutoPET 2023 dataset, the method achieved a Dice score of 0.5417, with FPV and FNV of 0.82 ml and 0.25 ml respectively on the preliminary test set, demonstrating robust performance on diverse, low-uptake, and anatomically complex lesions.

ABSTRACT

Automated segmentation of cancerous lesions in PET/CT images is a vital initial task for quantitative analysis. However, it is often challenging to train deep learning-based segmentation methods to high degree of accuracy due to the diversity of lesions in terms of their shapes, sizes, and radiotracer uptake levels. These lesions can be found in various parts of the body, often close to healthy organs that also show significant uptake. Consequently, developing a comprehensive PET/CT lesion segmentation model is a demanding endeavor for routine quantitative image analysis. In this work, we train a 3D Residual UNet using Generalized Dice Focal Loss function on the AutoPET challenge 2023 training dataset. We develop our models in a 5-fold cross-validation setting and ensemble the five models via average and weighted-average ensembling. On the preliminary test phase, the average ensemble achieved a Dice similarity coefficient (DSC), false-positive volume (FPV) and false negative volume (FNV) of 0.5417, 0.8261 ml, and 0.2538 ml, respectively, while the weighted-average ensemble achieved 0.5417, 0.8186 ml, and 0.2538 ml, respectively. Our algorithm can be accessed via this link: https://github.com/ahxmeds/autosegnet.

Motivation & Objective

• To address the challenge of accurate automated lesion segmentation in whole-body FDG-PET/CT scans, which are complicated by lesion diversity, low uptake, and physiological tracer uptake in organs.
• To overcome limitations of traditional thresholding and deep learning models trained on small, private datasets by leveraging the publicly available AutoPET 2023 challenge dataset.
• To improve segmentation robustness in the presence of class imbalance and noisy predictions using a novel loss function combining Generalized Dice and Focal Loss.
• To enhance model generalization and performance through 5-fold cross-validation and model ensembling strategies (average and weighted-average).

Proposed method

• Trained a 3D Residual UNet with two input channels (PET and CT) and two output channels (background and lesion) using the MONAI library.
• Applied a comprehensive preprocessing pipeline including CT intensity clipping, SUV conversion, body cropping, and isotropic resampling to 2.0 mm isotropic voxels.
• Employed extensive data augmentation including random spatial cropping, 3D translations, axial rotations, scaling, elastic deformations, gamma correction, and Gaussian noise.
• Formulated a Generalized Dice Focal Loss (GDFL) that combines class-weighted Generalized Dice Loss with Focal Loss to suppress easy examples and improve segmentation of rare or low-uptake lesions.
• Optimized the GDFL using Adam with cosine annealing learning rate scheduling over 300 epochs, with model selection based on highest validation DSC.
• Performed inference using a sliding-window approach with 192×192×192 patches and 50% overlap, followed by model ensembling via average and weighted-average methods using validation DSC as weights.

Experimental results

Research questions

• RQ1 Can a 3D Residual UNet trained with Generalized Dice Focal Loss achieve robust lesion segmentation across diverse cancer types and lesion characteristics in whole-body PET/CT scans?
• RQ2 How does the combination of Generalized Dice and Focal Loss improve segmentation performance compared to standard Dice loss in the presence of class imbalance and low-uptake lesions?
• RQ3 To what extent does 5-fold cross-validation and model ensembling enhance generalization and performance on the AutoPET 2023 test set?
• RQ4 Can data augmentation strategies effectively improve model robustness to anatomical and intensity variations in clinical PET/CT data?

Key findings

• The average ensemble achieved a Dice similarity coefficient (DSC) of 0.5417, FPV of 0.8261 ml, and FNV of 0.2538 ml on the AutoPET 2023 preliminary test set.
• The weighted-average ensemble achieved identical DSC (0.5417) but slightly improved FPV (0.8186 ml) and identical FNV (0.2538 ml), indicating better false positive control.
• On the 5-fold cross-validation, the mean DSC ranged from 0.61 to 0.64 across folds, with median DSC values between 0.71 and 0.74, indicating strong fold-wise consistency.
• The model demonstrated strong performance on low-uptake and small lesions, as evidenced by the low FNV (0.2538 ml), suggesting effective detection of missed lesions.
• The highest-performing fold (fold 0) achieved a DSC of 0.61 ± 0.26, with a median of 0.71, and FPV of 4.16 ± 7.79 ml, indicating some variability in false positive control across folds.
• The model was publicly released via GitHub, enabling reproducibility and further benchmarking in the medical imaging community.