Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation

Neerav Karani; Ertunc Erdil; Krishna Chaitanya; Ender Konukoglu

Source-linked AI summary

Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation

Neerav Karani, Ertunc Erdil, Krishna Chaitanya, Ender Konukoglu

arXiv:2004.04668v4Apr 9, 2020eess.IVcs.CVcs.LGstat.ML

TL;DR

Scanner and protocol shifts can substantially degrade medical-image segmentation CNNs because test images may differ from the training distribution. The paper adapts a shallow normalization module for each test image using a denoising-autoencoder prior over plausible labels, while keeping the segmentation network fixed. Across multi-center MRI datasets spanning brain, heart, and prostate, test-time adaptation consistently improved performance, although the method depends on assumptions about the shift and corruption model.

Problem
Medical-image segmentation CNNs lack robustness when test images differ from training images in scanner or protocol, limiting deployment across acquisition settings.
Method
A shallow image-normalization CNN is adapted per test image using a denoising autoencoder to encourage plausible segmentations, while the downstream segmentation CNN remains fixed.
Results
Across multi-center MRI datasets for brain, heart, and prostate, test-time adaptation consistently improved segmentation performance over the evaluated alternatives.
Takeaways & Limitations
The architecture is agnostic to the deep segmentation CNN and shows promise for robustness to unseen scanner and protocol variations.
Takeaways & Limitations
The method assumes target-domain prediction errors resemble the chosen noising process and may be vulnerable when source and target images differ substantially in modality or protocol.

Abstract

from arXiv · show

Convolutional Neural Networks (CNNs) work very well for supervised learning problems when the training dataset is representative of the variations expected to be encountered at test time. In medical image segmentation, this premise is violated when there is a mismatch between training and test images in terms of their acquisition details, such as the scanner model or the protocol. Remarkable performance degradation of CNNs in this scenario is well documented in the literature. To address this problem, we design the segmentation CNN as a concatenation of two sub-networks: a relatively shallow image normalization CNN, followed by a deep CNN that segments the normalized image. We train both these sub-networks using a training dataset, consisting of annotated images from a particular scanner and protocol setting. Now, at test time, we adapt the image normalization sub-network for \emph{each test image}, guided by an implicit prior on the predicted segmentation labels. We employ an independently trained denoising autoencoder (DAE) in order to model such an implicit prior on plausible anatomical segmentation labels. We validate the proposed idea on multi-center Magnetic Resonance imaging datasets of three anatomies: brain, heart and prostate. The proposed test-time adaptation consistently provides performance improvement, demonstrating the promise and generality of the approach. Being agnostic to the architecture of the deep CNN, the second sub-network, the proposed design can be utilized with any segmentation network to increase robustness to variations in imaging scanners and protocols. Our code is available at: \url{https://github.com/neerakara/test-time-adaptable-neural-networks-for-domain-generalization}.

1 Introduction

Medical image segmentation CNNs perform well in-distribution but are not robust to scanner and protocol differences between training and test images. The work targets this domain-shift problem while building on CNN segmentation performance.

CNN segmentation methods achieve top performance across several challenges and can approach inter-expert variability for some anatomies and modalities.
Scanner and protocol variation between training and test images remains a key barrier to large-scale clinical adoption.
Domain generalization trains a robust mapping from one or more labelled source domains for application to unseen target domains without retaining the source dataset during inference.
Transporting a trained CNN for domain-generalized inference can reduce the data-sharing burden associated with sharing source datasets across institutions.

2 Related work

Prior robustness and domain-generalization methods primarily promote domain-invariant features, augment training data, or impose structural constraints on predicted segmentations. These approaches often rely on multiple source domains or handcrafted assumptions.

Domain Invariant Features: Domain-invariant feature methods reduce reliance on domain-specific signals through autoencoding, discrepancy regularization, adversarial learning, or meta-learning.
Domain Invariant Features: A common disadvantage of domain-invariant feature approaches is requiring access to multiple source domains during training.
Data Augmentation: Data augmentation methods simulate domain variation using heuristic transformations, data-generation knowledge, worst-case searches, or multiple-source interpolation.
Imposing Shape Constraints during Training: Shape-constraint methods impose anatomical or topological restrictions on predicted segmentations.

3 Method

The method concatenates an adaptable image-normalization CNN with a fixed segmentation CNN, then adapts normalization per test image using a denoising autoencoder prior over plausible segmentations. The approach assumes scanner and protocol shifts can be handled through low-level intensity and contrast changes and that predicted errors resemble the DAE’s corruption process.

Architecture: The segmentation network is decomposed as Z = Sθ(Nφ(X)), with normalization Nφ preceding the normalized-image-to-segmentation network Sθ.
Test-Time Adaptation: During test-time adaptation, Nφ is updated for each image while Sθ remains fixed at its source-domain-trained parameters.
Assumptions: The approach assumes scanner and protocol shifts preserve anatomical structures and that target-domain prediction errors can be represented by the chosen segmentation noising process.
DAE Prior: The DAE maps corrupted segmentations toward clean segmentations resembling source-domain labels, providing the plausibility signal that drives adaptation.
Test-Time Adaptation: Adaptation minimizes dissimilarity between the segmentation prediction and its DAE-denoised version, iteratively making the prediction more plausible.
Design Rationale: The shallow normalization module has a small receptive field, limiting large structural alterations while allowing image-specific contrast changes.
DAE and 2D Compatibility: The DAE is modeled as a 3D CNN to learn relative anatomical locations and volumetric shape information, while a batching strategy extends adaptation to 2D segmentation CNNs.

4.1 Datasets

The experiments cover brain, prostate, and cardiac MRI across multiple datasets and domains. Across anatomies, segmentations remain similar between domains while scanner- and protocol-dependent image contrasts vary.

Experiments use multiple MRI datasets spanning the brain, prostate, and heart.The domain-generalization setting uses labelled data from one source domain, potentially containing images from multiple scanners.
Brain experiments use HCP-T1w as the source domain and ABIDE-Caltech-T1w and HCP-T2w as target domains.The brain task segments 15 labels, including background and anatomical structures.
Underlying organ structures are relatively consistent across domains, whereas images acquired with different scanners or protocols show varying contrasts.
Figure 2 includes three brain, three prostate, and two cardiac datasets with corresponding ground-truth segmentation maps.The prostate datasets include separate sub-gland labels in two datasets and whole-gland labels in the third.

4.2 Pre-processing

Images undergo bias-field removal and per-image intensity normalization before task-specific processing. Brain images additionally receive skull stripping, and all images and labels are spatially standardized before evaluation in original pixel size.

All images are bias-field corrected with N4 and undergo per-image 0-1 intensity normalization followed by intensity clipping.The normalization uses image-intensity percentiles before clipping values at 0 and 1.
Brain preprocessing additionally removes non-brain voxels through skull stripping by setting their intensities to zero.
Images and ground-truth labels are rescaled and cropped or padded to a fixed 256x256 image size.Fixed pixel sizes are 0.7mm2 for brain, 0.625mm2 for prostate, and 1.33mm2 for cardiac datasets.
Predicted segmentations are rescaled back and evaluated at their original pixel size to avoid experimental biases.

4.3 Implementation Details

The implementation separates shallow image normalization from segmentation and trains the segmentation CNN and DAE with different spatial dimensionalities. Data augmentation, DAE corruption, and atlas initialization support training and adaptation.

The image-normalization CNN uses three convolutional layers with 16, 16, and 1 output channels and trainable channel-specific Gaussian activations.
Training augmentations combine geometric and intensity transformations, with anatomy-specific rotations and flips added for cardiac images.Geometric transformations apply to images and segmentations, while intensity transformations apply only to images.
DAE training selects corruption hyperparameters through visual inspection and validates denoising on source-domain validation images corrupted 50 times.The selected maximum copied-patch size is nmax_2 = 20.
For brain T2w target images, an atlas-based initial optimization is used before switching optimization toward the DAE.The atlas is formed by voxel-wise averaging one-hot source-domain labels.
Performance is evaluated with Dice coefficient and 95th percentile Hausdorff distance, averaged across foreground labels, test images, and three runs.

4.4 Results

The results compare source-only, domain-generalization, post-processing, and test-time adaptation strategies using qualitative and quantitative evaluations. Test-time adaptation is reported to improve performance across datasets for brain and prostate anatomies.

Qualitative results: Figure 3 compares normalized images and predicted segmentations for SD, SD + DA, SD + DA + TTA, and TD training.Rows represent target-domain results for different anatomies, alongside test images and ground truth.
Compared methods: The source-domain model provides a baseline, while target-domain-specific training provides a benchmark for the generalization gap.
Domain generalization methods: Extensive data augmentation improves segmentation on unseen scanners and protocols but leaves a gap relative to separate target-domain training.
Domain generalization methods: DAE post-processing substantially improves over SD + DA on prostate data but degrades performance on brain datasets.

4.4.3 Test-Time Adaptation

Test-time adaptation improves segmentation across brain and prostate datasets and remains competitive with UDA without requiring labelled source data during adaptation. For cardiac data, augmentation already produced strong performance that TTA preserved but did not improve.

TTA provided substantial performance gains over competing methods across brain and prostate datasets.Qualitative improvements included correcting contextually misplaced predictions, completing organ shapes, and removing outliers.
TTA improved significantly over SD + DA + Post-Proc. on 3 of 5 datasets.Significance was assessed with a paired permutation test using 100000 permutations.
TTA preserved the cardiac baseline performance but could not further improve it.The SD training with data augmentation already yielded fairly good cardiac segmentations.
Using the same target images for adaptation and testing did not produce additional improvements.The evaluation considered target images used for adaptation as well as separate test images.
TTA achieved comparable results to the best UDA methods without using labelled source data during adaptation.The comparison covered scanner-related and larger domain shifts under the reported experimental settings.

4.4.5 Analysis experiments

Analysis experiments show that restricting adaptation to the normalization network is important, while DAE post-processing alone is less effective than adapting the network itself. The adaptation converged reliably in experiments, although multiple DAE passes could worsen brain segmentation.

Parameter adaptation: Adapting all segCNN parameters reduced Dice accuracy while improving Hausdorff distance in nearly all tested domain shifts.The authors attribute the Dice decline to less accurate organ edges despite more plausible shapes and fewer outliers.
Parameter adaptation: Freezing most parameters prioritizes accurate organ edges over removing extreme outliers through adaptation.The authors note that remaining outliers can be addressed with other post-processing steps if needed.
DAE post-processing: Multiple passes through the DAE did not improve segmentation as much as test-time adaptation and worsened accuracy on brain datasets.The comparison included 1, 10, and 100 DAE passes alongside the adapted prediction.
Convergence: Test-time adaptation converged across more than 100 test volumetric images, despite lacking a theoretical convergence guarantee.The experiments covered multiple anatomies, target domains, and repeated runs.
Convergence: Dice against ground truth correlated with Dice between DAE inputs and outputs, supporting the latter as a criterion for selecting normalization parameters.This relationship was reported for the convergence analysis across target domains.

5 Discussion

The discussion presents TTA as a method for robustness to unseen scanners and protocols, while identifying assumptions about DAE reliability and possible improvements to its training and normalized representations. Convergence was observed empirically but is not theoretically guaranteed.

The method targets cross-scanner and cross-protocol robustness on images from completely unseen scanners or protocols.It combines test-time CNN adaptation with denoising autoencoders to increase segmentation plausibility.
The approach assumes that incorrect unseen-domain segmentations resemble the DAE training input distribution and that DAE outputs are reliable.The authors identify the heuristic label-noising strategy as an area for improvement.
Empirical convergence was observed across experiments, but theoretical convergence is not guaranteed.Figure 5 relates adaptation progress to segmentation Dice and DAE input-output Dice.
The DAE is trained to output one clean segmentation from a noisy input, so treating its outputs as posterior samples remains an assumption.The discussion suggests that training for multiple possible denoised segmentations could benefit performance.
Additional constraints on the intermediate normalized representation could further guide test-time adaptation.The normalized representation is not expected to be a direct translator from the target domain to the source domain.

6 Conclusion

The method targets scanner- and protocol-related domain shifts in medical image segmentation through per-image normalization adaptation guided by denoising autoencoders.

It introduces an adaptable per-image normalization module within a segmentation CNN.
Denoising autoencoders drive test-time adaptation by incentivizing plausible segmentation predictions.
Experiments across multiple datasets and anatomies demonstrate the method's promise and generality relative to data augmentation, meta-learning, and unsupervised domain adaptation.