Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation

Abstract

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}.

arXiv commentPublished in Medical Image Analysis journal: https://doi.org/10.1016/j.media.2020.101907

References in corpus (18)

• Adam: A Method for Stochastic Optimization
• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• A Survey on Deep Learning in Medical Image Analysis
• Efficient Multi-Scale 3D CNN with Fully Connected CRF for Accurate Brain Lesion Segmentation
• Convolutional Neural Networks for Medical Image Analysis: Full Training or Fine Tuning?
• Efficient Inference in Fully Connected CRFs with Gaussian Edge Potentials
• Interactive Medical Image Segmentation using Deep Learning with Image-specific Fine-tuning
• Anatomically Constrained Neural Networks (ACNN): Application to Cardiac Image Enhancement and Segmentation
• Domain Generalization via Model-Agnostic Learning of Semantic Features
• Generalizing to Unseen Domains via Adversarial Data Augmentation
• Generalizing Across Domains via Cross-Gradient Training
• SynSeg-Net: Synthetic Segmentation Without Target Modality Ground Truth
• Generalized Denoising Auto-Encoders as Generative Models
• Test-Time Training with Self-Supervision for Generalization under Distribution Shifts
• Deep Mean-Shift Priors for Image Restoration
• Learning Shape Priors for Robust Cardiac MR Segmentation from Multi-view Images
• Nonlinear Markov Random Fields Learned via Backpropagation
• Denoising Auto-encoding Priors in Undecimated Wavelet Domain for MR Image Reconstruction

Cited by in corpus (19)

• SynthSeg: Segmentation of brain MRI scans of any contrast and resolution without retraining
• A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts
• Robust machine learning segmentation for large-scale analysis of heterogeneous clinical brain MRI datasets
• Domain Generalization for Medical Image Analysis: A Review
• UPL-SFDA: Uncertainty-aware Pseudo Label Guided Source-Free Domain Adaptation for Medical Image Segmentation
• Uncertainty-Induced Transferability Representation for Source-Free Unsupervised Domain Adaptation
• Unsupervised Knowledge-Transfer for Learned Image Reconstruction
• FUTURE-AI: Guiding Principles and Consensus Recommendations for Trustworthy Artificial Intelligence in Medical Imaging
• Generalizable Pancreas Segmentation via a Dual Self-Supervised Learning Framework
• Unsupervised Domain Adaptation in the Absence of Source Data
• A Dual-Task Synergy-Driven Generalization Framework for Pancreatic Cancer Segmentation in CT Scans
• Harmonizing Flows: Leveraging normalizing flows for unsupervised and source-free MRI harmonization
• Multi-Stream Cellular Test-Time Adaptation of Real-Time Models Evolving in Dynamic Environments
• Re-using Adversarial Mask Discriminators for Test-time Training under Distribution Shifts
• Test-Time Adaptation for Nighttime Color-Thermal Semantic Segmentation
• Source-Free Domain Adaptive Fundus Image Segmentation with Denoised Pseudo-Labeling
• A Principled Approach to Failure Analysis and Model Repairment: Demonstration in Medical Imaging
• Generative Self-training for Cross-domain Unsupervised Tagged-to-Cine MRI Synthesis
• Quantification of Predictive Uncertainty via Inference-Time Sampling