Pay Attention to the Atlas: Atlas-Guided Test-Time Adaptation Method for Robust 3D Medical Image Segmentation

Jingjie Guo,Weitong Zhang,Matthew Sinclair,Daniel Rueckert,Chen Chen
2024-02-27
Abstract:Convolutional neural networks (CNNs) often suffer from poor performance when tested on target data that differs from the training (source) data distribution, particularly in medical imaging applications where variations in imaging protocols across different clinical sites and scanners lead to different imaging appearances. However, re-accessing source training data for unsupervised domain adaptation or labeling additional test data for model fine-tuning can be difficult due to privacy issues and high labeling costs, respectively. To solve this problem, we propose a novel atlas-guided test-time adaptation (TTA) method for robust 3D medical image segmentation, called AdaAtlas. AdaAtlas only takes one single unlabeled test sample as input and adapts the segmentation network by minimizing an atlas-based loss. Specifically, the network is adapted so that its prediction after registration is aligned with the learned atlas in the atlas space, which helps to reduce anatomical segmentation errors at test time. In addition, different from most existing TTA methods which restrict the adaptation to batch normalization blocks in the segmentation network only, we further exploit the use of channel and spatial attention blocks for improved adaptability at test time. Extensive experiments on multiple datasets from different sites show that AdaAtlas with attention blocks adapted (AdaAtlas-Attention) achieves superior performance improvements, greatly outperforming other competitive TTA methods.
Computer Vision and Pattern Recognition,Machine Learning
What problem does this paper attempt to address?
### Problems Addressed by the Paper The paper primarily addresses the issue of performance degradation in medical image segmentation caused by the inconsistency between the target domain data and the training data distribution. Specifically: 1. **Data Distribution Discrepancy**: Medical images are typically acquired using equipment from different companies at various medical centers with different scanning protocols. This leads to variations in data distribution, causing models to perform poorly on unseen target domains. 2. **High Labeling Cost**: Directly annotating target domain data to fine-tune the model is very costly, making it impractical. 3. **Privacy Issues**: Most unsupervised domain adaptation methods require access to target domain data for training, but these data are difficult to obtain due to privacy concerns. To address these issues, the authors propose a novel graph-based Test-Time Adaptation (TTA) method called AdaAtlas. This method requires only a single unlabeled test sample as input and adjusts the segmentation network by minimizing a graph-based loss. Additionally, unlike most existing TTA methods that are limited to batch normalization blocks, AdaAtlas further leverages channel and spatial attention blocks to enhance test-time adaptability. Extensive experiments on multiple datasets demonstrate that AdaAtlas, after introducing the attention mechanism (AdaAtlas-Attention), achieves significant performance improvements, far surpassing other competitive TTA methods.