DART: Diffusion for Anatomical Reshaping with Text

DART is one text-conditioned 3D diffusion model that learns the shapes of 377 anatomical classes from segmentation probability maps. It corrects the output of any segmentation model, given the probability map and a description of the class. It never sees the image, so a model trained on CT predictions also corrects MRI predictions and lesion classes it was never trained on.

probability map ─► threshold 0.5 ─► bounding box + 10% ─► crop & resize ─► noise to t* ─► denoise ─► resample ─► mask
                                                                         "pancreas, elongated lobulated gland, ..."

Install

pip install -e .          # Python >= 3.9, PyTorch >= 2.1 (CUDA, ROCm, Apple MPS or CPU)

Correct a prediction

from dart import DART

dart = DART.from_pretrained("Koushik45048545309/DART")   # local folders and best.pt files work too
mask = dart.correct(prob, structure="pancreas")         # prob: [z,y,x] array in [0,1]; binary masks also work

dart correct --model Koushik45048545309/DART --input pancreas.nii.gz --structure pancreas --out pancreas_dart.nii.gz
dart correct --model Koushik45048545309/DART --input totalseg_case/ --out corrected/   # one file per class

A class outside the vocabulary just needs a description: dart.correct(prob, description="left kidney cyst, rounded fluid-filled lesion, within or on the surface of the left kidney", tstar=350).

Train and adapt

I want to…	Guide
Use DART from Python or the command line	docs/inference.md
Prepare my own data (probability maps, masks, nnU-Net outputs)	docs/data.md
Train from scratch	docs/training.md
Fine-tune the released model, add classes, adapt to a new segmenter or modality	docs/finetuning.md
Choose the noise level t* for each class	docs/noise-levels.md
Measure the gain against reference labels	docs/evaluation.md
Understand the checkpoint files / convert old checkpoints	docs/checkpoints.md
Run on a Slurm cluster (LUMI scripts included)	docs/hpc.md
Publish code and checkpoints on Hugging Face	docs/huggingface.md

Every command is listed by dart -h; every option by dart <command> -h.

Checkpoints

checkpoints/main/ is loaded by default. All trained runs, the data-scaling runs and the baselines are listed in checkpoints/README.md.

Repository

dart/           package: model.py (U-Net + text conditioning), diffusion.py, crop.py, data.py, train.py,
                inference.py, evaluate.py, checkpoint.py, hub.py, cli.py
configs/        dart128.yaml (paper, Table 7), dart64.yaml, finetune.yaml
scripts/        lumi/ (setup, training and inference Slurm jobs), upload_to_hf.sh
docs/           the guides above

Citation

Paper under double-blind review; citation to follow.