deboraJ23/AI_MRI

auto_detect_breast_mri

Malignancy classification on breast MRI, comparing an abbreviated protocol (AP) against the full dynamic protocol (FDP). The importable package is auto_detect_breast_mri; every command line entry point lives in scripts/.

"AP" and "abrv" throughout this repository means abbreviated protocol: the native scan plus the first subtraction image, instead of the full dynamic series, which is referred to as "FDP" or "full".

Getting started

With uv, which reads pyproject.toml and installs the package itself:

uv sync                                  # creates .venv from uv.lock, package included
cp config.example.yaml config.yaml       # config*.yaml is gitignored
$EDITOR config.yaml                      # fill in your paths
export AIMRI_CONFIG=$PWD/config.yaml
uv run python scripts/train/train.py resnet18_abrv -b 2 -n 1 -c 0 -d 0 -t -e

With conda or plain pip instead:

pip install -r requirements.txt          # or: conda env create -f ai_mri.yml
pip install -e . --no-deps               # makes `import auto_detect_breast_mri` work
cp config.example.yaml config.yaml
$EDITOR config.yaml
export AIMRI_CONFIG=$PWD/config.yaml

Both routes read the same two files: dependencies come from requirements.txt and the version from auto_detect_breast_mri/__init__.py, which pyproject.toml pulls in dynamically.

On a GPU cluster, install the torch build that matches the site's CUDA before uv sync (or point uv at the matching index), otherwise the default PyPI wheel decides which CUDA runtime you get.

The editable install is what lets the entry points under scripts/ find the package: Python puts the script's own directory on sys.path, never the working directory, so without it every script fails with ModuleNotFoundError: No module named 'auto_detect_breast_mri'. --no-deps keeps pip from reinstalling torch and friends over an environment that already works. export PYTHONPATH=$PWD from the repository root does the same for a single shell, without installing anything; that is what jobs/cross_validation.sbatch.example does.

Everything site specific — paths, cluster account, wandb project, CSV column names — lives in that YAML file, never in the code. Command line arguments always take precedence over it, so an invocation that passes all paths explicitly needs no config file at all:

python scripts/train/train.py resnet18_abrv <data_root> <metadata_file> <split_root> -b 32 -n 60 -c 0 -d 0 -t -e
python scripts/train/train.py resnet18_abrv -b 32 -n 60 -c 0 -d 0 -t -e     # same, paths from the config

$AIMRI_CONFIG may point anywhere, so several site configs can live side by side (for example config.hpc.yaml for the cluster and config.local.yaml for a workstation). Only config.yaml and config.yml in the repository root are found automatically.

Try it without patient data

python tools/make_synthetic_dataset.py --out synthetic_data --masks
export AIMRI_CONFIG=synthetic_data/config.yaml
python scripts/train/train.py resnet18_abrv -b 2 -n 1 -c 0 -d 0 -t -e

This writes eight synthetic examinations in both supported folder layouts, a matching metadata export and split files, plus a ready to use config. The volumes are random noise with a label-dependent blob, so a model can actually overfit them. CI runs the tests against this same dataset.

Data contract

Image folders

Two layouts are supported and may be mixed under one data_root (data_root may also be a comma separated list of roots):

Bilateral — one folder per examination, side encoded in the file name:

<data_root>/<examination_id>/
    L-Dyn_0.nii.gz   L-Sub_1.nii.gz   L-Sub_2.nii.gz  ...  L-T2.nii.gz
    R-Dyn_0.nii.gz   R-Sub_1.nii.gz   ...
    L-mask_breast_nn.nii.gz   R-mask_breast_nn.nii.gz     # only for segmentation/cropping

Unilateral — one folder per breast, side in the folder name, Pre/Post_X file names:

<data_root>/<examination_id>_left/
    Pre.nii.gz   Post_1.nii.gz   Sub_1.nii.gz  ...  T2.nii.gz
    mask_breast_nn.nii.gz                                  # only for segmentation/cropping
<data_root>/<examination_id>_right/
    ...

Pre is mapped onto Dyn_0 and Post_X onto Dyn_X, so both layouts end up with the same canonical sequence names. A folder name ending in 01 is also matched against the ID without those two digits, which some exports append.

Sequence and protocol names

Canonical sequence names: Dyn_0 (native), Dyn_1 … Dyn_4 (post contrast), Sub_1 … Sub_4 (subtraction, i.e. Dyn_X − Dyn_0) and T2.

protocol	expands to
abbreviated	Dyn_0, Sub_1
full	Dyn_0, Sub_1, Sub_2, Sub_3, Sub_4, T2
sub	Sub_1
a list of names	exactly those, in the given order

Dyn_0 must come first when present; the model's input channel count equals the number of sequences. Model names encode the protocol: *_abrv → abbreviated, *_full → full, *_sub → sub.

Metadata export

One row per examination, given as metadata_file (.csv or .xlsx), with per-side label columns:

exam_id	exam_id	pat_id	Malign_Left	Malign_Right
10001	10001	90001	0	1

Split files

One row per breast, under split_root, one folder per outer fold. <k> is the outer fold, <s> the inner (validation) fold; the test set has no inner fold:

<split_root>/fold<k>/stratified_training_set-f<k>-<s>.csv
<split_root>/fold<k>/stratified_evaluation_set-f<k>-<s>.csv
<split_root>/fold<k>/stratified_test_set-f<k>.csv
<split_root>/fold<k>/stratified_training_set-f<k>-<s>_frac0.25.csv   # training fractions

exam_id	pat_id	Side	Malign
10001	90001	Left	0
10001	90001	Right	1

The first column is read positionally as the examination ID and must match the image folder names. Splits are grouped by patient ID, so no patient appears in two folds.

Column names are configurable

The column names above are the defaults of the original UKA export. Point the code at a differently named export with the columns: block of your config file — no code change needed:

columns:
  examination_id: StudyID            # in the split files
  metadata_examination_id: StudyID   # in the metadata export
  patient_id: PatientID
  side: Side
  criterion: Cancer
  criterion_left: Cancer_Left
  criterion_right: Cancer_Right
  side_left: L
  side_right: R

Pipeline

1. Preprocess — crop with the trained UNet and split each volume into left/right halves: python -m auto_detect_breast_mri.preprocessing.split_breasts --unet-path <unet.pth> (raw DICOM first goes through auto_detect_breast_mri.preprocessing.anonymize_dicom).
2. Split — patient-grouped, label-stratified folds plus nested training fractions: python scripts/data_utils/make_splits.py and python scripts/data_utils/make_split_fractions.py (verify the nesting with python tools/check_split_fractions.py).
3. Train — one model per fold: python scripts/train/train.py <model_name> -c <fold> -d <subfold> -t -e On SLURM: python jobs/run_trainings.py --nets resnet18_abrv --folds 0 1 2 3 4 --dry-run
4. Compare protocols — abbreviated against full, with DeLong p-value and GradCAM heatmaps: python scripts/validation/compare_protocols.py <model_name> ... -o <checkpoint pattern> # checkpoint pattern will be formatted with model_key and fold, so use something like /{}_fold={}.pth depending on where you have model_key and fold in your filename
5. Non-inferiority — out-of-fold predictions, then the clustered bootstrap: python scripts/validation/predict_oof.py resnet18 <data_root> <metadata.csv> <split_files_root> -c 0 -m '<path_to_pretrained_models_root>/{model_key}_fold={fold}_subfold=0_frac=0.05_FINAL.pth' then python scripts/validation/analyze_noninferiority.py <output_root_from_predict_oof> <path to metadata.csv> --bca --skip_secondary
6. Per indication subgroups — exploratory, on the same out-of-fold predictions: python scripts/validation/analyze_subgroups.py <output_root_from_predict_oof> <path to metadata.csv> Reports per indication the AUC of both protocols and the paired AP minus FDP difference with a patient-clustered interval. Give the codes readable names with an indication_labels: block in the site config. Exploratory only: the subgroups are small, nothing is adjusted for multiplicity across them, and no margin is tested.

Layout

auto_detect_breast_mri/
├── config.py            site configuration ($AIMRI_CONFIG), path and column resolution
├── data/                datasets, transforms, split generation, dataloaders, metadata/NIfTI IO
├── models/              ResNets, the cropping UNet, checkpoint IO
├── training/            training and evaluation loops, LR schedules, shared CLI
├── evaluation/          metrics, GradCAM, clustered bootstrap, dataset statistics
└── preprocessing/       DICOM anonymisation, cropping and left/right splitting
scripts/                 every command line entry point, verb-named
├── data_utils/          splits, metadata fixes, label checks, dataloader/transform inspection
├── train/               training entry points (classifier, pretrained baseline, cropping UNet)
└── validation/          inference, evaluation, non-inferiority analysis, plots and heatmaps
jobs/                    SLURM templates and launchers (copy the .example files)
tools/                   synthetic dataset generator, repository checks
tests/                   unit tests only

Keeping patient data out of the repository

• config*.yaml (except the example), jobs/site.env and *.sbatch are gitignored: they are the only place where absolute paths, cluster accounts and personal directories belong.
• .gitignore also excludes NIfTI/DICOM/CSV/XLSX files, checkpoints, plots and wandb/.
• python tools/check_no_private_paths.py fails on absolute paths, cluster accounts, personal names and bare 7–10 digit numbers (which could be examination IDs). Install it as a pre-commit hook with pip install pre-commit && pre-commit install; CI runs it on every push.
• Weights & Biases: run configs and names are visible to everyone with access to the project, and heatmap captions contain patient keys. Set wandb_mode: offline (or disabled) in your config when that is not acceptable.