Here, we provide a high-quality pytorch implementation of GLCIC introduced by Iizuka et. al.
We tested our scripts in the following environment.
- Python: 3.5, 3.6
- torch: 0.4.1, 1.0.0
- torchvision: 0.2.1
- tqdm: 4.24.0
- Pillow: 5.2.0
- numpy: 1.14.5 (important)
- pyamg: 3.3.2 (important)
- scipy: 1.1.0
- GPU: Geforce GTX 1080Ti (12GB RAM) X 1
If you would like to run our training script, we recommend you to use a more than middle-range GPU such as GTX 1070 or GTX 1080(Ti).
Pretrained model: download
Training config file: download
The pretrained model was trained with CelebA dataset,
and the training config file saves training settings in json format.
Left: input image
Center: input image of GLCIC
Right: output image of GLCIC
# in {path_to_this_repo}/GLCIC-pytorch/,
$ python predict.py model_cn_step400000 config.json images/test_1.jpg out.jpg
Here, we introduce how to train a model using CelebA dataset.
First, download the dataset (i.e., img_align_celeba.zip) from this official link.
Second, run the following commands.
$ unzip img_align_celeba.zip
$ mv img_align_celeba/ {path_to_this_repo}/GLCIC-pytorch/datasets/
$ cd {path_to_this_repo}/GLCIC-pytorch/datasets/
$ python make_dataset.py img_align_celeba/
Originally, all the images are stored in img_align_celeba/
,
and the last command splits the dataset into two subsets; training dataset (80%) and test dataset (20%). All the training images are stored in img_align_celeba/train/
, while
the other images are in img_align_celeba/test/
.
Just run the following command.
# in {path_to_this_repo}/GLCIC-pytorch/,
$ python train.py datasets/img_align_celeba results/demo/
Training results (trained models and inference results at each snapshot period) are to be saved in results/demo/
.
The training procedure consists the following three phases.
- In phase 1, only Completion Network (i.e., generator) is trained.
- In phase 2, only Context Discriminator (i.e., discriminator) is trained, while Completion Network is frozen.
- In phase 3, Both of the Completion Network and the Context Discriminator are jointly trained.
Under default settings, the numbers of training steps during phase 1, phase 2, and phase 3 are 90,000, 10,000, and 400,000. Each snapshot period is set to 18,000, 2,000, and 80,000, respectively. Bach size is set to 16, while input image size is 160 x 160 (all the input images are rescalled so that the minumum side is 160, then randomly cropped to 160 x 160 images).
Basically, hyper parameters and the model architecture is exactly the same as described in the paper, but we changed the batch size from 96 to 16 due to lack of GPU memory.
To train a model with your own dataset, first you have to make a dataset directory in the following format.
dataset/ # the directory name can be anything.
|____train/
| |____XXXX.jpg # png images are also OK.
| |____OOOO.jpg
|____test/
|____oooo.jpg
|____xxxx.jpg
Images in dataset/train/
are used for training models, while
images in dataset/test/
are used to perform test inpainting at each
snapshot period.
# in {path_to_this_repo}/GLCIC-pytorch/,
$ mv dataset/ datasets/
$ python train.py datasets/dataset/ results/result/ [--cn_input_size] [--ld_input_size] [--steps_1] [--steps_2] [--steps_3] [--snaperiod_1] [--snaperiod_2] [--snaperiod_3] [--bsize]
Training results for each training phase (trained models and test completion results at each snapshot period) are to be stored in results/result/
.
Arguments
<dataset>
(required): Path to the dataset directory.<result>
(required): Path to the result directory.[--cn_input_size]
: Input size of Completion Network (default: 160). All the input images are rescalled so that the length of the minimum side = cn_input_size, then randomly cropped to cn_input_size x cn_input_size images.[--ld_input_size]
: Input size of Local Discriminator (default: 96).[--init_model_cn]
: Path to the pre-trained model of Completion Network. This is used as the initial parameters (default: None).[--init_model_cd]
: Path to the pre-trained model of Context Discriminator (default: None).[--steps_1]
: Training iterations in phase 1 (default: 90,000).[--steps_2]
: Training iterations in phase 2 (default: 10,000).[--steps_3]
: Training iterations in the last phase (default: 400,000).[--snaperiod_1]
: Snapshot period in phase 1 (default: 18,000).[--snaperiod_2]
: Snapshot period in phase 2 (default: 2,000).[--snaperiod_3]
: Snapshot period in the last phase (default: 80,000).[--max_holes]
: The max number of holes (default: 1).[--hole_min_w]
: The minimum width of a hole (default: 48).[--hole_max_w]
: The max width of a hole (default: 48).[--hole_min_h]
: The minimum height of a hole (default: 96).[--hole_max_h]
: The max height of a hole (default: 96).[--bsize]
: Batch size (default: 16).[--optimizer]
: 'adadelta' or 'adam' (default: 'adadelta').[--num_gpus]
: 1 or 2. If this parameter is set to 1, both Completion Network and Context Discriminator are run on a single GPU. On the other hand, if it is set to 2, Completion Network and Context Discriminator are run on two different gpus separately (default: 1).
Example: If you'd like to train a model with batch size 24, and the other parameters are default values, run the following command.
# in {path_to_this_repo}/GLCIC-pytorch/,
$ python train.py datasets/dataset results/result --bsize 24
Suppose you've finished train a model and the result directory is {path_to_this_repo}/GLCIC-pytorch/results/result
, run the following command.
# in {path_to_this_repo}/GLCIC-pytorch/,
$ python predict.py results/result/phase_3/model_cn_step{step_number} results/result/config.json <input_img> <output_img> [--max_holes] [--img_size] [--hole_min_w] [--hole_max_w] [--hole_min_h] [--hole_max_h]
Arguments
<input_img>
(required): Path to an input image.<output_img>
(required): Path to an output image.[--img_size]
: Input size of Completion Network (default: 160). The input image are rescalled so that the length of the minimum side = img_size, then randomly cropped to a img_size x img_size image.[--max_holes]
: The max number of holes (default: 1).[--hole_min_w]
: The minimum width of a hole (default: 48).[--hole_max_w]
: The max width of a hole (default: 48).[--hole_min_h]
: The minimum height of a hole (default: 96).[--hole_max_h]
: The max height of a hole (default: 96).
Example: If you'd like to make a inference with a input image {path_to_this_repo}/GLCIC-pytorch/input.jpg
and create an output image {path_to_this_repo}/GLCIC-pytorch/output.jpg
, run the following command.
# in {path_to_this_repo}/GLCIC-pytorch/,
$ python predict.py results/result/phase_3/model_cn_step{step_number} results/result/config.json input.jpg output.jpg