Source-linked AI summary
Inconsistency-aware Uncertainty Estimation for Semi-supervised Medical Image Segmentation
Yinghuan Shi, Jian Zhang, Tong Ling, Jiwen Lu, Yefeng Zheng, Qian Yu, Lei Qi, Yang Gao
TL;DR
Semi-supervised medical image segmentation remains difficult because medical images pose substantial challenges and precise pixel-wise annotation is laborious. CoraNet estimates uncertainty from inconsistent predictions under different misclassification costs, separately trains certain and uncertain regions, and outperforms comparison methods across several segmentation tasks.
Problem
Medical image segmentation is challenging, while fully supervised training requires time-consuming, laborious, and potentially subjective pixel-wise delineation.
Method
CoraNet estimates pixel-wise uncertainty from inconsistent predictions across cost-sensitive settings and uses separate self-training for certain and uncertain regions through CRM, C-SN, and UC-SN.
Results
CoraNet outperforms comparison methods in segmentation accuracy across CT pancreas, MR endocardium, and other medical image segmentation tasks, including DSC, recall, and Hausdorff distance on CT pancreas.
Takeaways & Limitations
The paper provides an uncertainty definition based directly on classification outputs without requiring a predefined boundary-aware assumption, alongside a separate treatment of certain and uncertain regions.
Abstract
from arXiv · showhide
In semi-supervised medical image segmentation, most previous works draw on the common assumption that higher entropy means higher uncertainty. In this paper, we investigate a novel method of estimating uncertainty. We observe that, when assigned different misclassification costs in a certain degree, if the segmentation result of a pixel becomes inconsistent, this pixel shows a relative uncertainty in its segmentation. Therefore, we present a new semi-supervised segmentation model, namely, conservative-radical network (CoraNet in short) based on our uncertainty estimation and separate self-training strategy. In particular, our CoraNet model consists of three major components: a conservative-radical module (CRM), a certain region segmentation network (C-SN), and an uncertain region segmentation network (UC-SN) that could be alternatively trained in an end-to-end manner. We have extensively evaluated our method on various segmentation tasks with publicly available benchmark datasets, including CT pancreas, MR endocardium, and MR multi-structures segmentation on the ACDC dataset. Compared with the current state of the art, our CoraNet has demonstrated superior performance. In addition, we have also analyzed its connection with and difference from conventional methods of uncertainty estimation in semi-supervised medical image segmentation.
I. INTRODUCTION
Semi-supervised medical image segmentation is hindered by costly annotations and uncertainty-estimation strategies that may produce unreliable pseudo-label initialization. CoraNet addresses these issues by identifying uncertainty through cost-sensitive prediction inconsistency and separately training certain and uncertain regions.
- Motivation: Medical image segmentation remains challenging because organ and tissue images exhibit large shape variation, low contrast, and irregular body movement.
- Motivation: Manual pixel-wise delineation is time-consuming, laborious, and subject to carelessness and fatigue, limiting supervised deployment in clinical scenarios.
- Research gap: Existing semi-supervised methods commonly estimate uncertainty using thresholds on softmax outputs, treating lower confidence as higher uncertainty.
- Research gap: Unreliable initial pseudo-labels can propagate errors, motivating a new uncertainty-estimation method beyond softmax-based confidence.
- Proposed approach: CoraNet identifies uncertain regions from inconsistent predictions under different misclassification costs and separately segments certain and uncertain regions.
- Proposed approach: Its CRM, C-SN, and UC-SN components are alternatively trained end-to-end, and the method outperforms baselines across CT pancreas, MR endocardium, and MR ACDC tasks.
A. Medical Image Segmentation
Medical image segmentation is challenging because anatomical variation, low contrast, and motion complicate robust prediction, while manual delineation is costly and subjective. Semi-supervised methods reduce labeling demands, but CoraNet proposes region-level uncertainty modeling based on cost-sensitive prediction inconsistency.
- Semi-supervised Learning: Previous semi-supervised segmentation methods commonly estimate uncertainty from pixel or voxel entropy within entropy-minimization or consistency-regularization frameworks.
- CoraNet Framework: CoraNet models region-level uncertainty with a conservative-radical module, then separately trains certain and uncertain regions using complementary self-training strategies.
- Cost-sensitive Uncertainty: Pixels with identical predictions under conservative and radical cost settings form certain regions, whereas inconsistent predictions identify uncertain regions.
B. Certainty-aware Prediction: CRM
CRM uses shared-encoder, separate-decoder models trained with conservative and radical misclassification costs. Their prediction disagreement produces masks distinguishing certain from uncertain regions.
- Cost-sensitive Losses: The conservative and radical losses encode different class misclassification costs, with α set to 5 in the implementation.
- CRM Design: CRM trains normal, conservative, and radical segmentation sub-models with a common encoder and separate decoders.
- CRM Design: The total CRM loss combines the normal, conservative, and radical training objectives across labeled samples.
- Mask Construction: The model obtains conservative and radical predictions for each unlabeled image and compares them to construct certainty masks.
- Mask Construction: Pixel-wise XOR marks prediction inconsistency between conservative and radical outputs, identifying the uncertain region.
C. Separate Self-training: C-SN and UC-SN
CoraNet separately self-trains certain and uncertain regions. Certain predictions serve directly as pseudo-labels, while uncertain regions use a mean-teacher mechanism for more reliable label assignment.
- C-SN for certain region segmentation: C-SN uses predictions in the certain region as pseudo-labels to self-train the main segmentation model.
- UC-SN for uncertain region segmentation: UC-SN avoids directly using low-confidence uncertain predictions as pseudo-labels and instead introduces a mean-teacher model for label assignment.
- C-SN for certain region segmentation: Certain-region self-training can provide efficient learning and complement labeled samples when labeled data are limited.
- UC-SN for uncertain region segmentation: The student and teacher models are constrained to produce consistent predictions on unlabeled samples.
- Overall training procedure: C-SN and UC-SN are alternated with CRM updates in an end-to-end training procedure summarized by Algorithm 1.
1) Network Architecture:
CoraNet’s two auxiliary decoders are designed to produce disagreement for uncertainty estimation, rather than consistency for regularization. They use distinct conservative and radical objectives and are discarded after generating uncertain masks.
- Network Architecture:: CoraNet’s conservative and radical decoders estimate prediction differences that generate region-level uncertainty masks before being discarded during testing.
- Network Architecture:: Unlike prior multi-decoder consistency methods, CoraNet optimizes its two decoders with different conservative and radical objective functions.
- Network Architecture:: Previous works commonly use identical decoder objectives with perturbations or data augmentation to encourage consistent predictions.
2) Efficiency:
CoraNet adds lightweight decoder branches during training but discards them at test time, preserving U-Net-level inference efficiency.
- Efficiency:: The model adds only two decoders, each with two 3×3 convolutions and one 1×1 convolution.The teacher network uses moving-averaged student weights, so it introduces no additional parameters.
- Efficiency:: At test time, CoraNet discards the conservative and radical decoders, leaving inference as efficient as U-Net.
- Efficiency:: In CT pancreas segmentation, training parameters are 29.6M for U-Net and 29.8M for CoraNet, while both use 29.6M at test time.
- Efficiency:: Equal weighting of the certain-region and consistency losses achieves satisfactory performance, partly because their magnitudes depend on region size and misclassification ratio.
3) Weight Setting:
The loss-weighting discussion motivates a fixed 1:1 balance between certain-region and consistency losses, while the implementation uses a four-level U-Net backbone and staged self-training.
- Weight Setting:: Certain regions contain many pixels with relatively few errors, whereas uncertain regions contain fewer pixels with relatively more errors.
- Weight Setting:: The two losses are often roughly equal, so CoraNet fixes their weighting at 1:1 to avoid an extra hyperparameter.Extremely large or small weights deteriorate performance.
- Weight Setting:: The backbone is a four-level U-Net with 64 channels in its final level, using convolutional blocks, batch normalization, ReLU, pooling, and transposed-convolution upsampling.
- Weight Setting:: CoraNet pretrains for 30 epochs, then jointly trains on labeled and unlabeled data for 100 epochs, relabeling unlabeled data every five epochs.
B. Datasets
The paper evaluates CoraNet across CT pancreas, MR endocardium, and multi-structure ACDC segmentation, examining initialization quality, uncertainty masks, and separate self-training.
- B. Datasets: The evaluation covers CT pancreas, MR endocardium, and ACDC segmentation tasks.
- B. Datasets: The CT pancreas dataset contains 82 contrast-enhanced abdominal CT volumes with physician-provided delineations.
- B. Datasets: The MR endocardium dataset contains 7,980 images from 33 individuals, targeting the left-ventricle cavity boundary.
- B. Datasets: The ACDC dataset contains 100 cine MR sequences and requires segmenting right-ventricle cavities, left-ventricle cavities, and myocardium.
- B. Datasets: CRM is evaluated for reliable initialization, while separate self-training is examined through changing uncertainty masks during training.
- B. Datasets: CRM is compared with softmax-confidence thresholds of 0.5, 0.7, and 0.9 for assigning pseudo-labels to sufficiently confident pixels.
- B. Datasets: CRM outperforms the thresholded confidence methods on all PPV, TPR, and CSI metrics except pancreas PPV.
- B. Datasets: On pancreas data, separate self-training reaches a 67.01% mean Dice score versus 61.23% for mean teacher and 62.33% for self-training.
D. Results on CT Pancreas
On CT pancreas segmentation, CoraNet is evaluated against supervised and semi-supervised baselines, ablated components, alternative uncertainty estimators, and visual examples in 2D and 3D settings.
- D. Results on CT Pancreas: The comparison includes U-Net, Π-model, temporal ensembling, mean teacher, and uncertainty-aware mean teacher baselines.
- D. Results on CT Pancreas: The evaluation uses DSC as the main metric, with precision, recall, and Hausdorff distance as supplementary metrics.
- D. Results on CT Pancreas: CoraNet outperforms comparison methods in DSC, recall, and Hausdorff distance, although temporal ensembling achieves better precision.
- D. Results on CT Pancreas: CoraNet uses 40% of all samples as labeled training data, with the remaining training samples unlabeled and 20% reserved for testing.
- D. Results on CT Pancreas: Across labeled-to-unlabeled ratios of 1:4, 1:2, and 1:1, CoraNet achieves more consistent and stable DSC performance than comparison methods.
- D. Results on CT Pancreas: The ablation evaluates the main segmentation loss, self-training, and masked self-training variants, including versions excluding uncertain-region training.
- D. Results on CT Pancreas: The complete method benefits overall segmentation performance, and comparisons with masked settings validate the efficacy of CRM’s certain-region mask.
- D. Results on CT Pancreas: Additional experiments compare 3D backbones and replace CoraNet’s mask generation with test augmentation, MC dropout, and aleatoric uncertainty.
E. Results on MR Endocardium
On MR endocardium, CoraNet outperforms the comparison methods, especially in DSC, while its segmentation remains relatively stable across tested misclassification costs. The study also illustrates more accurate results and identifies α=5 as the best tested setting.
- Comparative performance: CoraNet outperforms U-Net, Π-model, temporal ensembling, mean teacher, and UA-MT on MR endocardium, especially in DSC.Evaluation uses DSC, precision, recall, and Hausdorff distance with a 1:4 labeled-to-unlabeled ratio.
- Comparative performance: CoraNet produces more accurate MR endocardium segmentation examples than the compared methods.Figure 7 uses colored curves to compare ground truth, five baselines, and CoraNet.
- Misclassification-cost analysis: Segmentation results remain relatively stable when the misclassification cost α is set to 2, 5, or 10.The experiment evaluates these settings on both CT pancreas and MR endocardium datasets.
- Misclassification-cost analysis: α=5 is identified as the best setting because values that are too large or too small may introduce segmentation noise.The comparison includes α=1 with different decoder initializations but equal misclassification cost.
- Misclassification-cost analysis: Figure 8 visualizes initial uncertainty masks for α values of 2, 5, and 10 on the pancreas and endocardium datasets.Yellow curves indicate ground truth, while red curves indicate uncertain segmentation.
F. ACDC Segmentation
On the ACDC multi-class task, CoraNet is evaluated against multiple semi-supervised baselines under a 20%-labeled setting and 2D processing. It achieves higher segmentation accuracy and more precise difficult-boundary delineation than mean teacher in visual examples.
- Experimental setting: The ACDC comparison includes pseudo label, VAT, mean teacher, DCT-Seg, ADVNET, Mutual Information, consistency regularization, MSE, KL, and Peng et al.’s method.A fully supervised setting is also listed as an upper-bound reference.
- Experimental setting: CoraNet is evaluated for right-ventricle, left-ventricle, and myocardium segmentation using binary sub-tasks whose outputs are integrated by confidence-weighted majority voting.The method uses random rotation, flipping, and random crops during augmentation.
- Quantitative results: CoraNet achieves higher DSC than the comparison methods on ACDC under the shared evaluation setting.The reported DSC values are averages with standard deviations over three runs using different random seeds, and the comparison follows the baselines’ 2D setting.
- Qualitative results: CoraNet obtains more precise segmentation on some difficult boundary regions than mean teacher in visual ACDC examples.Different colors indicate different segmentation classes.
- Overall result: Across CT pancreas, MR endocardium, and multi-class ACDC segmentation, the proposed method outperforms related baselines in segmentation accuracy.The conclusion summarizes the paper’s evaluation across multiple medical image segmentation tasks.