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

Conference: ECCV 2026
Paper: Official paper page · Paper PDF
Code: https://github.com/yuyi1005/AtlasSegFM
Area: Medical Imaging
Keywords: One-shot segmentation, anatomical atlas, medical image registration, test-time adaptation, adaptive fusion

TL;DR¶

AtlasSegFM converts one annotated anatomical atlas into automatic prompts for a frozen segmentation foundation model and learns voxel-wise fusion at test time, raising SegRap Dice from 50.02 for five-click nnInteractive to 69.19 while retaining a dependence on anatomical correspondence between support and query.

Background & Motivation¶

Medical segmentation foundation models reduce the need to train a separate network for every organ, but they do not eliminate ambiguity about the intended task. A click identifies a location without necessarily specifying the structure's extent, boundary convention, or required subdivisions. Models such as nnInteractive and MedSAM2 can interpret such prompts using familiar anatomy from pretraining, yet small organs-at-risk and unusual anatomical contexts may not match those expectations. Repeating the interaction for every patient also limits efficiency when many scans share the same clinical segmentation requirement.

In-context learning offers a richer task description: an example image with its desired mask. Methods such as UniverSeg and Iris learn to use these examples through dedicated cross-task pretraining, which requires substantial curated supervision and does not directly exploit every newly available foundation model. Classical atlas registration has the complementary advantage of transferring a label without learning the target category first. Its weakness is geometric: imperfect alignment produces misplaced boundaries and loses local detail even when the transferred mask captures the correct overall anatomy.

AtlasSegFM assigns different responsibilities to these two sources rather than training another universal segmenter. The atlas specifies the task and global structure, while the foundation model examines the query image for local refinement. “One-shot” means one annotated support case per context, not a single forward pass without parameter updates. Core idea: use a registered atlas both as an automatic prompt and as a structural prediction, then learn spatially adaptive fusion with a frozen foundation model using a cycle back to the only available support annotation.

Method¶

Overall Architecture¶

The inputs are one support image with its segmentation mask and an unlabeled query volume; the output is a target-structure mask in query coordinates. For a multiclass dataset, each category is treated as a separate segmentation context rather than assuming one prompt jointly specifies all classes. The pipeline comprises atlas registration and prompt generation, frozen foundation-model refinement, and atlas–model adaptive fusion. Registration provides both the coarse query-space mask and the transformation needed to return a prediction to support coordinates for supervision.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Support image and label<br/>Query image"] --> B["Atlas Registration<br/>and Prompt Generation"]
    B --> C["Frozen Foundation-Model<br/>Refinement"]
    B -->|Atlas mask| D["Atlas–Model<br/>Adaptive Fusion"]
    C -->|Model prediction| D
    D --> E["Query segmentation"]
    D -.->|Warp back to support space| F["Support-label supervision"]
    F -.->|Update fusion estimator only| D

The segmentation foundation model remains frozen throughout. The registration network and lightweight fusion estimator are optimized during inference, so this is not an optimization-free method. The feedback loop in the diagram concerns fusion; registration has its own support–query image-similarity objective. For the non-promptable vesselFM, the second stage uses the query image directly, and atlas information still enters through fusion.

Key Designs¶

1. Atlas Registration and Prompt Generation: express the task in query coordinates

A support mask cannot simply be overlaid on another patient because pose, scale, and anatomy differ. The method applies rigid and affine preregistration before nonrigid alignment with the Recursive Deformable Pyramid (RDP) network. RDP uses a dual-stream encoder and a pyramid decoder to estimate deformation from coarse to fine. The contribution here is its use as a test-time optimizer for the current atlas–query pair, rather than a newly proposed registration architecture or a requirement for a target-domain training collection. Preregistration removes large global discrepancies before the network attempts finer deformation.

The objective in the paper's Equation (2) is:

\[ T^\star=\arg\min_T\mathcal{L}_{\mathrm{reg}}\bigl(T(X_{\mathrm{atlas}}),X_{\mathrm{query}}\bigr). \]

The optimization compares images, not query segmentation labels; the implementation uses normalized cross-correlation (NCC). Applying the resulting transformation to the support annotation produces a coarse atlas mask in query space. This mask communicates location, extent, and plausible shape, making it a more explicit task example than a sparse point. Its usefulness depends on correspondence: alignment cannot create a lesion concept that the support anatomy does not contain.

The mask is converted to the interface supported by the chosen foundation model. A click is derived from the centroid of its largest connected component, a bounding box from its middle axial slice, or the mask itself is used directly. Point and box prompts follow the single-slice protocol: nnInteractive infers the 3D volume, whereas MedSAM2 propagates from the middle slice in both directions. Mask prompts are the default in the reported promptable-model experiments. Crucially, this does not require a manually annotated query mask; the dense prompt is transferred automatically from the support case.

2. Frozen Foundation-Model Refinement: revise the transferred label using image evidence

The registered atlas is structurally informative but may have inaccurate boundaries or locally misplaced anatomy. For nnInteractive and MedSAM2, the query image and generated prompt enter the frozen foundation model, which predicts a new mask using the actual query appearance. Registration therefore supplies an informed initialization rather than the final segmentation decision. The annotation budget changes from repeated per-query interactions to one complete support annotation per context; these costs should not be treated as equivalent merely by counting clicks.

Refinement is not guaranteed to correct every atlas error. A foundation model may prefer an anatomical whole over an artificial subdivision encoded by the support annotation, or it may poorly represent a rare structure. The framework therefore retains the atlas mask as an independent source for the next stage rather than discarding it after prompting. The vascular branch uses vesselFM without prompts: its query-only prediction is combined with the atlas afterward, so improvements there cannot be attributed to atlas prompting of vesselFM.

3. Atlas–Model Adaptive Fusion: learn where each source is more reliable

A fixed blend assumes that atlas and foundation-model reliability are spatially uniform. That assumption is inappropriate when the atlas captures global structure but misses a boundary, or when the model captures appearance but misunderstands the requested extent. A lightweight 3D encoder predicts a voxel-wise gate, with sigmoid constraining its values to \([0,1]\). Larger gate values favor the atlas; smaller values favor the foundation-model prediction. The local text extraction loses operators in Equation (4), and Equations (1) and (3) are also damaged; no unverifiable fusion equation is reconstructed here, while the direction and purpose of the gate are explicit in Section 3.3.

The estimator receives more than two hard masks. Its feature volume includes the normalized query image, both predictions in logit form, their signed logit difference, disagreement between the soft masks, and each prediction's entropy map. Image appearance provides anatomical evidence, disagreement identifies contested regions, and entropy supplies uncertainty cues. These features are processed jointly rather than treating the lower-entropy source as automatically correct. The resulting “reliability” gate is learned through adaptation; the paper does not establish that it is a statistically calibrated probability of correctness.

Supervision comes from returning the fused query prediction to support space. After inverse warping, a Dice loss compares that prediction with the available support label, and only the fusion estimator is updated. This is not a direct comparison of the support label with itself: the prediction has passed through deformation, foundation-model refinement, fusion, and inverse mapping. Nevertheless, the supervision still depends on the same support case and on registration quality, rather than providing an independent check against query truth.

A Worked Example¶

Consider segmenting a target pelvic structure in a Pengwin CT volume using one annotated support CT. Rigid, affine, and nonrigid registration transfer the support label into the new scan, giving a mask that specifies the target but may miss its precise boundary. The transferred mask prompts frozen nnInteractive to recover local structure from the query image. The fusion estimator then compares both predictions alongside query appearance and produces a spatial gate. During adaptation, each fused prediction is returned to support coordinates for supervision; the delivered result remains the fused mask in query coordinates. This example explains the data flow without misrepresenting dataset-average scores as successive measurements on a single patient.

Loss & Training¶

Registration uses NCC with learning rate \(10^{-4}\) for 300 iterations on the current support–query pair.
Fusion uses the cycle-transformation Dice loss with learning rate \(10^{-5}\) for 100 iterations.
The segmentation foundation model is frozen, and query ground-truth masks are not used in either optimization stage.
No additional training dataset is required, but both learnable components incur test-time optimization rather than only feed-forward inference.

Key Experimental Results¶

Main Results¶

Section 4 evaluates HaN-Seg, SegRap, Pengwin, AVT, and the in-house Fe-MRA and BrainRT datasets. The protocol uses five-fold cross-validation with one support example drawn from the remaining validation set for each fold, treating each category as a context. The authors state that all six datasets are excluded from the nnInteractive training corpus. The selected results below come from Tables 1–3: Dice and normalized surface Dice (NSD) are percentages, higher is better, and differences are Dice percentage points. These are benchmark results rather than prospective clinical validation.

Dataset / task	Foundation model used by AtlasSegFM	Comparator	Comparator Dice ↑	AtlasSegFM Dice ↑	Difference	AtlasSegFM NSD ↑
HaN-Seg / organs-at-risk (Table 1)	nnInteractive	nnInteractive-5	55.29	63.71	+8.42	77.80
SegRap / organs-at-risk (Table 1)	nnInteractive	nnInteractive-5	50.02	69.19	+19.17	87.39
Pengwin / pelvic bones (Table 1)	nnInteractive	nnInteractive-5	93.63	95.50	+1.87	96.96
BrainRT / organs-at-risk (Table 2)	nnInteractive	nnInteractive-5	39.09	77.07	+37.98	77.55
AVT / vessels (Table 3)	vesselFM	nnInteractive-5	83.38	81.34	-2.04	Not reported

Interactive-baseline clicks are sampled from query ground truth, providing an optimistic localization signal unavailable in deployment; AtlasSegFM generates its query prompts through registration. Conversely, AtlasSegFM receives a complete support annotation and performs test-time optimization, and the AVT row also compares different foundation models. The result is therefore not a uniform victory from changing prompts under identical annotation and compute budgets: its AVT Dice is lower than nnInteractive-5. On that dataset, however, Table 3 gives AtlasSegFM clDice of 72.04 and HD95 of 30.84, versus 70.52 and 43.69 for nnInteractive-5. Centerline Dice measures vessel connectivity agreement and is higher-is-better; HD95 measures boundary discrepancy and is lower-is-better, so overlap and structural metrics need not rank methods identically.

Ablation Study¶

The following compact selection uses the Pengwin and HaN-Seg columns of Table 4, with Dice in percent and higher being better. The first block studies registration, while the second starts from a standalone foundation model; the rows are not one uninterrupted single-factor training sequence.

Configuration (Table 4)	Pengwin Dice ↑	HaN-Seg Dice ↑
Atlas registration only	63.24	37.32
Add rigid preregistration	70.55	41.89
Also add affine preregistration	82.75	48.72
Standalone nnInteractive, one point	72.09	42.42
Use atlas-derived prompts	94.56	61.74
Add adaptive fusion	95.50	63.71

On Pengwin, atlas prompting adds 22.47 points over the one-point baseline, and fusion adds another 0.94 points; fusion adds 1.97 points on HaN-Seg. The evidence favors getting the context right before refining the local combination, rather than attributing the entire improvement to the final gate. Table 5 reports approximately 1.8 minutes on an RTX 4090 for one \(256\times256\times240\) BrainRT organs-at-risk image: 1.5 minutes for registration, 0.3 for fusion, and 0.01 for foundation-model inference. The 8.6M learnable parameters cover registration and fusion, not the total system including its frozen foundation model, and this timing is not an average over every dataset.

Key Findings¶

Atlas prompting produces a larger improvement than final fusion, while preregistration substantially strengthens the atlas itself; these components solve different problems.
Small-structure gains are substantial, including 77.07 Dice for BrainRT organs-at-risk, but segmentation accuracy alone does not demonstrate improved clinical decisions.
Table 3 reports AVT clDice as 72.04, whereas the adjacent prose says 72.07; this edition follows the table and does not repeat the unsupported claim of leading every Dice comparison.

Highlights & Insights¶

The atlas is an executable task example rather than merely a semantic description. Spatial alignment turns its extent and shape conventions into inputs that existing models can consume.
The atlas remains an independent prediction after serving as a prompt. This preserves a structural alternative when the foundation model departs from the intended task.
Cycle supervision reuses the only annotation for query-space adaptation. It saves query-label acquisition without removing assumptions about registration or support selection.

Limitations & Future Work¶

The authors explicitly identify lesion segmentation and major anatomical differences as weaknesses because they violate support–query structural consistency; proposed directions include multiple atlases and improved prompt generation.
Reported failures include merging the left and right mandibles and missing part of the aortic arch. These expose conflicts between annotation conventions and model priors, as well as failures of support–query matching.
As a reader interpretation, the cycle loss inherits support-label and deformation errors rather than independently validating query accuracy. Approximately 1.8 minutes of adaptation also should not be described as real-time interaction.
There is another unresolved numerical inconsistency: standalone vesselFM obtains Fe-MRA Dice of 60.31 in Table 3 but 55.28 in Table 4, so no improvement is calculated from that baseline here. The cache contains the full main method and experiments but no separate supplement; details such as the NSD tolerance are not specified in the inspected main text.

Versus RDP and classical atlas segmentation: AtlasSegFM reuses a registration architecture as a per-pair test-time optimizer and sends its transferred labels through a foundation model, gaining refinement at the cost of per-case optimization.
Versus nnInteractive and MedSAM2: The contribution is not another interactive backbone; it is automatic contextual prompting and retained atlas evidence for fusion, with the backbone's anatomical preferences still present.
Versus UniverSeg and Iris: Instead of training a dedicated in-context segmenter, the method connects existing foundation models to explicit spatial correspondence, exchanging retraining requirements for a structural-consistency assumption.

Rating¶

Novelty: 3.5/5 — A useful integration of established registration, foundation models, and test-time cyclic fusion rather than wholly new components.
Experimental Thoroughness: 3.5/5 — Six datasets and informative component studies, tempered by numerical inconsistencies and limited support-selection or deployment analysis.
Writing Quality: 3/5 — The pipeline and implementation are understandable, but table–text discrepancies and damaged extracted equations complicate verification.
Value: 4/5 — A practical direction for single-example customization of 3D anatomical segmentation, within tasks that admit meaningful registration.