AutoAIF¶

AutoAIF finds the arterial input function in brain DCE-MRI automatically, with no manual ROI drawing. It is a 3D U-Net in Keras/TensorFlow, trained on multi-site brain DCE-MRI cohorts, and it ships with pretrained weights — no training is required.

Given a 4D DCE series it predicts a vascular function curve and the 3D mask the curve was measured from.

Repository Paper — MRM 2025

Where this fits¶

AutoAIF covers step 3. The full DCEasy pipeline:

Citation¶

If AutoAIF contributes to published work, please cite:

Saca, L., et al. Automatic detection of arterial input function for brain DCE-MRI in multi-site cohorts. Magnetic Resonance in Medicine, 94(6), 2732–2744 (2025). PMID: 40808286

When to use it¶

Use AutoAIF when you have brain DCE-MRI and you want the AIF chosen the same way every time. See the arterial input function for the background.

Consider alternatives when:

Input images should be gzipped NIfTI (.nii.gz) in radiological orientation. Dimensions do not need to match the training data — everything is resampled to 256×256×32×32 internally and the outputs are resampled back — but the training set spanned roughly 208×256×40×50 to 320×320×14×64, and inputs far outside that range are worth resampling yourself first.

Requirements¶

• Python 3.9+
• TensorFlow 2.12+ / Keras 2.12+
• CUDA GPU recommended but not required

pip install matplotlib nibabel numba numpy pandas Pillow \
  pingouin plotly psutil scipy tensorboard tensorflow

Install¶

git clone https://github.com/petmri/AutoAIF.git && cd AutoAIF
python3 -m venv tf && source tf/bin/activate
pip install -r requirements.txt

Then fetch the pretrained weights. They are posted on github, download manually or run:

# ~470 MB
curl -L -o model_weight_huber1.h5 \
  https://github.com/petmri/AutoAIF/releases/latest/download/model_weight_huber1.h5

Worked single example¶

One subject, from a preprocessed DCE series to an AIF mask and figures:

source tf/bin/activate

python main_vif.py --mode inference \
  --input_path  /data/derivatives/DCEPrep/sub-01/sub-01_desc-hmc_DCE.nii.gz \
  --model_weight_path model_weight_huber1.h5 \
  --save_output_path /data/derivatives/AutoAIF/sub-01/ \
  --save_image 1

That writes four files into --save_output_path, named after the input:

Drop --save_image 1 and you get the two NIfTIs only, which is what you want in a batch run.

A cohort example¶

A simple shell or python loop is an easy way to process lots of images. Pass one image per call and loop in the shell:

source tf/bin/activate

for img in /data/derivatives/DCEPrep/sub-*/sub-*_desc-hmc_DCE.nii.gz; do
  sub=$(basename "$img" | cut -d_ -f1)
  python main_vif.py --mode inference \
    --input_path "$img" \
    --model_weight_path model_weight_huber1.h5 \
    --save_output_path "/data/derivatives/AutoAIF/$sub/"
done

The weights load once per invocation, which dominates runtime on short series. For a large cohort it is worth batching inside a single Python process instead.

Fine-tuning on your own data¶

The shipped weights were trained on multi-site brain DCE-MRI, and they generalize across scanners better than a single-site model would. Retraining is still worth it if your sequence, field strength or contrast protocol sits well outside that range — and the same entry point trains a model from scratch.

Training dataset layout¶

For training organize the data by site. Each site needs an images/ folder and a masks/ folder, with one mask per image under the same filename:

dataset
├── site1
│   ├── images
│   │   └── id_x.nii.gz
│   └── masks
│       └── id_x.nii.gz
├── site2
│   ├── images
│   │   └── id_x.nii.gz
│   └── masks
│       └── id_x.nii.gz
└── site3
    ├── images
    │   └── id_x.nii.gz
    └── masks
        └── id_x.nii.gz

Every folder directly under dataset is treated as a site, except those whose names begin with test or TF — which is what keeps the generated TFRecords directory from being picked up as one.

Running it¶

source tf/bin/activate

python main_vif.py --mode training \
  --dataset_path /path/to/dataset/ \
  --save_checkpoint_path /path/to/checkpoints/ \
  --epochs 200 \
  --batch_size 1

--epochs 200 and --batch_size 1 are the defaults and can be omitted. --model_name selects the architecture — best (the shipped one), attn or modified_attn. python main_vif.py -h lists everything.

The split is 80/10/10 train/validation/test per site, on a fixed seed. It is written to train_set.txt, val_set.txt and test_set.txt in the checkpoint directory on the first run and reused verbatim if those files are already there — so a resumed or repeated run keeps the same test set, and forcing a fresh split means deleting them. TFRecords are likewise only generated when they are not already present.

Training also writes log.txt, history.npy and a history.png loss curve into the checkpoint directory, and TensorBoard logs under logs/fit/.

Hyperparameter search¶

--mode hp_tuning sweeps convolution kernel sizes — the first/last layers and the body independently — and logs each trial to TensorBoard through the HParams plugin:

python main_vif.py --mode hp_tuning \
  --dataset_path /path/to/dataset/ \
  --save_checkpoint_path /path/to/tuning/

Each trial gets its own subdirectory named after its parameters, and a trial whose directory already exists is skipped — so an interrupted sweep resumes where it stopped rather than starting over. It is a full grid, so budget accordingly.