refactor and release code

JiahuiYu · Feb 28, 2018 · 8747474 · 8747474
1 parent 9dc43c5
commit 8747474
Show file tree

Hide file tree

Showing 32 changed files with 1,087 additions and 9 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,3 @@
+model_logs
+neuralgym_logs
+data
diff --git a/README.md b/README.md
@@ -1,15 +1,13 @@
 # Generative Image Inpainting with Contextual Attention
 
-[Paper](http://jhyu.me/resources/publications/yu2018-generative-inpainting-paper.pdf) | [ArXiv](https://arxiv.org/abs/1801.07892) | [Project](http://jhyu.me/posts/2018/01/20/generative-inpainting.html) | [Demo](http://jhyu.me/posts/2018/01/20/generative-inpainting.html#post)
+[CVPR 2018 Paper](http://jhyu.me/resources/publications/yu2018-generative-inpainting-paper.pdf) | [ArXiv](https://arxiv.org/abs/1801.07892) | [Project](http://jhyu.me/posts/2018/01/20/generative-inpainting.html) | [Demo](http://jhyu.me/posts/2018/01/20/generative-inpainting.html#post)
 
 <img src="https://user-images.githubusercontent.com/22609465/35317673-845730e4-009d-11e8-920e-62ea0a25f776.png" width="425"/> <img src="https://user-images.githubusercontent.com/22609465/35317674-846418ea-009d-11e8-90c7-652e32cef798.png" width="425"/>
 <img src="https://user-images.githubusercontent.com/22609465/35317678-848aa3fc-009d-11e8-84a5-01be01a31fc6.png" width="210"/> <img src="https://user-images.githubusercontent.com/22609465/35317679-8496ab84-009d-11e8-945c-e1f957b04288.png" width="210"/>
 <img src="https://user-images.githubusercontent.com/22609465/35347783-c5e948fe-00fb-11e8-819c-8212d4edcfd3.png" width="210"/> <img src="https://user-images.githubusercontent.com/22609465/35347784-c5f4242c-00fb-11e8-8e46-5ad224e15096.png" width="210"/>
 
 Example inpainting results of our method on images of natural scene (Places2), face (CelebA) and object (ImageNet). Missing regions are shown in white. In each pair, the left is input image and right is the direct output of our trained generative neural networks without any post-processing.
 
-**Training/testing code and models will be released soon. Please stay tuned.**
-
 ## Run train/test
 
 * Install [tensorflow](https://www.tensorflow.org/install/) and tensorflow toolkit [neuralgym](https://github.com/JiahuiYu/neuralgym) (`pip install git+https://github.com/JiahuiYu/neuralgym`).  
@@ -20,23 +18,23 @@ Example inpainting results of our method on images of natural scene (Places2), f
 * Resume training:
   * Modify `MODEL_RESTORE` flag in [inpaint.yml](/inpaint.yml). E.g., `MODEL_RESTORE: '20180115220926508503_places2_model'`
   * Run `python train.py`.
-* Run testing: `python test.py --image examples/input.png --mask examples/mask.png --checkpoint model_logs/your_model_dir`.
+* Run testing: `python test.py --image examples/input.png --mask examples/mask.png --output examples/output.png --checkpoint model_logs/your_model_dir`.
 
 ## Pretrained models
 
-[Places2]() | [CelebA]() | [CelebA-HQ]() | [ImageNet]()
+[Places2](https://drive.google.com/open?id=1M3AFy7x9DqXaI-fINSynW7FJSXYROfv-) | [CelebA](https://drive.google.com/open?id=1sP8ViF3mxUMN--xpKqonEeW9d8S8pJEo) | [ImageNet](https://drive.google.com/open?id=136APWSdPRAF7-XoS8sMBTLV-X3f-ogE0)
 
 Download the model dirs and put it under `model_logs/`. Run testing or resume training as described above. All models are trained with images of resolution 256x256 and largest hole size 128x128, above which the results may be deteriorated. We provide several example test cases. Please run:
 
 ```bash
 # Places2 512x680 input
-python test.py --image examples/input.png --mask examples/mask.png --checkpoint model_logs/your_model_dir
+python test.py --image examples/places2/wooden_input.png --mask examples/places2/wooden_mask.png --output examples/output.png --checkpoint_dir model_logs/release_places2_256
 # CelebA 256x256 input
-python test.py --image examples/input.png --mask examples/mask.png --checkpoint model_logs/your_model_dir
+python test.py --image examples/celeba/celebahr_patches_164036_input.png --mask examples/center_mask_256.png --output examples/output.png --checkpoint_dir model_logs/release_celeba_256/
 # CelebA-HQ 256x256 input
-python test.py --image examples/input.png --mask examples/mask.png --checkpoint model_logs/your_model_dir
+# Please visit CelebA-HQ demo at: jhyu.me/demo
 # ImageNet 256x256 input
-python test.py --image examples/input.png --mask examples/mask.png --checkpoint model_logs/your_model_dir
+python test.py --image examples/imagenet/imagenet_patches_ILSVRC2012_val_00000827_input.png --mask examples/center_mask_256.png --output examples/output.png --checkpoint_dir model_logs/release_imagenet_256
 ```
 
 ## TensorBoard

diff --git a/examples/celeba/celebahr_patches_163050_input.png b/examples/celeba/celebahr_patches_163050_input.png
diff --git a/examples/celeba/celebahr_patches_164036_input.png b/examples/celeba/celebahr_patches_164036_input.png
diff --git a/examples/celeba/celebahr_patches_164787_input.png b/examples/celeba/celebahr_patches_164787_input.png
diff --git a/examples/celeba/celebahr_patches_165118_input.png b/examples/celeba/celebahr_patches_165118_input.png
diff --git a/examples/celeba/celebahr_patches_165230_input.png b/examples/celeba/celebahr_patches_165230_input.png
diff --git a/examples/center_mask_256.png b/examples/center_mask_256.png
diff --git a/examples/imagenet/imagenet_patches_ILSVRC2012_val_00000827_input.png b/examples/imagenet/imagenet_patches_ILSVRC2012_val_00000827_input.png
diff --git a/examples/imagenet/imagenet_patches_ILSVRC2012_val_00008210_input.png b/examples/imagenet/imagenet_patches_ILSVRC2012_val_00008210_input.png
diff --git a/examples/imagenet/imagenet_patches_ILSVRC2012_val_00022355_input.png b/examples/imagenet/imagenet_patches_ILSVRC2012_val_00022355_input.png
diff --git a/examples/imagenet/imagenet_patches_ILSVRC2012_val_00025892_input.png b/examples/imagenet/imagenet_patches_ILSVRC2012_val_00025892_input.png
diff --git a/examples/imagenet/imagenet_patches_ILSVRC2012_val_00045643_input.png b/examples/imagenet/imagenet_patches_ILSVRC2012_val_00045643_input.png
diff --git a/examples/output.png b/examples/output.png
diff --git a/examples/places2/building_input.png b/examples/places2/building_input.png
diff --git a/examples/places2/building_mask.png b/examples/places2/building_mask.png
diff --git a/examples/places2/canyon_input.png b/examples/places2/canyon_input.png
diff --git a/examples/places2/canyon_mask.png b/examples/places2/canyon_mask.png
diff --git a/examples/places2/grass_input.png b/examples/places2/grass_input.png
diff --git a/examples/places2/grass_mask.png b/examples/places2/grass_mask.png
diff --git a/examples/places2/sunset_input.png b/examples/places2/sunset_input.png
diff --git a/examples/places2/sunset_mask.png b/examples/places2/sunset_mask.png
diff --git a/examples/places2/wooden_input.png b/examples/places2/wooden_input.png
diff --git a/examples/places2/wooden_mask.png b/examples/places2/wooden_mask.png
diff --git a/examples/style_transfer/bike.jpg b/examples/style_transfer/bike.jpg
diff --git a/examples/style_transfer/bike_style_out.png b/examples/style_transfer/bike_style_out.png
diff --git a/examples/style_transfer/bnw_butterfly.png b/examples/style_transfer/bnw_butterfly.png
diff --git a/inpaint.yml b/inpaint.yml
@@ -0,0 +1,79 @@
+# parameters
+DATASET: 'celebahq'  # 'tmnist', 'dtd', 'places2', 'celeba', 'imagenet', 'cityscapes'
+RANDOM_CROP: False
+VAL: False
+LOG_DIR: full_model_celeba_hq_256
+MODEL_RESTORE: ''  # '20180115220926508503_jyugpu0_places2_NORMAL_wgan_gp_full_model'
+
+GAN: 'wgan_gp'  # 'dcgan', 'lsgan', 'wgan_gp', 'one_wgan_gp'
+PRETRAIN_COARSE_NETWORK: False
+GAN_LOSS_ALPHA: 0.001  # dcgan: 0.0008, wgan: 0.0005, onegan: 0.001
+WGAN_GP_LAMBDA: 10
+COARSE_L1_ALPHA: 1.2
+L1_LOSS_ALPHA: 1.2
+AE_LOSS_ALPHA: 1.2
+GAN_WITH_MASK: False
+DISCOUNTED_MASK: True
+RANDOM_SEED: False
+PADDING: 'SAME'
+
+# training
+NUM_GPUS: 1
+GPU_ID: -1  # -1 indicate select any available one, otherwise select gpu ID, e.g. [0,1,3]
+TRAIN_SPE: 10000
+MAX_ITERS: 1000000
+VIZ_MAX_OUT: 10
+GRADS_SUMMARY: False
+GRADIENT_CLIP: False
+GRADIENT_CLIP_VALUE: 0.1
+VAL_PSTEPS: 1000
+
+# data
+DATA_FLIST:
+  # https://github.com/JiahuiYu/progressive_growing_of_gans_tf
+  celebahq: [
+    'data/celeba_hq/train_shuffled.flist',
+    'data/celeba_hq/train_shuffled.flist'
+  ]
+  # http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html, please to use RANDOM_CROP: True
+  celeba: [
+    'data/celeba/train_shuffled.flist',
+    'data/celeba/train_shuffled.flist'
+  ]
+  # http://places2.csail.mit.edu/, please download the high-resolution dataset and use RANDOM_CROP: True
+  places2: [
+    'data/places2/train_shuffled.flist',
+    'data/places2/static_view.flist'
+  ]
+  # http://www.image-net.org/, please use RANDOM_CROP: True
+  imagenet: [
+    'data/imagenet/train_shuffled.flist',
+    'data/imagenet/static_view.flist',
+  ]
+
+STATIC_VIEW_SIZE: 30
+IMG_SHAPES: [256, 256, 3]
+HEIGHT: 128
+WIDTH: 128
+MAX_DELTA_HEIGHT: 32
+MAX_DELTA_WIDTH: 32
+BATCH_SIZE: 16
+VERTICAL_MARGIN: 0
+HORIZONTAL_MARGIN: 0
+
+# loss
+AE_LOSS: True
+L1_LOSS: True
+GLOBAL_DCGAN_LOSS_ALPHA: 1.
+GLOBAL_WGAN_LOSS_ALPHA: 1.
+
+# loss legacy
+LOAD_VGG_MODEL: False
+VGG_MODEL_FILE: data/model_zoo/vgg16.npz
+FEATURE_LOSS: False
+GRAMS_LOSS: False
+TV_LOSS: False
+TV_LOSS_ALPHA: 0.
+FEATURE_LOSS_ALPHA: 0.01
+GRAMS_LOSS_ALPHA: 50
+SPATIAL_DISCOUNTING_GAMMA: 0.9