Defect-aware Feature Manipulation GAN

This repository is the official implementation of the following paper, which focuses on defect/anomaly generation for downstream tasks in industries:

Few-Shot Defect Image Generation via Defect-Aware Feature Manipulation
Yuxuan Duan, Yan Hong, Li Niu, Liqing Zhang
The 37th AAAI Conference on Artificial Intelligence (AAAI 2023)
https://arxiv.org/abs/2303.02389

Update Notes

• Jul. 26, 2024 ⚠️ Some issues have been marked as good first issue. You are welcome to check on them if you
  • encounter errors w.r.t. plugins, especially bias_act and upfirdn2d;
  • have insufficient VRAM;
  • fail to render good defect-free images in the first stage (e.g. completely wrong color or orientation).
• Dec. 21, 2023 Pretrained hazelnut_good.pkl and hazelnut_hole.pkl are now available here for a quick trial.

Getting Started

• ⚠️ This repository is based on the official NVIDIA implementaion of StyleGAN2-ADA. Follow the requirements of it before the steps below. (Docker is not supported. Windows might probably do with minor modification, yet have not been tested.)
• Additional python libraries: pip install scipy psutil lpips tensorboard.
• Clone the repository:

  git clone https://github.com/Ldhlwh/DFMGAN.git
  cd DFMGAN

Dataset

• Download the MVTec Anomaly Detection (MVTec AD) dataset and unzip the archive files under ./data. (If you wish to try your own datasets, organize the defect-free images, defect images and the corresponding masks in a similar way.)
• Preprocess the dataset images into zip files for easy StyleGAN loading: (e.g. object category hazelnut, defect category hole)

  # Defect-free dataset for Stage 1
  python dataset_tool.py --source ./data/hazelnut/train/good \
      --dest ./data/hazelnut_good.zip \
      --width 256 --height 256
  
  # Defect image & mask dataset for Stage 2
  python dataset_tool.py --source ./data/hazelnut/test/hole \
      --source-mask ./data/hazelnut/ground_truth/hole \
      --dest ./data/hazelnut_hole_mask.zip --width 256 --height 256

Stage 1: Pretrain

• Update on 12/21/2023 You may skip this stage by downloading the pretrained hazelnut_good.pkl provided here.
• Pretrain a StyleGAN2 model on defect-free images, using the default configuration auto: (e.g. object category hazelnut)

  python train.py --data ./data/hazelnut_good.zip \
      --outdir runs/hazelnut_good \
      --gpus 2 --kimg 3000
  
  # If training for 3000 kimgs are not enough, you may resume the pretraining by
  python train.py --data ./data/hazelnut_good.zip \
      --outdir runs/hazelnut_good \
      --gpus 2 --kimg 3000 --resume runs/hazelnut_good/path/to/the/latest/model.pkl
      
  # You may also try different values for the following settings
  # --gpus: number of GPUs to be used

• Check the qualitative & quantitative results under ./runs/hazelnut_good/*. Chose a good model for the transfer in Stage 2. You may optionally make a copy as ./pkls/hazelnut_good.pkl for easy loading.

Stage 2: Transfer

• Update on 12/21/2023 You may skip this stage by downloading the pretrained hazelnut_hole.pkl provided here.
• Transfer the pretrained model to defect images with the defect-aware feature manipulation process: (e.g. object category hazelnut, defect category hole)

  python train.py --data ./data/hazelnut_hole_mask.zip \
      --outdir runs/hazelnut_hole --resume pkls/hazelnut_good.pkl \
      --gpus 2 --kimg 400 --snap 10 --transfer res_block_match_dis
  
  # You may also try different values for the following settings
  # --gpus: number of GPUs to be used
  # --lambda-ms: weight for the mode seeking loss
  # --dmatch-scale: the number of base channel/max channel of D_match

• The above process will by default compute FID@5k, KID@5k and Clustered LPIPS@1k on-the-fly per --snap ticks (i.e. $4 \times$--snap kimgs). You may alter the metric list with --metrics.
• Check the qualitative & quantitative results under ./runs/hazelnut_hole/*. You may optionally make a copy of a good model as ./pkls/hazelnut_hole.pkl for easy loading.

Inference: Defect Image Generation

• Generate 100 random defect images: (e.g. object category hazelnut, defect category hole)

  python generate.py --network pkls/hazelnut_hole.pkl \
      --output gen_img/hazelnut_hole
      
  # You may also try different values for the following settings
  # --seeds: specify the random seeds to be used
  # --num: number of generated images (only when --seeds is unspecified)
  # --gen-good: (flag) generate defect-free images along
  # --gen-mask: (flag) generate masks along

• Or, if you just wish to have a glimpse of the performance of a model, run the following command:

  python generate.py --network runs/hazelnut_hole/path/to/a/model.pkl --cmp

  to generate triplets of defect-free image, mask and defect image like Fig. 4 under the same directory with model.pkl, named cmp<kimg>.png.

Citation

If you find DFMGAN helpful to your research, please cite our paper:

@inproceedings{Duan2023DFMGAN,
  title     = {Few-Shot Defect Image Generation via Defect-Aware Feature Manipulation},
  author    = {Yuxuan Duan and Yan Hong and Li Niu and Liqing Zhang},
  booktitle = {AAAI},
  year      = {2023}
}

Acknowledgements

• The work was supported by the National Science Foundation of China (62076162), and the Shanghai Municipal Science and Technology Major/Key Project, China (2021SHZDZX0102, 20511100300).
• This repository have used codes from StyleGAN-ADA and LPIPS.