cs.CVOct 20, 2026
SaveDART: Scaling a Shape Representation Across Anatomical Classes and Imaging Modalities
Organizations: Johns Hopkins University · ELLIS Institute Finland · Aalto University · Johns Hopkins Medicine
Abstract
Current state-of-the-art segmentation models still produce anatomically implausible shapes, such as broken vessels and fragmented glands, on the structures that are hardest to segment. Learned shape representations capture what structures look like and can help correct these errors, but most cover a fixed set of anatomical classes. Extending this coverage across classes and imaging modalities is costly because it often requires new voxel-level annotations. We propose Diffusion for Anatomical Reshaping with Text (DART), which learns a shared shape representation from segmentation probabilities. It uses anatomical descriptions to correct these probabilities without receiving medical images. Across 13 public benchmarks, absolute Dice gains reach 31% for individual gallbladder cases in CT and 10% for individual prostate cases in MRI. Using existing segmentation predictions, we scale a single shape representation from 50 to 377 anatomical classes without new voxel-level annotations for correction training. The shared representation also corrects four lesion classes absent from training and transfers from CT to MRI without retraining. Corrected labels also improve segmentation models, raising one model's Dice from 57.1% to 67.5%. The same pipeline supports annotations for 45,000 CT scans. These results suggest that sharing shape information within and across classes can reduce the need for new voxel-level annotations when extending segmentation to more anatomical classes and imaging modalities. The model is available at https://huggingface.co/Koushik45/DART-Seg
Figures & tables
Figure 1: Segmentation quality is uneven across anatomy. Best reported Dice per class, from the Touchstone benchmark ( Bassi et al., 2024 ) and TotalSegmentator ( Wasserthal et al., 2023 ) . Large solid organs such as the liver and spleen exceed . Most other structures, including tubular structures, small glands and hollow organs, score lower. The outlined structures, all below , are the ones DART is designed to help; the dots indicate that its vocabulary extends to classes.
Figure 2: Overview of DART. Training (top): a segmentation probability map predicted by a segmentation model is cropped, resized to , and corrupted with Gaussian noise. A 3D U-Net learns to reverse the corruption. Conditioning (middle, blue): an anatomical description is encoded by a frozen text encoder, projected by a trainable MLP, added to the timestep embedding, and injected into every residual block. Inference (bottom): probabilities are thresholded at only to locate the crop. Their original values within the crop are resized, noised up to , and denoised under the same description. The final result is thresholded to obtain a segmentation mask. Blue components are contributed by this work.
Figure 3: On five hard structures, DART outperforms the two strongest baselines on every metric. Predictions from TotalSegmentator, AtlasNet and WholeBodySeg on the external benchmarks, averaged over the three models and over the benchmarks that annotate each structure. Each panel compares DART with the two baselines with the highest mean Dice for that structure, named in its legend. Dice, IoU, clDice and NSD (within mm) use the left axis (higher is better); HD95 in mm uses the right axis (lower is better).
Figure 4: One frozen model transfers to new tasks, and its corrected labels train better models. (a) Dice before and after DART on the seven transfer tasks with MedSAM input, six in MRI and one in CT (pancreatic lesion); four are lesion classes, and none of the seven appears in training. (b) Seven of the fifteen configurations of Table 8 , trained on original and on corrected labels with every other factor fixed; nnU-Net is the five-level DynUNet.
Figure 5: The model keeps improving with more training examples. (a) Dice after DART (on a – scale) for anatomical classes when the model is trained on to of the available training examples: individual classes (thin lines), their mean with a confidence interval (heavy line) and a saturating power-law fit to the mean (dashed; Appendix I ). The mean curve passes the mean input Dice of at about of the training examples. (b) The same runs grouped by anatomy.
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Composition of the training corpus. (a) Volumes contributed by each public dataset and the pooled total. (b) Annotated classes per dataset. (c) Number of anatomical classes in the training vocabulary, compared with the number of classes that other mask-based models can be conditioned on; SegRefiner is class-agnostic and Post-DAE is not a diffusion model, so this panel compares vocabulary size only, not quality.
Figure 7: DART raises mean Dice on every external benchmark. Mean Dice of the input (light) and after DART (dark). The axis starts at . CT inputs are the mean over TotalSegmentator, AtlasNet and WholeBodySeg, except KiTS23, which uses TotalSegmentator on its abnormal cases. MRI inputs come from TotalSegmentator MRI, and CHAOS averages its CT and MRI scans. FLARE22 and MSD pancreas appear in Figure 3 .
| Anatomical class | TotalSegmentator | + DART | Gain |
| Adrenal glands | 64.70 | 70.46 | |
| Duodenum | 60.44 | 64.50 | |
| Esophagus | 76.51 | 80.18 | |
| Gallbladder | 66.85 | 70.20 | |
| Colon | 76.43 | 79.54 | |
| Pancreas | 79.67 | 82.74 |
Table 2: DART on TotalSegmentator predictions. Dice before and after DART on held-out WORD and AMOS22 cases, sorted by gain. is the standard deviation of the per-case paired gain. Mean Dice improves for every evaluated class.
| Anatomical class | AtlasNet | + DART | Gain |
| Rectum | 72.13 | 75.64 | |
| Adrenal glands | 69.43 | 72.88 | |
| Esophagus | 76.89 | 78.41 | |
| Gallbladder | 79.14 | 80.59 | |
| Duodenum | 72.75 | 73.70 | |
| Bladder | 85.60 | 86.35 |
Table 3: DART on AtlasNet predictions. Same protocol as Table 2 . Mean Dice improves for thirteen of fourteen classes; the intestine loses points, which is within the per-case spread.
Figure 8: DART improves the predictions of two automatic segmentation models. (a) Dice gain per class over each model’s own prediction, on held-out WORD and AMOS22 cases with the per-class unchanged. (b) Mean gain against the two controls that use the same inputs: the identity path with the diffusion model removed, and morphological closing with its structuring element swept per class on the validation split.
Figure 9: DART follows what the description says, not its wording. Change in Dice relative to the matching description, on the predictions of the automatic segmentation models of Appendix D , with the noise level and seed held fixed. (a) A paraphrase that keeps the anatomy but changes the words barely changes Dice ( points); substituting another structure’s description costs points. Dots are individual anatomical classes. (b) The same, per class. Every class loses Dice under a mismatched description; none moves by more than points under a paraphrase.
| Anatomical class | Paraphrased | Mismatched |
| Adrenal glands | ||
| Bladder | ||
| Colon | ||
| Duodenum | ||
| Esophagus | ||
| Femur (left) | — |
Table 4: Description sensitivity. Change in Dice (points) relative to the matching description, averaged over the two segmentation models of Appendix D . ‘—’ means the class was not included in that run.
Figure 10: Boundary and fragmentation metrics. (a) Average symmetric surface distance before and after DART on held-out WORD cases. (b) Connected components per mask on the same cases. (c) Mean Dice on DAP-Atlas for the MedSAM input and after DART.
| Training vocabulary | Test classes | Dice |
| 50 | 50 | 73.85 |
| 377 | 50 | 74.20 |
Table 5: Expanding the vocabulary preserves mean Dice on existing classes. Mean Dice on the same anatomical classes, on TotalSegmentator, AMOS22 and WORD test cases, for models trained on anatomical classes and on the full vocabulary.
| Model | Dice | Gain |
| Input (before DART) | 68.71 | — |
| Index-conditioned | 69.10 | |
| Text-conditioned (DART) | 74.20 |
Table 6: Anatomical descriptions outperform class indices. Replacing the text embedding with a learned class embedding indexed by an integer class identifier, with everything else fixed, reduces the gain from to Dice points on TotalSegmentator, AMOS22 and WORD test cases. The evaluation set includes classes outside the index-conditioned model’s vocabulary; see the text below.
Figure 11: Selected noise levels follow anatomy. (a) Validation-selected for the training vocabulary, stacked by anatomical group. (b) Mean per group. Compact solid organs need the least noise; muscles and vessels need the most. The per-class values are released with the code.
| Architecture | |
| Backbone | 3D MONAI DiffusionModelUNet |
| Resolution levels | 4, with channels |
| Input | 1-channel probability map, cropped and resized to |
| Text conditioning | |
| Text encoder | all-MiniLM-L6-v2 , frozen |
| Text embedding | 384-dim, computed once and stored |
Table 7: Training and inference configuration.
| Configuration | Original labels | Corrected labels | Gain |
| Original-label arm at least 45 Dice | |||
| Attention U-Net (wide) | 62.3 | 72.7 | |
| Attention U-Net | 57.1 | 67.5 | |
| Residual U-Net | 57.2 | 66.3 | |
| U-Net | 52.6 | 60.4 | |
| DynUNet (4 levels) | 67.2 | 71.5 | |
Table 8: Corrected labels improve all fifteen segmentation model configurations. Dice of each configuration trained on the original labels and on the same labels after DART, with one shared training schedule and every other factor fixed. Configurations are grouped by whether the original-label arm reaches Dice.
Figure 12: Architecture of DART-Seg ( Rokuss et al., 2026 ) . Left: a frozen text encoder embeds a text prompt naming the target structure, and a trainable Prompt Decoder produces text features that modulate the image decoder. Right: at each decoder level, upsampled features are concatenated with the encoder skip connection and modulated by the text features before the segmentation head, with deep supervision at every level. We use the architecture unmodified and train it on DART-corrected training labels.
| Group | Anatomical class | Description |
| Solid viscus | pancreas | pancreas, elongated lobulated gland, retroperitoneum behind the stomach from duodenal curve to splenic hilum |
| Hollow organ | colon | colon, long convoluted haustrated hollow tube, framing the peripheral abdomen from cecum to sigmoid |
| Gland | adrenal_gland_left | left adrenal gland, small thin comma-shaped gland with slender limbs, superior pole of the left kidney medial to the spleen |
| Vessel | aorta | aorta, thick-walled large arterial tube arching then descending, from the left ventricle through the mediastinum to the iliac bifurcation |
| Muscle | autochthon_left | left autochthon, long paraspinal erector spinae muscle mass, alongside the left vertebral column from sacrum to skull base |
| Bone | clavicula_left | left clavicle, slender S-shaped long bone, horizontal at the left base of the neck from sternum to scapula |
Table 9: Example anatomical descriptions , one per anatomical group. The last row is a held-out lesion class, described at inference only.
| Method | Adrenal L | Adrenal R | Aorta | Atr. app. L | Autoch. L | Autoch. R |
| MedSAM (input) | 78.6 | 65.6 | 85.1 | 74.6 | 79.1 | 79.8 |
| Single-pass U-Net | 78.6 | 65.6 | 84.5 | 74.6 | 78.8 | 79.5 |
| SegRefiner † | 72.3 | 64.6 | 74.9 | 70.2 | 75.1 | 74.9 |
| Post-DAE | 49.0 | 39.0 | 53.6 | 41.7 | 62.4 | 62.7 |
| Hole filling | 78.6 | 65.6 | 85.1 | 74.6 | 79.1 | 79.8 |
| Morph. closing | 79.3 | 66.6 | 83.8 | 76.0 | 81.8 | 82.2 |
Table 10: Dice for all methods on the TotalSegmentator anatomical classes. Read in blocks of six; each column is one anatomical class. Every method starts from the MedSAM predictions. Best per column in bold blue , second best in blue . ‘—’ means the method produces no result (the two ASM variants cover of the classes). † reads image intensity. DART is the only method whose Dice is never below the input.
More in this topic
Medical Image Segmentation
Medical image segmentation delineates organs, lesions, and other structures in scans to support measurement, diagnosis, treatment planning, and research.
See all 786 papers