Generalizability issues with deep learning models in medicine and their potential solutions: illustrated with cone-beam computed tomography (CBCT) to computed tomography (CT) image conversion

This study used the CycleGAN architecture we developed previously (figure 4). It consists of two discriminators (DiscriminatorA and DiscriminatorB) and two generators (GeneratorA and GeneratorB). GeneratorA is to generate sCT from CBCT and GeneratorB is to generate synthesized CBCT (sCBCT) from CT, while DiscriminatorA is to distinguish between sCT and CT and DiscriminatorB is to distinguish between sCBCT and CBCT. Generators and discriminators are trained together and work against each other to minimize objectives. Besides typical adversarial loss in every GAN architecture, CycleGAN introduces the concept of cycle consistency loss to avoid paired dataset for training. A CBCT image going through GeneratorA and then GeneratorB is supposed to generate an image equal to the CBCT image (${\text{GeneratorB}}\left[ {{\text{GeneratorA}}\left[ {{\text{CBCT}}} \right]} \right] = {\text{CBCT}}$), and vice versa (${\text{GeneratorA}}\left[ {{\text{GeneratorB}}\left[ {{\text{CT}}} \right]} \right] = {\text{CT}}$). This cycle consistency helps to constrain the mapping from one domain to another. In this experiment, we use the CycleGAN architecture with U-Net for generators and pacthGAN for discriminators to convert CBCT images to sCT images with less noise and other artifacts. For more details, the reader is kindly referred to (Liang et al 2019).

Zoom In Zoom Out Reset image size

Since the data distribution could be different for different anatomical sites, scanning protocols, and vendors' scanners, it is unknown whether a model trained on one dataset works for another dataset. To investigate the generalizability of the CycleGAN model trained on different CBCT datasets, we split the seven datasets into source dataset (H&N1 and H&N2) and target dataset (Prostate1, Prostate2, Prostate3, Cervix and Pancreas) to mimic a situation where CBCT scans come from different clinical environments. A flowchart of the generalizability experiments is shown in figure 5(a). The source dataset consists of the H&N1 and H&N2 datasets, which represented data collected from one circumstance. The model trained on the source dataset is called source model. We then directly applied the source model to five target datasets without re-training and compared the results with that from the source model applied in source datasets to illustrate the generalizability problem that one might encounter. The number of patients used for training the source model and the number of testing patients for each dataset can be referred to table 1.

Zoom In Zoom Out Reset image size

To solve the generalizability problem mentioned above, we investigated three potential solutions: target model, combined model, and adapted model, shown in figures 5(b)–(d) respectively. The target model is trained on a target dataset starting from scratch, so each target dataset has its own unique target model. The combined model is trained on the combined source and target datasets starting from scratch. The source model trained on the H&N1 and H&N2 datasets was fine-tuned on the Prostate1, Prostate2, Prostate3, Cervix, and Pancreas datasets separately to get an adapted model for each target dataset. No layers in the architecture were frozen, and all layers were updated in the fine-tuning process. This strategy is commonly used in transfer learning for adapting an old model to a new domain when the training data from the new domain is limited.

The number of training, validation, and testing patients for each dataset used in these experiments are shown in table 1 except Prostate2 training dataset. For fair comparison, we randomly selected 15 out of 39 patients from Prostate2 dataset for training in experiments 2–4. For all three solutions, all hyper-parameters—including layer architecture, batch size and number of training iterations—were kept the same, except the learning rate. We used the grid search method to find the optimal learning rate for each target dataset (Reed and Marksii 1999, Bengio 2012, Goodfellow et al 2016). The search range of the learning rate was from 9 × 10⁻³ to 2 × 10⁻⁶.

The performance of the target, combined, and adapted models might depend on the size of the training dataset in the target domain. We conducted an experiment using Prostate2 dataset to investigate this effect, illustrated in figure 5(e). We trained the target, combined, and adapted models by randomly picking 5, 10, 15, 27 and 39 patients from the Prostate2 training dataset. In the target and adapted models, the total number of patients for training was 5, 10, 15, 27, and 39 separately. In the combined model, because the target dataset is combined with the source dataset, the total number of patients for training was 99, 104, 109, 121, and 133 respectively. Eleven patients were used to test all the models, the same as in sections 3.1 and 3.2. We repeated the whole process including training and testing 10 times for statistical analysis. We then analyzed the model performance in correlation with training data size for each model.

In addition to visually evaluating the image quality, we also evaluated the model performance by using mean absolute error (MAE), root mean square error (RMSE), structural similarity index (SSIM), and signal-to-noise ratio (SNR), the four most widely used metrics (Wang et al 2004, Al-Obaidi 2015) for measuring the similarity between generated sCT images and the dCT images. Since lack of real ground truth, dCT images serve as the gold standard in this experiment to evaluate image quality by the four similarity measure metrics in the following sections. Assume $sCT\left( {x,y} \right)$ to be the HU value of the sCT images generated from the CycleGAN models and $dCT\left( {x,y} \right)$ to be the HU value of the deformably registered reference CT images. All the images have the same size [512 512].

MAE measures the difference between two images, as follows:

$$\begin{equation}MAE = \frac{1}{{{n_x}{n_y}}}\mathop \sum \limits_0^{{n_x} - 1} \mathop \sum \limits_0^{{n_y} - 1} \left| {dCT\left( {x,y} \right) - sCT\left( {x,y} \right)} \right|.\end{equation} \tag{ 1 }$$

RMSE also measures the average errors between two images, but it gives relatively high weight to large errors, as follows:

$$\begin{equation}RMSE = \sqrt {\frac{1}{{{n_x}{n_y}}}\mathop \sum \limits_0^{{n_x} - 1} \mathop \sum \limits_0^{{n_y} - 1} {{\left( {dCT\left( {x,y} \right) - sCT\left( {x,y} \right)} \right)}^2}.} \end{equation} \tag{ 2 }$$

SSIM measures the similarity between two images based on the human visual system, (Wang et al 2003) that is

$$\begin{equation}SSIM = \frac{{\left[ {\left( {2{\mu _1}{\mu _2} + {c_1}} \right)\left( {2{\sigma _{12}} + {c_2}} \right)} \right]}}{{\left[ {\left( {\mu _1^2 + \mu _2^2 + {c_1}} \right)\left( {\sigma _1^2 + \sigma _2^2 + {c_2}} \right)} \right]}},\end{equation} \tag{ 3 }$$

where ${\mu _1}$ is the average pixel value of the dCT image, ${\mu _2}$ is the average pixel value of the sCT image, $\sigma _1^2$ is the pixel variance of the dCT, $\sigma _2^2$ is the pixel variance of the sCT, ${\sigma _{12}}$ is the pixel covariance of the dCT and sCT, ${c_1} = {\left( {{k_1}L} \right)^2}$, ${c_2} = {\left( {{k_2}L} \right)^2}$, L is the dynamic range of the pixel values (L = 4095 in our case), ${k_1} = 0.01$ and ${k_2} = 0.03$. SSIM ranges from 0 to 1, and higher values indicate greater similarity between two images.

SNR compares the level of a desired signal to the level of background noise and is defined as

$$\begin{equation}SNR = 10lo{g_{10}}\left[\frac{{\mathop \sum \nolimits_0^{{n_x} - 1} \mathop \sum \nolimits_0^{{n_y} - 1} {{\left( {dCT\left( {x,y} \right)} \right)}^2}}}{{\mathop \sum \nolimits_0^{{n_x} - 1} \mathop \sum \nolimits_0^{{n_y} - 1} {{\left( {dCT\left( {x,y} \right) - sCT\left( {x,y} \right)} \right)}^2}}}\right].\end{equation} \tag{ 4 }$$

MAE and RMSE are given in HU. SSIM is unitless. SNR is given in dB.

The quantitative evaluation results including MAE, RMSE, SSIM, and SNR of the sCT generated by the source model on source datasets are shown in table 2 and serve as a baseline of the performance of the source model on target datasets. The dispersion of evaluation results from the source model on source datasets are expressed in boxplots shown in figure 6. The generated sCT images and their corresponding CBCT and dCT images from the H&N1 and H&N2 testing datasets are shown in supplementary figures 1 and 2 (available online at stacks.iop.org/MLST/2/015007/mmedia) for visual evaluation. We can see that the trained CycleGAN model can generate sCT images that are very similar to the reference dCT images and much better than the CBCT images when the model is applied to the same dataset that it was trained on.

Zoom In Zoom Out Reset image size

We quantitatively evaluated the performance of the source, target, combined, and adapted models in terms of MAE, RMSE, SSIM and SNR methods for each of the following target datasets: Prostate1, Prostate2, Prostate3, Cervix, and Pancreas. The geometric mean and SD of the similarity scores for every case and its relative improvement from source model are shown in table 3. Two-tailed Wilcoxon signed-rank tests were used for pairwise patient samples comparison due to non-normally distributed data and small sample size. The calculated p-values are shown with boxplots in figures 7–8, and supplementary figures 8–10. The sCT images generated by the source, target, combined, and adapted models and their corresponding CBCT and dCT images for each target dataset are shown in supplementary figures 3–7 for visual evaluation.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

In the Prostate1 dataset, which comes from Varian OBI scanners, like the source datasets (H&N1 and H&N2), the sCT images generated by the source model tend to have higher MAE geometric mean of 21.9 (SD 7.6) compared to the target model with 20.1 (5.7), the combined model with 19.3 (5.3), and the adapted model with 19.7 (5.8). In SSIM evaluation, they tend to have lower SSIM geometric mean of 88.8 (5.5) compared to the target model with 91.2 (53.0), the combined model with 91.9 (2.6), and the adapted model with 91.6 (2.9). The p-values calculated between source vs target model, source vs combined model, and source vs adapted model in both MAE and SSIM evaluations are less than 0.05. However in RMSE and SNR evaluations, no statistically significant difference were found in these group comparisons. This indicates that the source model performed reasonably well on this target dataset, and the target, combined and adapted models can only improve a little. Among the three updated models, no statistically significant difference were found between adapted vs target model in RMSE and SNR evaluations, and between adapted vs combined model in all four similarity measures. These results suggest that the three updated models performed comparably. These results are compatible with supplementary figure 3 for visual evaluation. Thus, when applying the source model to a dataset which it has never seen before (different anatomical sites), but coming from the same vendor's scanners, the source model can generate good quality sCT images from CBCT, while the target model, combined model, and adapted model slightly improve upon this performance by 8.2%, 12.0%, and 10.3% respectively, if evaluated in MAE similarity measure.

In Prostate2 dataset, which come from Elekta XVI (Versa) scanners, the source model has the highest MAE and RMSE geometric means of 46.8 (SD 5.3) and 112.4 (8.2), and the lowest SSIM and SNR geometric means of 75.1 (3.7) and 17.4 (1.0) among all the models. Statistical tests between source vs target model, source vs combined model, and source vs adapted model in four similarity measures all show p-values less than 0.05. These results indicate that the source model performed much worse in this target dataset (Prostate2) than in the Prostate1 target dataset. In contrast, the adapted model has the lowest MAE and RMSE geometric means of 22.4 (4.7), and highest SSIM and SNR geometric means of 88.4 (4.4) and 21.5 (1.3) compared to the target model and combined model with MAE of 28.4 (4.4) and 28.0 (4.5), RMSE of 82.5 (6.8) and 84.3 (5.8), SSIM of 85.9 (4.1) and 86.3 (3.9), and SNR of 20.1 (0.9) and 20.0 (0.9) respectively. Statistical tests between adapted vs target model, and adapted vs combined models in four similarity measures all show p-values less than 0.05. All these results lead us to conclude that the adapted model has the best performance among the three updated models. Upon visual evaluation of supplementary figure 4, it also indicates that the source model performs poorly on the Prostate2 dataset, while the target, combined, and adapted models can improve its performance by 39.4%, 40.3%, and 52.2% respectively, if evaluated in MAE similarity measure. Among all the updated models, the adapted model performs the best. In the Prostate3, Cervix, and Pancreas datasets, which come from Elekta XVI (Versa) or XVI (Agility) scanners, similar patterns were also noted for the four models' performance in supplementary figures 8–10. For visual evaluation, one can refer to supplementary figures 5–7. Thus, when applying the source model to datasets which it has never seen before and been collected from different anatomical sites and different vendors' scanners, the source model fails to generate good quality sCT images, and its MAE, RMSE, SSIM, SNR scores are substantially worse than in the previous scenario (Prostate1). All three updated models greatly outperform the source model, and the adapted model always performs best.

Figure 9 shows that the performance of the target, combined, and adapted models depends on the size of the training dataset in the target domain (Prostate2). The three models performed similarly with a large number (2730) of training images. As the number of training images decreased, the accuracy of the combined model decreased the fastest, while the accuracy of the adapted model decreased the slowest. Because the combined model was trained on the Prostate2, H&N1 and H&N2 datasets, the training data for the combined model becomes unbalanced when there is much less Prostate2 data than H&N1 and H&N2 data. The knowledge previously learned from the H&N1 and H&N2 data gives the adapted model a good starting point to update features, so it requires less training data to achieve good performance. Therefore, with a small number of training images, the adapted model is clearly superior to the target and combined models. The adapted model still achieved good performance when trained on only 1050 images from 15 patients. These results suggest that model fine-tuning is the best option to guarantee good performance and robust training regardless of the amount of data available.

Zoom In Zoom Out Reset image size