zhangzjn / OCR-GAN

[TIP 2023] Omni-frequency Channel-selection Representations for Unsupervised Anomaly Detection

OCR-GAN — Official PyTorch Implementation

For any inquiries, please contact Yufei Liang at [e-mail] or [e-mail]

Using the Code

Requirements

This code has been developed under Python3.7, PyTorch 1.2.0 and CUDA 10.0 on Ubuntu 16.04.

# Install python3 packages
pip install -r requirements.txt

Datasets

Download MVTec, and the dataset should be copied into ./data directory, and should have the following directory & file structure:

data
├──metal_nut
│   ├── test
│   │   ├── good
│   │   │   └── 000.png
│   │   │   └── 001.png
│   │   │   ...
│   │   │   └── n.png
│   │   ├── bad
│   │   │   └── 000.png
│   │   │   └── 001.png
│   │   │   ...
│   │   │   └── m.png
│   ├── train
│   │   ├── good
│   │   │   └── 000.png
│   │   │   └── 001.png
│   │   │   ...
│   │   │   └── t.png

Inference

• Download pretraind NetG for the class "metal_nut" in MVTec dataset to the path output/ocr_gan_aug/metal_nut/train/weights/netG_best.pth.
• Download pretraind NetD for the class "metal_nut" in MVTec dataset to the path output/ocr_gan_aug/metal_nut/train/weights/netD_best.pth.

python test.py --dataset metal_nut --isize 256 --model ocr_gan_aug --load_weights

Training

Train OCR-GAN model.

python train_all.py --dataset all --isize 256 --niter 200 --model ocr_gan_aug --batchsize 32

Citation

If our work is helpful for your research, please consider citing:

@article{liang2022omni,
  title={Omni-frequency Channel-selection Representations for Unsupervised Anomaly Detection},
  author={Liang, Yufei and Zhang, Jiangning and Zhao, Shiwei and Wu, Runze and Liu, Yong and Pan, Shuwen},
  journal={arXiv preprint arXiv:2203.00259},
  year={2022}
}

Acknowledgements

We thank the great work GANomaly for providing assistance for our research.