A U-Net Convolutional Neural Network with Multiclass Dice Loss for Automated Segmentation of Tumors and Lymph Nodes from Head and Neck Cancer PET/CT Images

Abstract

We implemented a 2D U-Net model with an ImageNet-pretrained ResNet50 encoder for performing segmentation of primary tumors (GTVp) and metastatic lymph nodes (GTVn) from PET/CT images provided by the HEad and neCK TumOR segmentation challenge (HECKTOR) 2022. We utilized a multiclass Dice Loss for model training which was minimized using the AMSGrad variant of the Adam algorithm optimizer. We trained our 2D models on the axial slices of the images in a 5-fold cross-validation setting and stacked the 2D predictions axially to obtain the predicted 3D segmentation masks. We obtained mean aggregate Dice similarity coefficients (mean DSC$_{\text {agg}}$) of 0.6865, 0.6689, 0.6768, 0.6792, and 0.6726 on the 5 validation sets respectively. The model with the best performance on the validation set (validation split 1) was chosen for evaluating segmentation masks on the test set for submission to the challenge. Our model achieved a mean DSC$_{\text {agg}}$ = 0.6345 on the test set, with DSC$_{\text {agg}}$(GTVp) = 0.6955 and DSC$_{\text {agg}}$(GTVn) = 0.5734. The implementation can be found under our Github repository.