PASS:Test-Time Prompting to Adapt Styles and Semantic Shapes in Medical Image Segmentation

TL;DR

PASS introduces a test-time adaptation framework for medical image segmentation that jointly adapts style and shape features using learned prompts, significantly improving performance across diverse datasets without retraining.

Contribution

The paper proposes a novel prompt-based TTA method that leverages shape-aware prompts and style adaptation, addressing limitations of existing methods in medical image segmentation.

Findings

PASS outperforms state-of-the-art TTA methods on multiple datasets.

The shape-aware prompts effectively handle shape variability across domains.

The input decorator dynamically adjusts style prompts based on input data.

Abstract

Test-time adaptation (TTA) has emerged as a promising paradigm to handle the domain shifts at test time for medical images from different institutions without using extra training data. However, existing TTA solutions for segmentation tasks suffer from (1) dependency on modifying the source training stage and access to source priors or (2) lack of emphasis on shape-related semantic knowledge that is crucial for segmentation tasks.Recent research on visual prompt learning achieves source-relaxed adaptation by extended parameter space but still neglects the full utilization of semantic features, thus motivating our work on knowledge-enriched deep prompt learning. Beyond the general concern of image style shifts, we reveal that shape variability is another crucial factor causing the performance drop. To address this issue, we propose a TTA framework called PASS (Prompting to Adapt Styles and Semantic shapes), which jointly learns two types of prompts: the input-space prompt to reformulate the style of the test image to fit into the pretrained model and the semantic-aware prompts to bridge high-level shape discrepancy across domains. Instead of naively imposing a fixed prompt, we introduce an input decorator to generate the self-regulating visual prompt conditioned on the input data. To retrieve the knowledge representations and customize target-specific shape prompts for each test sample, we propose a cross-attention prompt modulator, which performs interaction between target representations and an enriched shape prompt bank. Extensive experiments demonstrate the superior performance of PASS over state-of-the-art methods on multiple medical image segmentation datasets. The code is available at https://github.com/EndoluminalSurgicalVision-IMR/PASS.