[ICCV 2025] Teaching AI the Anatomy Behind the Scan:
Addressing Anatomical Flaws in Medical Image Segmentation with
Learnable Prior

We introduce the Anatomy-Informed Cascaded Network (AIC-Net), a novel model for medical image segmentation that enhances the global interpretation of CT/MRI scans by incorporating a learnable anatomical prior. This prior is adaptable to a patient’s specific anatomy through differentiable spatial deformation, allowing the model to make more informed, anatomy-aware predictions.

Installation

1. Software Requirements

We run AIC-Net on a system running Debian GNU/Linux 11, with Python 3.12.4, PyTorch 2.3.1, and CUDA 12.4. For a full list of software packages and version numbers, see the Conda environment file environment.yml.

2. Hardware Requirements

To train/test AIC-Net, as well as other baseline models we require you to have at least:

• GPU memory >= 32GB on smaller backbone models (UNet, DeepLabV3+)
• GPU memory >= 80GB on smaller backbone models (UNETR, UNETR-Swin)

Also we strongly recommend using SSD storage to increase the data loading speed.

3. Installation Guide

We recommend installation of the required packages using the conda package manager, available through the Anaconda Python distribution. After installing Anaconda, install AIC-Net with the following command ：

warning

you cannot clone anonymous repo so please manually download for now.

git clone https://anonymous.4open.science/r/aic-net-A86C
cd aic-net-A86C
conda env create -f environment.yml
conda activate prior-seg 

Training

1. Dataset Download

Download TotalSegmentator dataset via followling links:

• TotalSegmentator

2. Setting Up Dataset

Totalsegmentator dataset has mutiple label files for each scan. run /data/preprocess/total_segmentor/main.py to merge the label files into one:

--create organ dataset
python data/preprocess/total_segmentor/main.py\
  --load_path [path/to/raw/dataset]\
  --save_path [project_path/data/datasets/TotalSegmentator]
  --label_type organ

--create vertebrae dataset
python data/preprocess/total_segmentor/main.py\
  --load_path [path/to/raw/dataset]\
  --save_path [project_path/data/datasets/TotalSegmentator]
  --label_type vertebrae

3. Initial prior weight download

For faster model convergence, its good to initialize the learnable prior. download the prior weights from :

• Organ
• Vertebrae

and move the prior to /data/datasets/TotalSegmentator/organ/ or /data/datasets/TotalSegmentator/vertebrae/.

4. Training and Testing

Before running a script, you have to add aic-net project as one of python paths

-- add project path as one of pythonpaths
export PYTHONPATH="${PYTHONPATH}:$HOME/aic-net-A86C"

To train and test AIC-Net, run:

python scripts/train.py\
  --model_name [PriorSeg]\
  --estimator_base_name [FasterUNet|DeepLabPlus|UNETRVit|UNETRSwin]

To train and test cascaded baseline, run:

python scripts/train.py\
  --model_name [PriorSeg]\
  --estimator_base_name [FasterUNet|DeepLabPlus|UNETRVit|UNETRSwin]\
  --no_prior

To train and test vanilla baseline, run:

python scripts/train.py\
  --model_name [PriorSeg]\
  --estimator_base_name [FasterUNet|DeepLabPlus|UNETRVit|UNETRSwin]\
  --no_prior\
  --local_only