| --- |
| pipeline_tag: image-segmentation |
| license: apache-2.0 |
| --- |
| |
| # AortaSeg24: Hierarchical Semantic Learning for Multi-Class Aorta Segmentation |
|
|
| [](https://arxiv.org/abs/2511.14187) |
| [](https://huggingface.co/spc819/AortaSeg24) |
| [](https://github.com/PengchengShi1220/AortaSeg24) |
| [](https://aortaseg24.grand-challenge.org/) |
|
|
| This repository hosts the trained nnU-Net checkpoints of our solution to the |
| [AortaSeg24 Challenge](https://aortaseg24.grand-challenge.org/), which achieved |
| **🥈 2nd place in the Final Test Phase** ([official rankings](https://grand-challenge.org/forums/forum/multi-class-segmentation-of-aortic-branches-and-zones-in-cta-726/topic/challenge-final-rankings-2826/)). |
|
|
| Training and inference code (including the customized nnU-Net fork, the hierarchical |
| semantic loss and the cbDice loss) is at |
| **[github.com/PengchengShi1220/AortaSeg24](https://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](https://github.com/PengchengShi1220/fractal-softmax) and the |
| [cbDice](https://github.com/PengchengShi1220/cbDice) (centerline boundary Dice) loss. |
| See the paper for details. |
|
|
| ## Requirements |
|
|
| **Python Version:** 3.10 (same environment as the code repository) |
|
|
| ### Installation |
|
|
| ```bash |
| 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): |
| |
| ```bash |
| 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): |
|
|
| ```bash |
| 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 |
| 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 |
| 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 |
| 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](https://aortaseg24.grand-challenge.org/). |
| - 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](https://github.com/MIC-DKFZ/nnUNet), |
| [fractal softmax](https://github.com/PengchengShi1220/fractal-softmax) and |
| [cbDice](https://github.com/PengchengShi1220/cbDice). We thank the AortaSeg24 organizers |
| and the grand-challenge.org team for the dataset and the evaluation platform. |
|
|
