Dementia Diagnosis with Multi-Modal Vision Transformers Using MRI and PET

Official Pytorch Implementation of Paper - 💎 DiaMond: Dementia Diagnosis with Multi-Modal Vision Transformers Using MRI and PET - Accepted by WACV 2025

Installation

1. Set up the environment

make env

2. Activate the environment

conda activate diamond

Note: If the environment already exists, it will not be recreated. To force update use:

make env-update

Data

We used data from Alzheimer's Disease Neuroimaging Initiative (ADNI) and Japanese Alzheimer's Disease Neuroimaging Initiative (J-ADNI).

To prepare the data:

1. Download ADNI data manually and place in data/raw/{train,val,test} directories

2. Run data processing:

make data

Data for training, validation, and testing will be processed into separate HDF5 files, using the following hierarchical format:

1. First level: A unique identifier, e.g. image ID.
2. The second level always has the following entries:
   1. A group named MRI/T1, containing the T1-weighted 3D MRI data.
   2. A group named PET/FDG, containing the 3D FDG PET data.
   3. A string attribute DX containing the diagnosis labels: CN, Dementia/AD, FTD, or MCI, if available.
   4. A scalar attribute RID with the patient ID, if available.

Usage

The package uses PyTorch. To train and test DiaMond, execute the src/train.py script. The configuration file of the command arguments is stored in config/config.yaml. The essential command line arguments are:

• --dataset_path: Path to HDF5 files containing either train, validation, or test data splits.
• --img_size: Size of the input scan.
• --test: True for model evaluation.

After specifying the config file, simply start training/evaluation by:

python src/train.py

Airflow Integration

DiaMond supports Apache Airflow for automated workflow orchestration. The following Make commands are available:

# Initialize Airflow environment
make airflow-init

# Start Airflow web server
make airflow-start

# Start Airflow scheduler
make airflow-scheduler

# Stop all Airflow processes
make airflow-stop

# Trigger the data processing DAG
make airflow-trigger

# Check status of DAGs
make airflow-status

# Set all required Airflow variables from your current configuration
make set-airflow-vars

# Run the complete DiaMond pipeline (processing + training)
make airflow-run-pipeline

# Open the Airflow web interface to monitor pipeline execution
make airflow-monitor

These commands enable you to schedule and monitor the entire DiaMond workflow from data preprocessing to model training using Airflow's powerful orchestration capabilities.

Misc

This project uses:

• Make for automation.
• Airflow for workflow orchestration.
• Ruff for code linting and formatting.
• unittest for unit testing.
• Hypothesis for property-based testing.
• PyTorch for deep learning models.
• Scikit-learn for metrics and evaluation.
• Wandb for experiment tracking.
• Nibabel for neuroimaging data handling.
• PyDICOM for DICOM file handling.
• dicom2nifti and dcm2niix for DICOM to NIFTI conversion.
• einops for tensor operations.
• tqdm for progress bar visualization.