Part-aware Prompted Segment Anything Model for Adaptive Segmentation

TL;DR

P^2SAM tackles the challenge of adapting segmentation models to new patients with limited annotated data by introducing a one-shot, in-context segmentation framework. It leverages a part-aware prompt mechanism that decomposes a patient’s prior anatomy into multiple foreground parts and uses their interactions with the current target image to generate robust prompts for a promptable model like SAM. A distribution-guided retrieval strategy selects the optimal number of parts at test time, mitigating prompt ambiguity and improving generalization. Empirically, P^2SAM achieves substantial gains in patient-adaptive NSCLC and polyp segmentation, attains a state-of-the-art 95.7% mean IoU on PerSeg, and demonstrates notable generalization to natural-image personalization tasks, underscoring its practical potential for cross-domain, data-efficient segmentation without fine-tuning.

Abstract

Precision medicine, such as patient-adaptive treatments assisted by medical image analysis, poses new challenges for segmentation algorithms in adapting to new patients, due to the large variability across different patients and the limited availability of annotated data for each patient. In this work, we propose a data-efficient segmentation algorithm, namely Part-aware Prompted Segment Anything Model ($P^2SAM$). Without any model fine-tuning, $P^2SAM$ enables seamless adaptation to any new patients relying only on one-shot patient-specific data. We introduce a novel part-aware prompt mechanism to select multiple-point prompts based on the part-level features of the one-shot data, which can be extensively integrated into different promptable segmentation models, such as SAM and SAM 2. Moreover, to determine the optimal number of parts for each specific case, we propose a distribution-guided retrieval approach that further enhances the robustness of the part-aware prompt mechanism. $P^2SAM$ improves the performance by +8.0% and +2.0% mean Dice score for two different patient-adaptive segmentation applications, respectively. In addition, $P^2SAM$ also exhibits impressive generalizability in other adaptive segmentation tasks in the natural image domain, e.g., +6.4% mIoU within personalized object segmentation task. The code is available at: https://github.com/Zch0414/p2sam

Figures (16)

• Figure 1: Illustration of SAM's ambiguity property. The ground truth is circled by a red dashed circle; the predicted mask is depicted by a yellow solid line.
• Figure 2: Illustration of two patient-adaptive segmentation tasks. $\text{P}^2\text{SAM}$ can segment the follow-up data by utilizing one-shot prior data as multiple-point prompts. Prior and predicted masks are depicted by a solid yellow line.
• Figure 3: Illustration of the part-aware prompt mechanism. Masks are depicted by a yellow solid line. We first cluster foreground features in the reference image into part-level features. Then, we select multiple-point prompts based on the cosine similarity ($\otimes$ in the figure) between these part-level features and target image features. A colorful star, matching the color of the corresponding part, denotes a positive-point prompt, while a gray star denotes a negative-point prompt. These prompts are subsequently fed into the promptable decoder to do prediction.
• Figure 4: Illustration of $\text{P}^2\text{SAM}$'s improvement. Wasserstein distances between the priors and results are shown in white.
• Figure 5: Illustration of the distribution-guided retrieval approach.