For any inquiries, please contact Yufei Liang at yufeiliang@zju.edu.cn
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
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
- 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
Train OCR-GAN model.
python train_all.py --dataset all --isize 256 --niter 200 --model ocr_gan_aug --batchsize 32