Image Inpainting

A PyTorch Implementation of the paper, Image Inpainting for Irregular Holes Using Partial Convolutions. Architecture may not be an exact match of due to the limited description of hyperparameters and architecture details.

NOTE: See lightning branch for original pytorch-lightning version

TODO: Update README with latest pytorch ignite training and evaluating commands

Instructions

sudo conda install ignite -c pytorch-nightly
sudo conda install pip
sudo pip install neptune-client

Results

Currently results shows training done on the Goolge Landmark v2 Dataset on a single P100 GPU.

Note: Training for these results were only for 4.5 hours, while original paper trained the model on a V100 GPU for 10 days.

Example Results (so far)

4.5 Hours of Training
tv loss is very low at the moment, might need to increase the scaling factor.
highest loss is style out.

Validation Set

Currently optimizing training and hyperparameters before training for full duration longer.

Training the model

Use command below for more parameter details.

python main.py -h

Training from scratch

python main.py \
    --train_dir "/content/data/train" \
    --valid_dir "/content/data/test" \
    --mask_dir "/content/data/irregular_mask/disocclusion_img_mask" \
    --save_path "/content/gdrive/My Drive/image-inpainting" \
    --train_percent_check 1.0 \
    --val_check_interval 0.5

Continue from pretrained model

python main.py \
    --version_number 1 \
    --checkpoint_dir "/content/gdrive/My Drive/image-inpainting/default/version_0/checkpoints/_ckpt_epoch_2.ckpt" \
    --train_dir "/content/data/train" \
    --valid_dir "/content/data/test" \
    --mask_dir "/content/data/irregular_mask/disocclusion_img_mask" \
    --save_path "/content/gdrive/My Drive/image-inpainting" \
    --train_percent_check 1.0 \
    --val_check_interval 0.5

References

TODO

Create architecture
Implement training architecture
Save the model to an external location
Document code
Train for full duration
Add mixed precision support
Update dataloader to use nvidia-dali
Refactor code to use pytorch-lightning
Create Colab prediction implementation

Small Dataset Test Script

cd dataset
mkdir train
cd train
aws s3 --no-sign-request cp s3://open-images-dataset/tar/validation.tar.gz .
tar -xzf validation.tar.gz
rm validation.tar.gz
mkdir test
cd validation
mv `ls | head -100` ../test

ryanwongsa/Image-Inpainting