Inclusive-Exclusive Model Training Framework to Jointly Perform Semantic Segmentation and Uncertainty Map Estimation
Hossein Arabi, Habib Zaidi · 2022
In unsupervised or semi-supervised approaches, large quantity of unlabeled datasets are readily available which could be effortlessly used to develop segmentation models. A key factor in the success of the unsupervised approaches is a framework to generate the uncertainty maps for the unlabeled data/samples based on which the performance of the model is enhanced through refining the training dataset. In this light, we set out to propose/develop a model training framework that jointly estimates the segmentation probability map and uncertainty/confidence map to distinguish between accurate and inaccurate segmentation outcomes. The proposed method relies on seven parallel deep learning models trained on the same training dataset using seven loss functions with varying levels of exclusive (conservative) to inclusive (radical) criteria. To this end, seven different cross-entropy loss functions were developed to give different sensitivity (errors) to the voxels to be included in or excluded from the target structures. The average of the probability maps generated by the seven models would serve as a single probability map of the target structure and the voxel-wise standard deviation or variation of the seven probability maps would be regarded as an uncertainty map. The same residual convolutional architecture was used for the seven networks with only different loss functions from the most exclusive (conservative) to the most inclusive (radical). For evaluation of this framework, the kidney dataset from the Decathlon challenge (kidney segmentation from CT images) and BraTS19 dataset for lesion segmentation from MR images were employed. The visual inspection of the segmentation outcomes (either kidney from CT images or lesion from MR images) revealed the superior performance of the proposed training framework to the conventional approach with a single cross-entropy loss function. The Dice indices for kidney and lesion segmentation improved from 91.7±1.6 to 94.4±1.4 and 90.9±1.8 to 93.1±1.4, respectively, when using the inclusive-exclusive framework. Moreover, a high correlation was observed between the true labels and the voxel-wise uncertainty map obtained from the proposed framework. The estimated uncertainty map could be efficiently used to discriminate between accurate and inaccurate segmentation outcomes. Overall, the quantitative analysis demonstrated the overall improved segmentation performance using the proposed framework as well as the high correlation between the estimated uncertainty indices and the true labels.