rreezN / adlcv_anomaly

Anomaly detection with diffusion models

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

To Do list before Exam Deadline - 12th of May 2024

Code implementation

  • Transpose images to have the brain be vertically oriented as in the paper
  • Normalize images in the dataloader with torch transforms - between -1 and 1

To run on HPC

  • Train the correct model with a cosine scheduler for the variance

Evaluations

  • Classifier with accuracy and confusion matrix
  • Threshold the anomaly images (OTSU???) and compared to GT segmentations (Segmentation scores or IoU or DICE). Can also be measured in accuracy, e.g. by saying if the IoU is under some threshold we count it as detected and otherwise NOPE.
  • Script that can calculate average IoU or DICE score for all samples (using sampleloader.py)
  • FID scores for all samples
  • Assess the model visually with different L values - e.g. equidistantly run 10 different L values
  • Assess the model visually with different s values - e.g. equidistantly run 10 different s values

Poster

  • Less text, preferebly bullet points
  • Illustration of model architectures
  • In Methods - have the training algorithm or showcase what happens during the training of the DDPM/DDIM ()
  • In Methods - remove some of the main text
  • Evaluate choice of font size
  • Showing plots brain samples with s=0
  • Plot comparing L-values visually for one brain
  • Plot comparing S-values visually for one brain
  • Investigate COO of the data. African data might be different. Difference between Brats21 and Brats20.

Diffusion Models for Medical Anomaly Detection

We provide the Pytorch implementation of our MICCAI 2022 submission "Diffusion Models for Medical Anomaly Detection" (paper 704).

The implementation of Denoising Diffusion Probabilistic Models presented in the paper is based on openai/guided-diffusion.

Data

We evaluated our method on the BRATS2020 dataset, and on the CheXpert dataset. A mini-example how the data needs to be stored can be found in the folder data. To train or evaluate on the desired dataset, set --dataset brats or --dataset chexpert respectively.

Usage

We set the flags as follows:

MODEL_FLAGS="--image_size 256 --num_channels 128 --class_cond True --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16"
DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False"
TRAIN_FLAGS="--lr 1e-4 --batch_size 10"
CLASSIFIER_FLAGS="--image_size 256 --classifier_attention_resolutions 32,16,8 --classifier_depth 4 --classifier_width 32 --classifier_pool attention --classifier_resblock_updown True --classifier_use_scale_shift_norm True"
SAMPLE_FLAGS="--batch_size 1 --num_samples 1 --timestep_respacing ddim1000 --use_ddim True"

To train the classification model, run

python scripts/classifier_train.py --data_dir path_to_traindata --dataset brats_or_chexpert $TRAIN_FLAGS $CLASSIFIER_FLAGS

Example:
python scripts/classifier_train.py --data_dir data/brats/processed/training/training  --dataset brats --lr 1e-4 --batch_size 10 --image_size 256 --classifier_attention_resolutions 32,16,8 --classifier_depth 4 --classifier_width 32 --classifier_pool attention --classifier_resblock_updown True --classifier_use_scale_shift_norm True   
                                                                                                                                      

To train the diffusion model, run

python scripts/image_train.py --data_dir --data_dir path_to_traindata --datasaet brats_or_chexpert  $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS

Example:
python scripts/image_train.py --data_dir  data/brats/processed/training/training --dataset brats --image_size 256 --num_channels 128 --class_cond True --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --lr 1e-4 --batch_size 10

The model will be saved in the results folder.

For image-to-image translation to a healthy subject on the test set, run

python scripts/classifier_sample_known.py  --data_dir path_to_testdata  --model_path ./results/model.pt --classifier_path ./results/classifier.pt --dataset brats_or_chexpert --classifier_scale 100 --noise_level 500 $MODEL_FLAGS $DIFFUSION_FLAGS $CLASSIFIER_FLAGS  $SAMPLE_FLAGS 

A visualization of the sampling process is done using Visdom.

Comparing Methods

FixedPoint-GAN

We follow the implementation given in this repo. We choose λcls=1, λgpidrec=10, and train our model for 150 epochs. The batch size is set to 10, and the learning rate to 10-4.

VAE

We follow the implementation given in this repo and train the model for 500 epochs. The batch size is set to 10, and the learning rate to 10-4.

DDPM

For sampling using the DDPM approach, run

python scripts/classifier_sample_known.py  --data_dir path_to_testdata  --model_path ./results/model.pt --classifier_path ./results/classifier.pt  --dataset brats_or_chexpert --classifier_scale 100 --noise_level 500 $MODEL_FLAGS $DIFFUSION_FLAGS $CLASSIFIER_FLAGS 

About

Anomaly detection with diffusion models

License:MIT License


Languages

Language:Jupyter Notebook 92.6%Language:Python 7.4%