Each Test Image Deserves A Specific Prompt: Continual Test-Time Adaptation for 2D Medical Image Segmentation
About
Distribution shift widely exists in medical images acquired from different medical centres and poses a significant obstacle to deploying the pre-trained semantic segmentation model in real-world applications. Test-time adaptation has proven its effectiveness in tackling the cross-domain distribution shift during inference. However, most existing methods achieve adaptation by updating the pre-trained models, rendering them susceptible to error accumulation and catastrophic forgetting when encountering a series of distribution shifts (i.e., under the continual test-time adaptation setup). To overcome these challenges caused by updating the models, in this paper, we freeze the pre-trained model and propose the Visual Prompt-based Test-Time Adaptation (VPTTA) method to train a specific prompt for each test image to align the statistics in the batch normalization layers. Specifically, we present the low-frequency prompt, which is lightweight with only a few parameters and can be effectively trained in a single iteration. To enhance prompt initialization, we equip VPTTA with a memory bank to benefit the current prompt from previous ones. Additionally, we design a warm-up mechanism, which mixes source and target statistics to construct warm-up statistics, thereby facilitating the training process. Extensive experiments demonstrate the superiority of our VPTTA over other state-of-the-art methods on two medical image segmentation benchmark tasks. The code and weights of pre-trained source models are available at https://github.com/Chen-Ziyang/VPTTA.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Image Classification | PACS | Overall Average Accuracy88.12 | 230 | |
| Multi-class classification | VLCS | Acc (Caltech)97.25 | 139 | |
| Prostate Segmentation | HK (test) | DSC83.27 | 20 | |
| Prostate MRI Segmentation | Prostate MRI Dataset Domain D | Dice Coefficient83.65 | 11 | |
| Cardiac Image Segmentation | M&MS Domain D 1.0 (test) | ASSD (LV)3.29 | 11 | |
| Cardiac Image Segmentation | M&MS Domain C (target) | LV Dice82.15 | 11 | |
| Cardiac Image Segmentation | M&MS Domain C 1.0 (test) | ASSD (LV)4.24 | 11 | |
| Cardiac Image Segmentation | M&MS Average 1.0 (test) | ASSD4.25 | 11 | |
| Cardiac Image Segmentation | M&MS Domain D (target) | LV Dice84.17 | 11 | |
| Cardiac Image Segmentation | M&MS Domain B 1.0 (test) | ASSD (LV)3.86 | 11 |