/OCR-GAN

Source code of "Omni-frequency Channel-selection Representations for Unsupervised Anomaly Detection"

Primary LanguagePython

(Forked)OCR-GAN — Official PyTorch Implementation

For any inquiries, please contact Yufei Liang at yufeiliang@zju.edu.cn or 186368@zju.edu.cn

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