THOR_DDPM

Diffusion Models with Implicit Guidance for Medical Anomaly Detection

Cosmin Bercea • Benedikt Wiestler • Daniel Rueckert • Julia A. Schnabel

Official repository of the paper

Preprint

Citation

If you find our work useful, please cite our paper:

@misc{Bercea2024diffusion,
    title={Diffusion Models with Implicit Guidance for Medical Anomaly Detection},
    author={Cosmin I. Bercea and Benedikt Wiestler and Daniel Rueckert and Julia Schnabel},
    year={2024},
    month={3},
    eprint={2403.08464},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}

Abstract: Diffusion models have advanced unsupervised anomaly detection by improving the transformation of pathological images into pseudo-healthy equivalents. Nonetheless, standard approaches may compromise critical information during pathology removal, leading to restorations that do not align with unaffected regions in the original scans. Such discrepancies can inadvertently increase false positive rates and reduce specificity, complicating radiological evaluations. This paper introduces Temporal Harmonization for Optimal Restoration (THOR), which refines the de-noising process by integrating implicit guidance through temporal anomaly maps. THOR aims to preserve the integrity of healthy tissue in areas unaffected by pathology. Comparative evaluations show that THOR surpasses existing diffusion-based methods in detecting and segmenting anomalies in brain MRIs and wrist X-rays. Code: https://github.com/ci-ber/THOR_DDPM.

Setup and Run

The code is based on the deep learning framework from the Institute of Machine Learning in Biomedical Imaging: https://github.com/compai-lab/iml-dl

Framework Overview

1). Clone repository

git clone https://github.com/ci-ber/THOR_DDPM.git
cd THOR_DDPM

2). Install PyTorch

Optional create virtual env:

conda create --name thor python=3.8.0
conda activate thor

Example installation:

with cuda:

pip3 install torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html

w/o cuda:

pip3 install torch==1.9.1 torchvision==0.10.1 -f https://download.pytorch.org/whl/torch_stable.html

3). Install requirements

pip install -r pip_requirements.txt

4). Set up wandb (https://docs.wandb.ai/quickstart)

wandb login

Paste the API key from https://wandb.ai/authorize when prompted.

5). Download datasets

IXI • Atlas (Stroke) • Pediatric Wrist X-ray

Move the datasets to the target locations. You can find detailed information about the expected files and locations in the corresponding *.csv files under data/$DATASET/splits.

Alternatively you can use your own mid-axial slices of T1w brain scans with our pre-trained weights for Gaussian noise or Simplex noise or train from scratch on other anatomies and modalities.

For pediatric wrist X-rays you can use our pre-trained weights for Gaussian noise or Simplex noise.

6). Run the pipeline

Run the main script with the corresponding config like this:

python core/Main.py --config_path ./projects/thor/configs/brain/thor.yaml

Refer to the thor.yaml for the default configuration.

That's it, enjoy! :rocket:

ci-ber / THOR_DDPM

readme