OptTTA: Learnable Test-Time Augmentation for Source-Free Medical Image Segmentation Under Domain Shift

by   Devavrat Tomar, et al.

As distribution shifts are inescapable in realistic clinical scenarios due to inconsistencies in imaging protocols, scanner vendors, and across different centers, well-trained deep models incur a domain generalization problem in unseen environments. Despite a myriad of model generalization techniques to circumvent this issue, their broad applicability is impeded as (i) source training data may not be accessible after deployment due to privacy regulations, (ii) the availability of adequate test domain samples is often impractical, and (iii) such model generalization methods are not well-calibrated, often making unreliable overconfident predictions. This paper proposes a novel learnable test-time augmentation, namely OptTTA, tailored specifically to alleviate large domain shifts for the source-free medical image segmentation task. OptTTA enables efficiently generating augmented views of test input, resembling the style of private source images and bridging a domain gap between training and test data. Our proposed method explores optimal learnable test-time augmentation sub-policies that provide lower predictive entropy and match the feature statistics stored in the BatchNorm layers of the pretrained source model without requiring access to training source samples. Thorough evaluation and ablation studies on challenging multi-center and multi-vendor MRI datasets of three anatomies have demonstrated the performance superiority of OptTTA over prior-arts test-time augmentation and model adaptation methods. Additionally, the generalization capabilities and effectiveness of OptTTA are evaluated in terms of aleatoric uncertainty and model calibration analyses. Our PyTorch code implementation is publicly available at https://github.com/devavratTomar/OptTTA.


page 3

page 7

page 8

page 17

page 22

page 23

page 24

page 25


Single-domain Generalization in Medical Image Segmentation via Test-time Adaptation from Shape Dictionary

Domain generalization typically requires data from multiple source domai...

TeSLA: Test-Time Self-Learning With Automatic Adversarial Augmentation

Most recent test-time adaptation methods focus on only classification ta...

Test-time image-to-image translation ensembling improves out-of-distribution generalization in histopathology

Histopathology whole slide images (WSIs) can reveal significant inter-ho...

Improved Test-Time Adaptation for Domain Generalization

The main challenge in domain generalization (DG) is to handle the distri...

Scale-aware Test-time Click Adaptation for Pulmonary Nodule and Mass Segmentation

Pulmonary nodules and masses are crucial imaging features in lung cancer...

Decoupling Predictions in Distributed Learning for Multi-Center Left Atrial MRI Segmentation

Distributed learning has shown great potential in medical image analysis...

Fourier Test-time Adaptation with Multi-level Consistency for Robust Classification

Deep classifiers may encounter significant performance degradation when ...

Please sign up or login with your details

Forgot password? Click here to reset