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

The training data consisted of four cohorts, namely scans presenting lymphoma (145 cases), lung cancer (168 cases), melanoma (188 cases), and negative control patients (513 cases). Each of these 4 cohorts was individually split into 5 folds. The final 5 training folds were created by segregating together the images belonging to fold $f$ across the four cohorts, where $f=\{0,1,2,3,4\}$. No other dataset (public or private) was used in this work.

The CT images were first downsampled to match the coordinates of their corresponding PET images. The PET intensity values in units of Bq/ml were decay-corrected and converted to SUV. During training, we employed a series of non-randomized and randomized transforms to augment the input to the network. The non- randomized transforms included (i) clipping CT intensities in the range of [-1024, 1024] HU (ii) min-max normalization of clipped CT intensity to span the interval [0, 1], (iii) cropping the region outside the body in PET, CT, and mask images using a 3D bounding box, and (iv) resampling the PET, CT, and mask images to an isotropic voxel spacing of (2.0 mm, 2.0 mm, 2.0 mm) via bilinear interpolation for PET and CT images and nearest-neighbor interpolation for mask images.

On the other hand, the randomized transforms were called at the start of every epoch. These included (i) random spatial cropping of cubic patches of dimensions (192, 192, 192) from the images, (ii) 3D translations in the range (0, 10) voxels along all three directions, (iii) axial rotations by angle $\theta\in(-\pi/12,\pi/12)$, (iv) random scaling by a factor of 1.1 in all three directions, (v) 3D elastic deformations using a Gaussian kernel with standard deviation and offsets on the grid uniformly sampled from (0, 1), (vi) Gamma correction with $\gamma\in(0.7,1.5)$, and (vii) addition of random Gaussian noise with $\mu=0$ and $\sigma=1$. Finally, the augmented PET and CT patches were concatenated along the channel dimension to construct the final input for the network.

We used a 3D Residual UNet [10], adapted from the MONAI library [11]. The network consisted of 2 input channels, 2 output channels, and 5 layers of encoder and decoder (with 2 residual units per block) paths with skip-connections. The data in the encoder was downsampled using strided convolutions, while the decoder unsampled using transpose strided convolutions. The number of channels in the encoder part from the top-most layer to the bottleneck were 32, 64, 128, 256, and 512. PReLU was used as the activation function within the network. The network consisted of 19,223,525 trainable parameters.

We employed the binary Generalized Dice Focal Loss $\mathcal{L}_{\text{GDFL}}=\mathcal{L}_{\text{GDL}}+\mathcal{L}_{\text{FL}}$, where $\mathcal{L}_{\text{GDL}}$ is the Generalized Dice Loss [12] and $\mathcal{L}_{\text{FL}}$ is the Focal Loss [13]. The Generalized Dice Loss $\mathcal{L}_{\text{GDL}}$ is given by,

where $p_{ilj}$ and $g_{ilj}$ are values of the $j^{th}$ voxel of the $i^{th}$ cropped patch of the predicted and ground truth segmentation masks with class $l\in\{0,1\}$ respectively in a mini-batch size $n_{b}$ of the cropped patches and $N^{3}$ represents the total number of voxels in the cropped cubic patch of size $(N,N,N)$, where $N=192$. Here, $w_{l}=1/(\sum_{j=1}^{N^{3}}g_{ilj})^{2}$ represents the weight given to class $l$. The mini-batch size was set to $n_{b}=4$. Small constants $\epsilon=\eta=10^{-5}$ were added to the numerator and denominator, respectively to ensure numerical stability during training. The Focal Loss $\mathcal{L}_{\text{FL}}$ is given by,

where, $v_{0}=1$ and $v_{1}=100$ are the focal weights of the two classes, $\sigma(x)=1/(1+\text{exp}(-x))$ is the sigmoid function, and $\gamma=2$ is the focal loss parameter that suppresses the loss for the class that is easy to classify.

$\mathcal{L}_{\text{GDFL}}$ was optimized using the Adam optimizer. Cosine annealing scheduler was used to decrease the learning rate from $1\times 10^{-3}$ to $0$ in 300 epochs. The loss for an epoch was computed by averaging the $\mathcal{L}_{\text{GDFL}}$ over all batches. The model with the highest mean Dice similarity coefficient (DSC) on the validation fold $f$ was chosen for further evaluation, for all $f\in\{0,1,2,3,4\}$.

For the images in the validation set, we employed only the non-randomized transforms. The prediction was made directly on the 2-channel (PET and CT) whole-body images using a sliding-window technique with a window of dimensions $(192,192,192)$ and overlap=0.5. For final testing, the outputs of the 5 best models (obtained from 5-folds training) were ensembled via average and weighted average ensembling to generate the output mask. For the weighted average ensembling, the weights were chosen as the value of mean DSC of the respective validation fold. The final output masks were resampled to the coordinates of the original ground truth masks for computing the evaluation metrics.

The challenge employed three evaluation metrics, namely the mean DSC, mean false positive volume (FPV) and mean false negative volume (FNV). For a foreground ground truth mask $G$ containing $L_{g}$ disconnected foreground segments (or lesions) $\{G_{1},G_{2},...,G_{L_{g}}\}$ and the corresponding predicted foreground mask $P$ with $L_{p}$ disconnected foreground segments $\{P_{1},P_{2},...,P_{L_{p}}\}$, these metrics are defined as,

where $\delta(x):=1$ for $x=0$ and $\delta(x):=0$ otherwise. $v_{g}$ and $v_{p}$ represent the voxel volumes (in ml) for ground truth and predicted mask, respectively (with $v_{p}=v_{g}$ since the predicted mask was resampled to the original ground truth coordinates). The submitted algorithms were ranked separately for each of the three metrics and the final ranking was determined based on the formula: $0.5\times\text{rank}_{\text{DSC}}+0.25\times\text{rank}_{\text{FPV}}+0.25\times\text{rank}_{\text{FNV}}$. The function definitions for these metrics were obtained from the challenge GitHub page and can be accessed via this link.