Automatic Tumor Segmentation via False Positive Reduction Network for Whole-Body Multi-Modality PET/CT Images

AutoPET consists of a training dataset of 1,014 PET/CT scans derived from 900 patients acquired at the University Hospital Tübingen, Germany [4]. All images are in NifTI format. There are 513 scans without lesions, and 188, 168, and 145 scans are histologically proven with malignant melanoma, lung cancer, and lymphoma, respectively. In addition, all patients have clinical reports including cancer diagnosis, sex, and age. A separate test dataset was not released to the public and was only used for the evaluation. The test dataset has a preliminary test set of 5 studies for self-evaluation and a final test set of 200 studies for final ranking. The preliminary test set was part of the final test set where 100 studies were from the same hospital as the training database (University Hospital Tübingen) and the other 100 scans were acquired from the University Hospital of the LMU in Munich with a similar acquisition protocol. The tumor regions for all the training and testing datasets were annotated by two radiologists with more than 5 years of experience in Hybrid Imaging and experience in machine learning research. In this study, we did not use any external dataset to build and train our model.

Multiple pre-processing steps were applied. Firstly, to compress the usage of GPU memory, all the PET/CT image volumes were cropped into a patch size of 224×224 in the axial plane. Then the images were set to the SUV range of [0, 14.25] for PET and HU range of [-800, 400] for CT, and further mapped to [0, 1] via min-max normalization. Finally, for PET images, the input slices were normalized with the mean and standard deviation values of the entire training dataset, such that to adjust all the regions of interest (ROIs) to a notionally common scale based on the metabolism intensity of tumor regions. For the CT images, the input slices were normalized with the mean and standard deviation values of the individual patient.

Our false positive reduction network consists of two main modules, as shown in Fig. 1: a global segmentation module and a local refinement module.

Specifically, a ResNet50 [8] encoder is firstly pre-trained on the training data (tumor presented) via contrastive learning [2] using concatenated 3-channel PET/CT images as the input. Two channels of the input images are set to be PET while the rest of 1 channel is assigned to CT. Then, a U-Net based decoder combined in together with the pre-trained encoder forms the global segmentation module, where a combined loss $\operatorname{Loss}_{GSM}$ is used for training. $\operatorname{Loss}_{GSM}$ consists of a dice loss [11] and a cross-entropy loss [1], which is defined as:

where $p_{i}\in\left[0,1\right]$ is pixels of the predicted tumor probability map, $g_{i}\in\left[0,1\right]$ is the pixels of the ground truth tumor mask (label), and the sums run over all available N pixels of the segmentation.

The global segmentation module can coarsely annotate the tumor lesion regions with a probability map at a threshold of 0.5. Both the probability map and the binary segmentation prediction (thresholded results) are further concatenated with the corresponding input PET/CT images, which are fed into the local refinement module.

The backbone of the local refinement module is a 2D U-Net [15] with 5-channel imaging data as input (i.e., paired PET/CT images, global tumor probability map, and global binary segmentation prediction). A pixel-wise mean squared error (MSE) loss and a cross-entropy loss [1] are combined as $\operatorname{Loss}_{LRM}$ and used to compare the predicted segmentation outputs with the tumor ground-truth mask such that the false positive regions can be removed from the global segmentation with:

Our method was implemented with PyTorch [12] using one NVIDIA GeForce GTX 2080Ti GPU. Our model was initialized using the approach presented by He et al [7], and an adaptive-moment estimation with decoupled weight decay (AdamW) [10] was used for network optimization. During the training phase, the batch size was set to 8 and the learning rate was set to 0.0001 using a cosine annealing schedule. Data augmentation techniques were in real-time to avoid overfitting. The used data augmentation techniques are random rotation (90°, 180°, or 270°) in the axial axis and randomly flipping in one of the two axes (sagittal and coronal).

Furthermore, we only used PET/CT slices with tumors for the global segmentation module, while an equal number of PET/CT slices without tumors were sampled and added to the training of the local refinement module. All the training was terminated when no further change is in the total loss. In our method, the total loss was generally stable after 160 epochs. Our results on the testing set were obtained using an ensemble of four models trained on different splits of the training set. Three models were built using 3-fold cross-validation, the other one was built from the 3D full resolution nnU-Net [9]. For the segmentation output, the testing prediction was a weighted sum between the nnU-Net output and the direct average result of the 3 cross-validation models, where the weights are empirically set to 0.35 and 0.65 respectively.