1.

Introduction

Oral cancer is one of the leading causes of cancer-related deaths, especially in South Central Asia and Melanesia, accounting for 377,713 new cases and 177,757 new deaths in 2020, according to GLOBOCAN.¹ It is highly prevalent in India, where the incidence rate in males was 13.9 per 100,000 and the mortality rate in females was 7.7 per 100,000 in 2018.² In middle- and low-income countries such as India, the five-year survival rate is $<50\%$ due to late diagnosis, according to the paper published in 2014.³ Therefore, point-of-care screening platforms and algorithms are in great need.

Two important deep-learning applications in clinical research are accurate and automated medical image classification and segmentation. Deep learning has achieved state-of-the-art performance in medical image analysis, including classification and segmentation for cancer diagnosis. Many deep-learning methods have been used to perform the diagnosis of different cancers, including skin cancer,⁴ breast cancer,⁵ and oral cancer.⁶ Despite the state-of-the-art performance of deep learning in medical segmentation, it rarely provides an uncertainty estimation when making predictions. Deep-learning models are often considered black boxes due to a lack of theoretical understanding of their underlying mechanisms. To improve the reliability of deep-learning methods and use them for clinical applications, uncertainty estimation related to the model’s prediction is a key factor to consider. The Bayesian deep learning (BDL) model⁷ provides a framework to accomplish this task by modeling the posterior distribution. Bayesian networks learn a distribution over their weights instead of deterministic ones.

Some researchers have used BDL for different medical applications to quantify uncertainty. Liu et al.⁸, introduced deep spectral learning for optical imaging oximetry with uncertainty quantification. Chai et al.⁹ proposed a Bayesian deep multisource learning model that incorporates model uncertainty into glaucoma diagnosis. Rączkowski et al.¹⁰ introduced a Bayesian convolutional neural network (CNN) for classifying histopathological colorectal images with uncertainty measurements. Liu et al.¹¹ designed a spatial attentive BDL network for automatic segmentation of the peripheral and transition zones of the prostate with uncertainty estimation. Some pioneering works have also applied BDL to build more reliable deep-learning methods for diagnosis of diseases including, but not limited to, skin cancer,¹² oral cancer,¹³ and prostate cancer.¹⁴

In general, BDL research for medical applications is an active topic aimed at improving the robustness and reliability of CNNs. Therefore, we introduce an uncertainty estimation method for oral cancer image segmentation based on a Bayesian UNet architecture in this study. The dataset used in this study contains 492 white-light images that were captured using our customized oral cancer screening platform^15,16 and annotated by oral oncology specialists. There are several different types of oral potentially malignant disorders that have unique clinical features that can be observed under white-light illumination.¹⁷ Nonhomogeneous leukoplakia is one such disorder that commonly includes symptoms of white and/or red patches; small polypoid outgrowths; rounded, red, or white excrescences; and a wrinkled or corrugated surface appearance. The lesions of erythroplakia are usually irregular in outline and have a bright red velvety surface. Because these unique clinical features are the basis for diagnosis, machine learning algorithms need to find them in images to perform correct automatic diagnosis.¹⁷ We trained the model and evaluated the segmentation and uncertainty estimation performance using multiple metrics. We also compared models with MC dropout layers applied to only the contracting path or expansive path, or both, and built an efficient model by replacing the convolutional layers with separable convolutional layers. To the best of our knowledge, this is the first study leveraging BDL to enhance the reliability and understandability of results from deep learning-based oral cancer image segmentation.

2.

Material and Methods

2.2.

Bayesian Deep Learning

Despite their success in different medical tasks, one of the limitations of CNNs for medical applications is their inability to provide prediction uncertainties. The softmax output (predictive probabilities) obtained at the end of a CNN is often erroneously interpreted as model confidence. This is an unwise solution, however, as a model can be uncertain in its prediction even with a high softmax output. Uncertainty quantification is a key factor for the clinical application of deep learning methods because it can increase the reliability of results provided by these methods. BDL models provide a framework for estimating uncertainty by modeling the posterior distribution.^20–22

Bayesian networks are probabilistic models, not deterministic ones, that learn a distribution over their weights. Given training data $X$ and $Y$, they aim to learn the posterior distribution of the neural network’s weights $W$. The posterior distribution is often approximated using variational inference methods, such as Dropout variational inference. Monte Carlo (MC) dropout²³ can be considered using the Bernoulli distribution to approximate distributions over the network’s weights. The prediction distribution of a Bayesian deep network for a new input x* is modeled as

$$p(y^{*}|x^{*},X,Y)=\int p(y^{*}|x^{*},W)p(W|X,Y)\mathrm{d}W,$$

where $p(y^{*}|x^{*},W)$ is the Softmax function and $p(W|X,Y)$ is the posterior over the weights. The prediction is approximated by sampling the model multiple ($\sigma$) times. The uncertainty is obtained by calculating the variance:

$$p(y^{*}|x^{*},X,Y)\approx \frac{1}{\sigma }\sum _{i=1}^{\sigma }\mathrm{Softmax}(f{W}_i^{*}(x^{*})),$$

$$v=\frac{1}{\sigma }\sum _{\mathrm{i}=1}^{\sigma }{(p(y|x^{*},W_i)-p(y|x^{*},X,Y))}^2,$$

where $p(y|x^{*},W_i)$ represents $\sigma$ times softmax output with different weights $W_i$ of input $x^{*}$ and $p(y|x^{*},X,Y)$ is the predictive posterior mean of input $x^{*}$. This study applied MC dropout layers (with 0.5 rate) in each contracting block, in each expansive block, or in both paths simultaneously. The MC Dropout layer was placed following the max pooling layer in the contracting block, whereas in the expansive block, the MC Dropout layer was placed following the up-convolutional layer. The proposed UNet-based BDL model is shown in Fig. 1(a), which is the original model with conventional convolutional layers, and Fig. 1(b) shows the efficient model with depthwise separable convolution layers.

Fig. 1

The proposed BDL model for oral cancer image segmentation based on UNet with (a) conventional convolutional layers; (b) efficient depthwise separable convolution layers; (c) depthwise separable convolution filter; and (d) standard convolution filter.

2.3.

Dataset

The dataset used in this study contains 492 white-light images that were captured using our customized oral cancer screening platform^15,16 from patients attending the outpatient clinics of the Department of Oral Medicine and Radiology at the KLE Society Institute of Dental Sciences, the Head and Neck Oncology Department of Mazumdar Shaw Medical Center, and the Christian Institute of Health Sciences and Research, India. Institutional ethics committee approval was obtained from all participating hospitals and written informed consents were collected from all subjects enrolled. These images were annotated by oral oncology specialists from Mazumdar Shaw Medical Center, KLE Society Institute of Dental Sciences, and Christian Institute of Health Sciences and Research using MATLAB Image Labeler.²⁴ The oral potentially malignant lesion (OPML) and malignant lesion areas in these images were labeled. The dataset used in this study contains 396 positive samples that have OPML and malignant lesions and 96 negative samples (examples shown in Fig. 2). We performed 10-fold cross-validation in this study.

Fig. 2

Examples of the dataset used for this study. (a1)–(d1) White-light oral cavity images captured using our customized oral cancer screening platform. (a2)–(d2) Corresponding pixel-level annotations labeled by oral oncology specialists. (a)–(c) Three positive samples and (d) one negative sample. The OPML and malignant lesion areas are shown in red, and other areas are shown in gray.

3.

Experiments and Results

The training data was augmented multiple ($n=6$) times with horizontal/vertical flipping, rotation, zoom, brightness adjustment, and gamma correction. The Adam optimizer was used with a batch size of 16 in each experiment. All models were trained for 300 epochs, and the best model was saved after every epoch if there was a decrease in validation loss. Code implementation was made with Keras and Tensorflow backend (using the Python programming language), and the training was done on the high-performance computing platform of the University of Arizona.²⁶ The trained models were inferenced on a desktop computer with an Intel Xeon Silver 4114 CPU, an Nvidia 1080Ti GPU, and 32 GB of RAM. For uncertainty estimation, the models were sampled multiple ($\sigma =100$) times for each test image.

First, we trained a network with MC dropout layers applied to each contracting and expansive block using conventional convolutional layers. We evaluated the Bayesian deep neural network’s segmentation performance with 10-fold cross-validation and compared it with an original UNet model without any MC dropout layers (see Table 1). The model achieved 0.714 mean IoU, 0.796 weighted IoU, and 0.881 pixel accuracy, and it performed better than the original UNet model on the oral dataset. These results show that this model was able to segment the oral potential malignant lesion and malignant lesion areas from healthy tissue and background.

Table 1

Segmentation performance comparison of the Bayesian deep-learning network with original UNet (mean and standard deviations of the cross-validation).

Figure 3 shows examples of model predictions that include uncertainty estimations. Figures 3(a), 3(e), and 3(i) are three white-light oral cavity images. Figures 3(b), 3(f), and 3(j) show the doctor’s annotations. Figures 3(c), 3(g), and 3(k) present examples of uncertainty estimation of these cases. These uncertainty maps are obtained by sampling 100 predictions from the model and estimating the standard deviation for each pixel. Pixels displayed in bright green are associated with high uncertainty, and pixels displayed in dark blue are associated with high certainty. Figures 3(d), 3(h), and 3(l) show the results of uncertainty estimation as well as lesion segmentation. These heatmaps are obtained by combining the information of segmentation and uncertainty estimation. Pixels displayed in red are associated with high certainty of suspicious lesion areas, pixels displayed in dark blue are associated with high certainty of non-suspicious areas, and pixels displayed in green are associated with high uncertainty.

Fig. 3

Example results of uncertainty estimation and lesion segmentation using the proposed Bayesian deep learning model for oral cancer image segmentation. The first column shows the original white-light images, the second column shows the annotation by specialists, the third column shows the uncertainty estimation generated by the model, the fourth column shows the uncertainty estimation and lesion segmentation together in the form of a heatmap, and the fifth column shows the predicted binary label mask.

The cases shown in the first two rows of Fig. 3 indicate that the model has high confidence for most pixels in its prediction, except for pixels near lesion borders. This is reasonable as it is difficult to assess the lesion edges accurately even for experienced specialists. In addition to the edges, the model also shows uncertainty with (1) some suspicious areas that are not obvious and (2) some non-suspicious areas with feature changes. Indeed, these are the extra pieces of information that we expect the uncertainty estimation to divulge to help find challenging prediction areas. For example, although the model fails to segment parts of the suspicious areas in the last case (the last row of Fig. 3), the model shows high uncertainty on suspicious areas that are not obvious and low uncertainty on more obviously suspicious areas. These examples indicate the BNN model can produce pixel-level uncertainty estimation.

By removing some of the uncertain pixels and leaving these confusing areas for further inspection, the model can produce more accurate and reliable segmentation results on the remaining areas. We measured the change in pixel accuracy, mean IoU, and weighted IoU when removing pixels with uncertainty values higher than a specific level. By adjusting the level of uncertainty thresholding, we plotted the change of these three evaluation metrics in Fig. 4. We can see a continuous increase in all three evaluation metrics in response to a change in uncertainty thresholding.

Fig. 4

The change of pixel accuracy, mean IoU, and weighted IoU when removing pixels with uncertainty values higher than a specific level.

To check if this process removed too many uncertain pixels, we monitored the remaining pixel ratios change (1—removed pixel ratio) corresponding to different uncertainty thresholds (see Fig. 5). If we want $\sim 90\%$ of the pixels to remain, the model can achieve a pixel accuracy of 0.911, a mean IoU of 0.751, and a weighted IoU of 0.841. This is higher than 0.881/0.714/0.796, the result without removing uncertain pixels (see Table 2). This experiment was not trying to prove that this method could improve the accuracy by removing some uncertain pixels. The result only demonstrated that the model provides less accurate predictions in uncertain areas and may need more attention and further inspection. We expect the proposed method could provide extra pieces of uncertainty information in addition to the binary segmentation result to help find challenging prediction areas that need further inspection. The removed 10% pixels (uncertainty pixels as described above) were located mainly at the lesion boundary, and some suspicious areas that are not obvious as well as some non-suspicious areas with feature changes. This coincides with our expectations, as it is difficult to assess the lesion edges accurately even for experienced specialists.

Fig. 5

The remaining pixel ratios (1—removed pixel ratio) corresponding to different uncertainty thresholds.

Table 2

Segmentation performance comparison of the Bayesian deep learning network with and without removing uncertain pixels (mean and standard deviations of the cross-validation).

Because we were also curious about the influence of MC dropout layers when applied to the contracting or expansive paths of UNet, we trained two more models that either (1) only add MC dropout layers in the contracting blocks or (2) only add MC dropout layers in the expansive blocks. These two models were trained using the same parameter settings and with 10-fold cross-validation. The segmentation performance comparison is shown in Table 3. The segmentation performance of the first model is the best, which may be due to more dropout layers resulting in less overfitting. To evaluate the uncertainty estimation performance, we calculated and compared the p (accurate | certain), $p$ (uncertain | inaccurate), and PAvPU described in Sec. 2.4. The comparison is shown in Fig. 6. The model with MC Dropout layers applied to both contracting and expansive paths works better than the other two models on p (accurate | certain) and PAvPU, whereas the model with MC dropout layers applied only to contracting path works better than the other two models on $p$ (uncertain | inaccurate). For Fig. 6, the values of these metrics depended on three parameters (described in Sec. 2.4): the patch dimensions, the accuracy threshold, and the uncertainty threshold. We fixed the patch dimensions as $2\times 2$ and the accuracy threshold as 0.5. We then observed how these metrics varied with a change of uncertainty threshold. A model with a higher value of these metrics is a better performer.

Table 3

Segmentation performance comparison of three models by adding MC dropout layers on contracting blocks, expansive blocks, or both (mean and standard deviations of the cross-validation).

Fig. 6

Uncertainty estimation performance comparison of three models by adding MC dropout layers on contracting or expansive blocks or both, using $p$ (accurate | certain), $p$ (uncertain | inaccurate), and patch accuracy versus PAvPU.

Although the computing time of one single inference is insignificant, the models need to be sampled multiple times to estimate the uncertainty, so efficiency is a concern, especially for the outdated computers and portable devices commonly used in resource-limited settings. Therefore, we built a more compact and efficient model by replacing the convolutional layers of the original model with depthwise separable convolutional layers (described in Sec. 2.1). The new model was trained using the same settings and with 10-fold cross-validation. The efficient model is almost six times smaller than the original one (3.42 MB versus 20 MB). The computing speed when sampling 100 times is also faster using the desktop computer mentioned before (42 s versus 1 min 24 s). These improvements make our new efficient model more suitable for future implementation on portable devices. We were worried about whether the improvement of efficiency might affect segmentation and uncertainty estimation performance, so we also calculated the mean IoU, weighted IoU, and global pixel accuracy, as well as p (accurate | certain), $p$ (uncertain | inaccurate), and PAvPU and compared with the original model. The results are shown in Table. 4 and Fig. 7. This experiment shows that the efficient model reduced the computation time and size, albeit with some performance compromises. Although the performance compromises of the efficient model are not huge, the difference in accuracy is still important and not negligible for rigorous medical applications such as a treatment protocol design. Therefore, the efficient model will be considered for detection tasks in resource-limited settings. We are interested in implementing the efficient model on portable devices and comparing the performance with a model directly distilled from the UNet to be smaller.

Table 4

Segmentation performance comparison of the efficient and original models (mean and standard deviations of the cross-validation).

Fig. 7

Uncertainty estimation performance comparison of the efficient and original models, using $p$ (accurate | certain), $p$ (uncertain | inaccurate), and patch accuracy versus PAvPU.

4.

Conclusion

In this paper, we have proposed a Bayesian UNet architecture for uncertainty estimation of oral cancer image segmentation and have shown that the model achieves good segmentation accuracy with 10-fold cross-validation. By sampling the model multiple times, uncertainty maps of oral cancer images can be obtained. The uncertainty maps can provide more pixel-level information than segmentation predictions alone. From the results, we observe that the model is uncertain with (1) lesion borders, (2) some suspicious areas that are not obvious, and (3) some non-suspicious areas with feature changes when making the prediction. With this extra uncertainty information, the accuracy and reliability of the model’s prediction can be improved. In the experiments, we observed an improvement in pixel accuracy and mean IoU by removing uncertain pixels. This result reflects that the model provides less accurate predictions in uncertain areas that may need more attention and further inspection. To evaluate the segmentation uncertainty estimation of our models, we also used the metrics introduced by Mukhoti and Gal.²⁵ We experimentally compared three models with MC dropout layers applied to only the contracting path, to only the expansive path, and to both paths simultaneously. We also built and tested an efficient model by replacing the conventional convolutional layers with depthwise separable convolutional layers. The efficient model is almost six times smaller and two times faster than the original UNet, with small performance compromises, expanding its potential for future implementation on portable devices. In general, our proposed method can effectively segment the OPML and malignant lesion areas from healthy tissue and background, as well as estimate uncertainty when making the prediction.