AortaSeg24: Hierarchical Semantic Learning for Multi-Class Aorta Segmentation

This repository hosts the trained nnU-Net checkpoints of our solution to the AortaSeg24 Challenge, which achieved 🥈 2nd place in the Final Test Phase (official rankings).

Training and inference code (including the customized nnU-Net fork, the hierarchical semantic loss and the cbDice loss) is at github.com/PengchengShi1220/AortaSeg24.

Repository contents

The layout mirrors nnUNet_results/ exactly, so the download can be dropped in as-is:

Dataset824_AortaSeg24_CTA_50/
└── nnUNetTrainer_Hierarchical_CBDC_AortaSeg23_NoMirroring_500epochs__nnUNetPlans__3d_fullres/
    ├── dataset.json  plans.json  dataset_fingerprint.json
    └── fold_{0,1,2,3,4,all}/
        ├── checkpoint_final.pth      # 24-class full-resolution model (23 foreground classes + background)
        └── debug.json
Dataset825_AortaSeg24_CTA_bin_50/
└── nnUNetTrainer_CE_DC_CBDC_NoMirroring_500epochs__nnUNetPlans__3d_lowres/
    ├── dataset.json  plans.json  dataset_fingerprint.json
    └── fold_{0,all}/
        ├── checkpoint_final.pth      # 2-class low-resolution model (aorta vs. background)
        └── debug.json

• Dataset824_AortaSeg24_CTA_50 — full-resolution multi-class model (3d_fullres), 5 folds + fold_all.
• Dataset825_AortaSeg24_CTA_bin_50 — low-resolution binary model (3d_lowres), fold 0 + fold_all, used to localize the aorta before the multi-class stage.
• params.json — model metadata (framework, stage/dataset mapping, folds, class counts).
• Only weights and JSON configs are published here; training logs and curve images are not included.

Both models were trained for 500 epochs with mirroring disabled (*_NoMirroring_500epochs).

Method

A hierarchical two-stage pipeline:

1. Low-resolution binary segmentation — a 3d_lowres nnU-Net segments the whole aorta (all foreground vessels) on downsampled volumes.
2. ROI extraction — the binary prediction is expanded by a scaling factor m = 2.0 to obtain a generous region of interest.
3. Full-resolution multi-class segmentation — a 3d_fullres nnU-Net predicts the 24 classes (23 aortic segments/zones + background) inside the ROI, and the result is pasted back into the original image grid.

This hierarchical design reduces inference time by more than five-fold compared with a single-stage full-volume multi-class model, while keeping full-resolution accuracy inside the ROI.

The multi-class model is trained with hierarchical semantic learning through fractal softmax and the cbDice (centerline boundary Dice) loss. See the paper for details.

Requirements

Python Version: 3.10 (same environment as the code repository)

Installation

unzip nnUNet_hierarchical_cbdc_nnUNet_v2.5.zip      # from the GitHub repository
cd nnUNet_hierarchical_cbdc_nnUNet_v2.5
pip install -e .

Quick start with these checkpoints

1. Download the weights into nnUNet_results/ (the dataset folders must land directly inside it):

hf download spc819/AortaSeg24 --repo-type model --local-dir /path/to/nnUNet_results

2. Set the nnU-Net environment variables (as usual for nnU-Net):

export nnUNet_raw=/path/to/nnUNet_raw
export nnUNet_preprocessed=/path/to/nnUNet_preprocessed
export nnUNet_results=/path/to/nnUNet_results

3. Run the two-stage inference with the code repository:

bash test_run.sh        # edit the paths / fold inside before running

The inference entry point inference.py runs the low-resolution binary model, crops the expanded ROI and then applies the full-resolution multi-class model.

Training guidance

Training is also two-stage; run the scripts from the GitHub repository on the corresponding datasets:

bash task_run_3d_lowres_825.sh     # 1) Dataset825_AortaSeg24_CTA_bin_50  (2 classes)
bash task_run_3d_fullres_824.sh    # 2) Dataset824_AortaSeg24_CTA_50      (24 classes)

Prepare the datasets in the standard nnU-Net format (imagesTr/, labelsTr/, dataset.json) and preprocess them with nnUNetv2_plan_and_preprocess before training. The hierarchical semantic loss and cbDice are provided by the customized nnU-Net fork installed in the step above.

Inference guidance

bash test_run.sh                   # two-stage inference
bash save.sh                       # optionally package the container for submission

• Input: CTA volumes in NIfTI format (*_0000.nii.gz, standard nnU-Net naming).
• Output: label maps with 0 = background and 1–23 = aortic segments/zones; the exact class definition follows the AortaSeg24 challenge description.
• The ROI scaling factor (m = 2.0) and the fold selection are set in the inference script; the checkpoints published here cover folds 0–4 and fold_all for the multi-class model, and fold 0 / fold_all for the binary model.

Citation

@misc{shi2025hierarchicalsemanticlearningmulticlass,
      title={Hierarchical Semantic Learning for Multi-Class Aorta Segmentation}, 
      author={Pengcheng Shi},
      year={2025},
      eprint={2511.14187},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2511.14187}, 
}

If you use the cbDice loss, please also cite:

@inproceedings{shi2024centerline,
  title={Centerline Boundary Dice Loss for Vascular Segmentation},
  author={Shi, Pengcheng and Hu, Jiesi and Yang, Yanwu and Gao, Zilve and Liu, Wei and Ma, Ting},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={46--56},
  year={2024},
  organization={Springer}
}

@article{isensee2021nnu,
  title={nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation},
  author={Isensee, Fabian and Jaeger, Paul F and Kohl, Simon AA and Petersen, Jens and Maier-Hein, Klaus H},
  journal={Nature methods},
  volume={18},
  number={2},
  pages={203--211},
  year={2021},
  publisher={Nature Publishing Group}
}

Acknowledgements

This solution is built on nnU-Net, fractal softmax and cbDice. We thank the AortaSeg24 organizers and the grand-challenge.org team for the dataset and the evaluation platform.