Awesome

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

News

June 29, 2024

• The anomalous data for hazelnut has been updated since the previous version can only achieve AP<85. Current Version can achieve an AP score around 96. You can download the new version in the following link to Google Drive.
• ldm/models/diffusion/ddpm.py has been updated since previous validation step has some problem. But it does not influence the training results.

Todo (Latest update: 2024/08/11)

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

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.)

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

Compute IC-LPIPS

To compute IC-LPIPS for the generated dataset, please run:

python cal_ic_lpips.py --mvtec_path=$path_to_mvtec --gen_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}
}