OCR-GAN — Official PyTorch Implementation

For any inquiries, please contact Yufei Liang at [email protected]

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