Atlas is Your Perfect Context: One-Shot Customization for Generalizable Foundational Medical Image Segmentation

AI summary

Overview

Research area: Medical image segmentation, foundation models, one-shot / in-context learning, atlas-based registration.

Technical level: Intermediate (assumes familiarity with segmentation metrics, foundation models, and image registration concepts, but the core idea is explained intuitively).

One-sentence scope: The paper introduces AtlasSegFM, a framework that customizes off-the-shelf medical segmentation foundation models to a new clinical task using a single annotated example (an atlas) plus atlas-query registration and a lightweight test-time fusion module.

What This Paper Is About

Accurate segmentation of anatomical structures in medical images is essential for diagnosis and treatment planning, but interactive foundation models still need precise prompts and can fail in rare or underrepresented clinical contexts such as small organs-at-risk. The authors propose AtlasSegFM, which uses one annotated example ("atlas") and atlas-query registration to automatically generate context-aware prompts, then refines and fuses the result with a frozen foundation model. The goal is a lightweight, deployable, one-shot customization method that generalizes to new clinical contexts without training on a dedicated dataset.

Key Contributions

1. The authors investigate one-shot customization of segmentation foundation models, adapting a frozen foundation model to a new clinical context using a single annotated atlas, with no dedicated training set.
2. They present AtlasSegFM, a pipeline that combines the global anatomical consistency of atlas registration with the local refinement capability of foundation models, integrated through a lightweight adaptive fusion module at test time.
3. They evaluate on six public and in-house datasets spanning organs-at-risk, pelvic anatomy, and vascular trees, reporting consistent gains in accuracy and boundary quality, with the largest improvements in underrepresented anatomical contexts.
4. They show an inference-only design with 8.6M learnable parameters and report runtime and parameter comparisons against prior methods.

Main Findings

• Organs-at-risk and pelvic bone results: On HaN-Seg, SegRap, and Pengwin, the nnInteractive-based version of AtlasSegFM achieves the highest mean Dice among in-context models (63.71%, 69.19%, and 95.50%, respectively). This is compared against the supervised upper bound nnUNet (75.10, 80.20, 98.90).

• Gains over recent ICL methods: AtlasSegFM improves mean Dice on HaN-Seg by 26.14% over Iris and 29.72% over Tyche.

• Brain datasets: On BrainRT whole-brain segmentation, AtlasSegFM reaches a Dice of 91.24% and HD of 6.78. For organs-at-risk (e.g., optic nerves, brainstem), it achieves a Dice of 77.07%, NSD of 77.55%, and HD of 5.17.

• Vessel datasets: On AVT and Fe-MRA, AtlasSegFM attains Dice of 81.34% and 84.42%, clDice of 72.04% and 82.99%, and HD of 30.84 and 3.00, respectively.

• Interactive prompt weakness: Interactive methods perform well on common scenarios (e.g., 88.21% Dice for whole brain with 5-click interaction on BrainRT) but drop to 26.91% Dice for organs-at-risk.

• Ablation findings: Starting from atlas-only registration (63.24% Pengwin, 37.32% HaN-Seg, 35.73% Fe-MRA), adding rigid pre-registration improves to 70.55%, 41.89%, 66.43%; adding affine pre-registration reaches 82.75%, 48.72%, 81.64%. Prompts from the atlas raise the foundation-model baseline (72.09% Pengwin) to 94.56%, and fusion produces the final 95.50% Pengwin, 63.71% HaN-Seg, and 84.42% Fe-MRA.

• Prompt type comparison (Pengwin, nnInteractive): Atlas-derived prompts yield 92.60 Dice (click), 88.96 Dice (box), and 95.77 Dice (mask), with mask prompts performing best.

• Efficiency: Total inference time is 1.8 minutes per image on an NVIDIA 4090 GPU for a 256×256×240 BrainRT image, versus 18.1 minutes for SegGPT, 2.2 minutes for UniverSeg, and 2.1 minutes for Tyche. Total learnable parameters are 8.6M (registration 8.5M, fusion 0.1M), versus 1.4G for SegGPT and 1.2M for UniverSeg and Tyche.

• Robustness to imperfect alignment: Dice degrades more slowly for AtlasSegFM than for registration alone as MSE increases; in the highlighted case, the method improves Dice from 0.22 to 0.64 after foundation-model refinement and adaptive fusion.

• Overall average: Across the six datasets, the paper reports an average Dice of 80.35%, and a Dice of 73.60% in rare contexts underrepresented in foundation model pretraining (organs-at-risk).

• Failure cases: The method erroneously merges left and right mandibles on SegRap (annotation separates artificially) and fails to fully recover the aortic arch on AVT when support and query images differ substantially.

Methodology in Plain English

The pipeline works in three steps. First, the atlas (a single annotated image, or support pair) is aligned to the unlabeled query image using a registration network derived from the Recursive Deformable Pyramid (RDP), repurposed for test-time optimization. Rigid and affine pre-registration are applied first for coarse alignment, then the network is optimized with a normalized cross-correlation loss (learning rate 1e-4, 300 iterations) to produce a coarse mask. Second, this coarse mask is converted into a prompt—a click (centroid of the largest region), a box (bounding box at the middle z-slice), or a mask (the mask itself)—and fed to a frozen foundation model (nnInteractive for non-vascular tasks, vesselFM for vascular tasks; vesselFM does not accept prompts). Third, a lightweight adaptive fusion module combines the atlas mask and the foundation-model mask using a voxel-wise reliability gate, expressed as M_final = M_fm + K·(M_atlas − M_fm), where K is estimated by a small 3D encoder that takes the query image, both predictions, their signed logit difference, voxel-wise disagreement, and entropy maps as inputs. The fusion estimator is optimized at test time with a cycle-transformation Dice loss (it warps the fused prediction back to the support space and computes Dice against the support label), updating for 100 iterations at a learning rate of 1e-5 while the foundation model remains frozen.

Why This Matters

Impact on research: The paper shows that combining classical atlas-based registration priors with modern foundation models can match or approach fully supervised performance without training on the target domain, and it highlights that interactive foundation models degrade sharply on small or rare structures. This provides a template for one-shot customization of existing models rather than retraining new ones.

Real-world applications (bullets):

• Radiotherapy planning, especially organs-at-risk delineation (eyes, optic nerves, optic chiasm, brainstem) where under-segmentation risks radiation-induced damage.
• Vascular segmentation in angiography, including the aortic vessel tree (AVT) and lower-limb arterial and venous structures (Fe-MRA) relevant to conditions such as varicose veins and arterial thrombosis.
• Pelvic bone segmentation for pelvic anatomical planning.
• Head and neck and nasopharyngeal carcinoma radiotherapy, where organs-at-risk are densely packed and low-contrast.

Industry relevance: The method uses only a single annotated example per context, requires no dedicated training set, keeps the foundation model frozen, and runs on a single NVIDIA 4090 GPU with 8.6M learnable parameters and 1.8 minutes per image, making it deployable in clinical workflows where large-scale retraining is impractical.

Future Directions

• Multi-atlas strategies for greater robustness, particularly for patients whose anatomy differs markedly from the atlas.
• Refining prompt generation for more context-aware guidance, building on the click/box/mask comparison.
• Extending the framework beyond organs to lesion segmentation, which the authors note is limited by the dependence on structural consistency that focal lesions often lack.
• Improving registration and matching for cases where support and query images differ substantially, such as the AVT aortic arch failure and the SegRap mandible-merging failure.

Target Audience

Researchers and practitioners in medical image analysis, computer vision, and clinical imaging who work on foundation models, in-context or one-shot learning, image registration, or radiotherapy and vascular segmentation. It is most useful for readers with intermediate background in segmentation metrics (Dice, NSD, HD, clDice) and image registration, and for clinicians or engineers seeking lightweight, deployable customization of existing segmentation models.

Authors’ abstract

Accurate segmentation of anatomical structures in medical images is essential for diagnosis and treatment planning. While recent interactive segmentation foundation models enhance generalization through large-scale multimodal pretraining, they still depend on precise prompts and can fail in underrepresented clinical contexts (e.g., small organs-at-risk). We present AtlasSegFM, an atlas-guided framework that customizes off-the-shelf foundation models to new clinical contexts with a single annotated example. AtlasSegFM 1) performs atlas-query registration to generate context-aware prompts, 2) refines the segmentation with a frozen foundation model, and 3) applies a lightweight adaptive fusion module to combine atlas priors with foundation-model inputs and predictions. Extensive experiments on six public and in-house datasets across radiotherapy and vascular scenarios show consistent gains, with the largest improvements on small and delicate structures. AtlasSegFM provides a lightweight, deployable solution for one-shot customization of segmentation foundation models in real-world clinical workflows.