Source-linked AI summary
Transformation Consistent Self-ensembling Model for Semi-supervised Medical Image Segmentation
Xiaomeng Li, Lequan Yu, Hao Chen, Chi-Wing Fu, Lei Xing, Pheng-Ann Heng
TL;DR
Medical-image segmentation is constrained by scarce and costly labeled data. The paper proposes transformation-consistent self-ensembling with supervised and consistency losses to exploit labeled and unlabeled inputs, and reports superior performance across three challenging 2D/3D tasks, while assuming labeled and unlabeled data share a distribution.
Problem
Medical-image segmentation requires labeled data that are expensive and time-consuming to collect, motivating learning from limited labeled and arbitrary unlabeled data.
Method
TCSM_v2 uses a teacher-student self-ensembling framework with weighted supervised and unsupervised losses, enforcing transformation consistency for segmentation predictions.
Results
The method shows superior segmentation performance over existing methods across skin-lesion, optic-disc, and liver datasets, including 2D and 3D medical images.
Takeaways & Limitations
The proposed transformation-consistent strategy is reported as effective for semi-supervised medical-image segmentation and applicable to both 2D and 3D networks.
Takeaways & Limitations
The method assumes labeled and unlabeled data come from the same distribution, leaving domain shift for future domain-adaptation work.
Abstract
from arXiv · showhide
Deep convolutional neural networks have achieved remarkable progress on a variety of medical image computing tasks. A common problem when applying supervised deep learning methods to medical images is the lack of labeled data, which is very expensive and time-consuming to be collected. In this paper, we present a novel semi-supervised method for medical image segmentation, where the network is optimized by the weighted combination of a common supervised loss for labeled inputs only and a regularization loss for both labeled and unlabeled data. To utilize the unlabeled data, our method encourages the consistent predictions of the network-in-training for the same input under different regularizations. Aiming for the semi-supervised segmentation problem, we enhance the effect of regularization for pixel-level predictions by introducing a transformation, including rotation and flipping, consistent scheme in our self-ensembling model. With the aim of semi-supervised segmentation tasks, we introduce a transformation consistent strategy in our self-ensembling model to enhance the regularization effect for pixel-level predictions. We have extensively validated the proposed semi-supervised method on three typical yet challenging medical image segmentation tasks: (i) skin lesion segmentation from dermoscopy images on International Skin Imaging Collaboration (ISIC) 2017 dataset, (ii) optic disc segmentation from fundus images on Retinal Fundus Glaucoma Challenge (REFUGE) dataset, and (iii) liver segmentation from volumetric CT scans on Liver Tumor Segmentation Challenge (LiTS) dataset. Compared to the state-of-the-arts, our proposed method shows superior segmentation performance on challenging 2D/3D medical images, demonstrating the effectiveness of our semi-supervised method for medical image segmentation.
I. INTRODUCTION
The paper addresses scarce, costly medical-image annotations with a transformation-consistent self-ensembling method that regularizes pixel-level predictions using unlabeled data. It evaluates the approach across skin-lesion, optic-disc, and liver segmentation, including both 2D and 3D datasets.
- Motivation: Medical-image segmentation needs substantial labeled data, but expert annotation is scarce, expensive, and time-consuming, especially for melanoma and volumetric liver delineation.The paper motivates semi-supervised learning as a way to use limited supervision alongside unlabeled data.
- Method: The method combines supervised learning on labeled inputs with self-ensembling consistency regularization across differently regularized versions of the same input.Regularization includes randomized noise, dropout, and data transformations.
- Method: Transformation consistency adapts self-ensembling to segmentation by matching transformed inputs with correspondingly transformed pixel-level predictions.The strategy uses rotation and flipping because segmentation predictions should transform equivariantly with their inputs.
- Scope and evaluation: TCSM_v2 is designed to apply flexibly to both 2D and 3D convolutional neural networks.The paper evaluates skin-lesion, optic-disc, and liver segmentation tasks across representative 2D and 3D datasets.
- Results: Experiments report effectiveness over existing methods on three challenging medical-image segmentation tasks, including a new ISIC 2017 skin-lesion segmentation record.The evaluated tasks cover dermoscopy, retinal fundus, and volumetric CT images.
II. RELATED WORK
Recent semi-supervised medical image segmentation methods use iterative pseudo-labeling, generative models, and self-ensembling to learn from limited labels. This paper develops self-ensembling for both 2D and 3D medical image segmentation, using supervised and consistency-based regularization losses.
- Semi-supervised segmentation for medical images: Deep learning-based semi-supervised segmentation methods include iterative label refinement and generative approaches such as VAEs and GANs.These methods have been applied to medical image segmentation, including cardiac, optic cup, and other tasks.
- Self-ensembling for segmentation: The method extends self-ensembling to medical image segmentation, covering both two-dimensional and three-dimensional cases.This differentiates the approach from prior self-ensembling work focused on semi-supervised classification.
- Semi-supervised learning formulation: The proposed framework combines supervised loss on labeled inputs with regularization loss using both labeled and unlabeled inputs.The supervised term uses cross-entropy, while the regularization term encourages consistent predictions under perturbations.
- Self-ensembling: Self-ensembling generates targets from the network itself under different perturbations, including Gaussian noise, dropout, and randomized transformations.The teacher and student roles use two perturbed evaluations of the same input.
- Mean teacher framework: The teacher model is updated as an exponential moving average of student parameters, with α set to 0.999 in the experiments.A larger α makes the teacher rely more on its previous parameters.
B. Transformation Consistent Self-ensembling Model
The method introduces a transformation-consistent self-ensembling model for semi-supervised medical image segmentation.
- B. Transformation Consistent Self-ensembling Model: TCSM_v2 is presented as a transformation-consistent self-ensembling method for semi-supervised medical image segmentation.
1) Motivation:
Segmentation requires transformation-equivariant predictions, but ordinary convolutions are not generally equivariant to rotation or flipping. TCSM_v2 addresses this mismatch by aligning transformed inputs and prediction maps during consistency regularization.
- Motivation: Unlike classification, segmentation should transform its pixel-wise mask consistently when the input image is rotated or flipped.The desired relationship is transformation equivariance rather than transformation invariance.
- Motivation: Convolutions are not generally rotation- or flipping-equivariant, weakening randomized transformation regularization for segmentation.
- Transformation-consistent scheme: TCSM_v2 approximates transformation consistency by applying the transformation at the input and output spaces within the unsupervised regularization term.
- Transformation-consistent scheme: Each training input is evaluated twice: one prediction transforms the input before the student network, while the other transforms the teacher prediction map.Gaussian noise and network dropout may also perturb the two evaluations.
- Transformation-consistent scheme: The consistency loss minimizes the mean squared difference between the transformed student prediction and teacher prediction.
3) Loss function:
Training combines labeled-data classification loss with consistency regularization over labeled and unlabeled samples. The regularization weight ramps up during training while teacher predictions provide targets for consistency learning.
- Loss function: The overall loss combines cross-entropy supervision on labeled inputs with a regularization term using labeled and unlabeled inputs.The weighting factor controls the strength of regularization relative to supervision.
- Loss weighting: The regularization weight follows a Gaussian ramp-up function λ(T)=k·e^(−5(1−T)^2), with k empirically set to 1.0.
- Training schedule: Early training emphasizes supervised learning because λ(T) is small, while later training uses the model’s predictions to regularize unlabeled data.
- Prediction construction: Student predictions are formed by transforming the student output, whereas teacher predictions come from applying the transformation to the input before teacher inference.
- Teacher update: The teacher model is updated through exponential moving averaging of student parameters, with α=0.999.
- Transformation regularization: The transformation-consistency operation combines random scaling, horizontal flipping, and four 90° rotation choices during training.Labeled and unlabeled samples are selected evenly and randomly within each minibatch.
C. Technical Details of TCSM_v2
The method uses DenseUNet-based teacher and student models for 2D segmentation, with a 3D U-Net variant for liver CT and standard on-the-fly augmentation during training.
- 2D architecture: DenseUNet serves as both teacher and student for dermoscopy and retinal fundus segmentation, using decoder blocks and UNet-like skip connections.The decoder combines upsampling, convolution, batch normalization, and ReLU layers.
- Transformation-consistent operation: Figure 4 visualizes the transformation-consistent operation π within DenseUNet while omitting UNet connections and the decoder for simplification.
- 3D architecture: TCSM_v2 is generalized to 3D liver CT segmentation using a modified 3D U-Net with base filters set to 32.Training uses SGD with a 0.01 learning rate, batch normalization, and weighted cross-entropy loss.
- Data processing: Training uses on-the-fly flipping, rotation, and scaling augmentation, while inference removes transformations and performs a single test.
IV. EXPERIMENTS
Experiments evaluate the method across dermoscopy, retinal fundus, and volumetric CT segmentation datasets, including ISIC 2017, REFUGE, and LiTS.
- Experimental scope: The experiments cover dermoscopy images, retinal fundus images, and liver CT scans across multiple medical imaging modalities.
- Dermoscopy dataset: The ISIC 2017 dermoscopy dataset contains 2,000 annotated training images, 150 validation images, and 600 testing images.Images are resized to 248×248 using bicubic interpolation.
- Retinal fundus dataset: The retinal fundus dataset comes from the MICCAI 2018 Retinal Fundus Glaucoma Challenge.
- Comparison setup: Table I compares supervised learning, supervised learning with regularization, and semi-supervised learning using 50 labeled and 1,950 unlabeled dermoscopy images.
- Liver CT dataset: The LiTS dataset contains 131 training and 70 testing contrast-enhanced 3D abdominal CT scans acquired across six clinical sites.The scans vary in in-plane resolution from 0.55 mm to 1.0 mm and slice spacing from 0.45 mm to 6.0 mm.
B. Evaluation Metrics
Segmentation performance is measured with classification-derived overlap and pixel-level metrics, using dataset-specific evaluation choices.
- Dermoscopy metrics: Dermoscopy segmentation is evaluated with Jaccard index, Dice coefficient, pixel-wise accuracy, sensitivity, and specificity.
- Dermoscopy metrics: Pixel-wise accuracy is defined as AC = (TP + TN) / (TP + FP + TN + FN).TP, TN, FP, and FN denote true positives, true negatives, false positives, and false negatives.
- Dermoscopy metrics: Sensitivity and specificity are defined as SE = TP / (TP + FN) and SP = TN / (TN + FP).The formulas use true-positive, true-negative, false-positive, and false-negative counts.
- Dermoscopy metrics: Jaccard index and Dice coefficient are defined as JA = TP / (TP + FN + FP) and DI = 2·TP / (2·TP + FN + FP).
- Task-specific metrics: Optic disc segmentation uses JA, whereas liver segmentation uses Dice per case according to the LiTS challenge evaluation.
C. Experiments on Dermoscopy Image Dataset
Dermoscopy experiments show that transformation-consistent semi-supervised learning improves segmentation with few labels, while complementary regularization and teacher averaging provide additional gains.
- Quantitative and visual results with 50 labeled data: With 50 labeled and 1,950 unlabeled images, the semi-supervised method outperforms supervised learning on every evaluation metric.Using the same DenseUNet architecture, improvements include 4.07% on JA and 3.47% on DI.
- Quantitative and visual results with 50 labeled data: Supervised training with regularization also improves over supervised training across all evaluation metrics.
- Effectiveness of TCSM and TCSM_v2: Noise-and-dropout and transformation-consistent regularization each contribute gains, while their combination further enhances performance.
- Effectiveness of TCSM and TCSM_v2: Random scaling with ratio 0.1 improves semi-supervised results, and the generalized transformation-consistent strategy improves the mean-teacher variant.Weight-averaged consistency targets also improve results; the final model reaches 75.24% JA and 83.44% DI, surpassing the supervised baseline by 5.7%.
- Results under different number of labeled data: Semi-supervised methods consistently outperform supervised learning across different labeled/unlabeled settings, with larger gains when labeled data are scarce.The performance gap narrows as more labeled samples become available.
- Results under different number of labeled data: TCSM_v2 consistently improves TCSM across label and unlabeled-data settings, and the unsupervised loss improves performance even when all labels are used.The authors associate this improvement with more robust features from the consistency requirement.
4) Comparison with other semi-supervised segmentation methods:
The method is compared with semi-supervised segmentation approaches and challenge submissions across dermoscopy, fundus, and liver segmentation settings. It improves over supervised baselines and achieves strong benchmark performance on dermoscopy images.
- The comparison includes recent semi-supervised segmentation, adversarial-learning, and self-ensembling methods using the same network backbone and training strategies.The experiments use 50 labeled and 1950 unlabeled dermoscopy images for validation comparisons.
- 4.07% JA improvement is achieved by using unlabeled dermoscopy data compared with supervised training.This result is reported on the validation set with 50 labeled and 1950 unlabeled images.
- 1.6% JA improvement over the state-of-the-art method raises test performance from 76.5% to 78.1% on the ISIC 2017 challenge benchmark.DI and SE also improve by 1.1% and 3.7%, respectively.
- The semi-supervised method achieves the highest reported improvement under both 10% and 20% labeled-data settings for the fundus and liver datasets.The supplied comparison passages identify these settings and improvements for the corresponding experiments.
- The authors conclude that the overall semi-supervised learning method is effective beyond the supervised baseline.
D. Experiments on Retinal Fundus Image Dataset
Experiments evaluate optic disc and liver segmentation with limited labeled data. The method achieves the strongest reported improvements in both tasks and improves optic disc boundary capture qualitatively.
- D. Experiments on Retinal Fundus Image Dataset: REFUGE training images are split 9:1, with semi-supervised experiments using 10% and 20% labeled training images.The input images are preprocessed by subtracting the mean RGB values of the training dataset.
- D. Experiments on Retinal Fundus Image Dataset: 1.82% improvement is achieved under the 10% labeled training setting for optic disc segmentation, ranking highest among compared methods.The improvement under the 20% setting is also reported as the highest.
- D. Experiments on Retinal Fundus Image Dataset: The visual results indicate better capture of the optic disc boundary in fundus images.Figure 7 shows semi-supervised segmentation examples for fundus images and liver CT scans.
- E. Experiments on LiTS dataset: LiTS liver segmentation is evaluated with Dice per case, defined as the average Dice score per volume.Experiments use 10% and 20% labeled training settings.
- E. Experiments on LiTS dataset: 5.33% and 5.72% improvements are achieved under the 10% and 20% labeled settings, respectively, with higher performance than other compared methods.The results are obtained on liver segmentation from CT volumes.
V. DISCUSSION
The discussion frames transformation-consistent self-ensembling as a semi-supervised strategy for reducing annotation effort across 2D and 3D medical image segmentation. Its main scope boundary is the assumption that labeled and unlabeled data share a distribution.
- The work targets medical image segmentation with limited annotation by exploiting numerous unlabeled images.
- Transformation-consistent self-ensembling is the paper’s key strategy for semi-supervised segmentation.
- The method is applied with both 2D and 3D segmentation networks and is described as flexible across these data formats.Examples include dermoscopy and fundus images, as well as volumetric MRI and CT data.
- The method assumes labeled and unlabeled data come from the same distribution, leaving domain shift for future domain-adaptation work.
- The transformation design addresses rotation and scaling because convolutional networks are not inherently equivariant to these transformations.
- Experiments are averaged over three trials, which the authors use as an indication of robustness.