Work-in-progress Pytorch repository for the paper “Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images”. An official version will be released soon.

To pre-process BraTS data, after placing the nifti files in the folder data/brats2021, you can run:

python scripts/brats-preprocess.py 

To train a diffusion model on healthy samples:

python scripts/image_train.py 

To train an autoencoder to reconstruct the images:

python scripts/train_autoencoder.py -filepath_to_arguments_json_config scripts/baseline.json --experiment_name autoencoder

To train a baseline classifier to classify the images:

python scripts/train_baseline_classifier.py -filepath_to_arguments_json_config scripts/baseline.json --experiment_name baseline_classifier

To generate the saliency maps with ACAT (Adversarial Counterfactual Attention), employing the classifier and autoencoder trained in the previous steps:

python scripts/saliency_maps.py -filepath_to_arguments_json_config scripts/baseline.json --experiment_name create_saliency

To generate healthy samples with Dif-fuse, you can run the following code, making sure to set the correct path to the model that you wish to use:

python scripts/image_sample_dif-fuse.py --model_path results/test/model054000.pt