AnomalyDiffusion: Few-Shot Anomaly Image Generation with Diffusion Model (AAAI 2024)

Teng Hu^1#, Jiangning Zhang^2#, Ran Yi^1*, Yuzhen Du¹, Xu Chen², Liang Liu², Yabiao Wang², and Chengjie Wang^1,2.

(#Equal contribution,*Corresponding author)

¹Shanghai Jiao Tong University, ²Youtu Lab, Tencent

Project Page

Todo (Latest update: 2024/03/13)

• **Release the training code
• **Release the UNet checkpoints for testing anomaly detection accuracy
• **Release the data
• **Release checkpoints for anomalydiffusion.
• **Release the inference code

Data and checkpoints

The generated anomaly data and all the checkpoints can be downloaded from the following links. And some checkpoints should be placed at the corresponding directory. (Note that we have filtered out some data with poor generation effects. Therefore, some classes have relatively fewer samples.)

Data and Models	Download	Place at
Generated data	Google Drive	-
Checkpoints for anomaly generation model	Google Drive	logs/anomaly-checkpoints
Checkpoints for mask generation model	Google Drive	logs/mask-checkpoints
Checkpoints for anomaly localization	Google Drive	checkpoints/localization
Checkpoints for anomaly classification	Google Drice	checkpoints/classification

Overview

Anomalydiffusion is a few-shot anomaly generation model for anomaly inspection (detection, localization and classification). The overall process can be divided into the following 5 steps:

(1) Train the anomaly generation model and mask generation model;

(2) Generate anomalous masks by the mask generation model in step (1);

(3) Generate anomaly mask-image pairs by the anomaly generation model based on the generated masks in step (2);

(4) Train the anomaly detection (for both detection and localization) and classification model based on the anomalous image-mask pairs in step (3).

Prepare

(1) Prepare the environment

Ubuntu
python 3.8
cuda==11.8
gcc==7.5.0
conda env create -f environment.yaml
conda activate Anomalydiffusion

(2) Checkpoint for LDM

Download the official checkpoint of the latent diffusion model:

mkdir -p models/ldm/text2img-large/
wget -O models/ldm/text2img-large/model.ckpt https://ommer-lab.com/files/latent-diffusion/nitro/txt2img-f8-large/model.ckpt

Generating anomalous image-mask pairs

In this section, you can first train the anomaly generation model by (1). After that, you can run (2), which contains training mask generation models, generating anomalous masks and generating anomalous image-mask pairs.

(1) Train the anomaly generation model by:

CUDA_VISIBLE_DEVICES=$gpu_id python main.py --spatial_encoder_embedding --data_enhance
 --base configs/latent-diffusion/txt2img-1p4B-finetune-encoder+embedding.yaml -t 
 --actual_resume models/ldm/text2img-large/model.ckpt -n test --gpus 0, 
  --init_word anomaly  --mvtec_path=$path_to_mvtec_dataset

(2) Train the mask generation model and generate image-mask pairs by:

CUDA_VISIBLE_DEVICES=$gpu_id python run-mvtec.py --data_path=$path_to_mvtec_dataset

Test the anomaly inspection performance

After generating anomalous image-mask pairs, you can train and test the anomaly detection model (for both anomlay detection and localization) by:

python train-localization.py --generated_data_path $path_to_the_generated_data  --mvtec_path=$path_to_mvtec
python test-localization.py --generated_data_path $path_to_mvtec

you can train and test the anomaly classification model by:

python train-classification.py --mvtec_path=$path_to_mvtec --generated_data_path=$path_to_the_generated_data
python test-classification.py --mvtec_path=$path_to_mvtec --generated_data_path=$path_to_the_generated_data

Citation

If you make use of our work, please cite our paper:

@inproceedings{hu2023anomalydiffusion,
  title={AnomalyDiffusion: Few-Shot Anomaly Image Generation with Diffusion Model},
  author={Hu, Teng and Zhang, Jiangning and Yi, Ran and Du, Yuzhen and Chen, Xu and Liu, Liang and Wang, Yabiao and Wang, Chengjie},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  year={2024}
}