Git Product home page Git Product logo

zoom-to-inpaint's Introduction

Zoom-to-Inpaint: Image Inpainting with High Frequency Details

Reference code for the paper Zoom-to-Inpaint: Image Inpainting with High Frequency Details, presented at the New Trends in Image Restoration and Enhancement (NTIRE) Workshop held in conjunction with CVPR 2022.

If you find our repo useful, please consider citing our paper:

@inproceedings{kim2022zoomtoinpaint,
      title = {Zoom-to-Inpaint: Image Inpainting with High Frequency Details}, 
      author = {Kim, Soo Ye and Aberman, Kfir and Kanazawa, Nori and Garg, Rahul and Wadhwa, Neal and Chang, Huiwen and Karnad, Nikhil and Kim, Munchurl and Liba, Orly},
      booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops},
      year = {2022}
}

Requirements

This code was implemented using Tensorflow 2 with Python 3.6 under a Linux environment. The required libraries can be viewed in requirements.txt, and can be downloaded using the following command:

pip install -r requirements.txt

The training pipelines are implemented with tf.distribute.MirroredStrategy() for distributed learning with up to 8 GPUs on a single worker. Note that the training codes also work on a single GPU or CPUs without any modification of the code.

Training

Go through these steps to follow the training scheme in our paper:

  1. Pre-training steps:
  • Pre-train the coarse network with: python main.py pretrain --network_mode=coarse
  • Pre-train the refinement network with: python main.py pretrain --network_mode=refine
  • Pre-train the super-resolution network. python main.py pretrain --flagfile=pretrain_sr.cfg
  1. Train all components jointly in a GAN framework with small masks: python main.py train --flagfile=train_small_mask.cfg
  2. Train all components jointly in a GAN framework with large masks. python main.py train --flagfile=train_large_mask.cfg

Notes

  • Pre-training (stage 1)
    • Weights will be saved in: ./pretrain/[network_mode]/ckpt
    • Logs for Tensorboard and a text log file will be saved in: ./pretrain/[network_mode]/logs
    • [network_mode]: coarse, refine, sr
  • Main training (stage 2 & 3)
    • Weights will be saved in: ./train/[mask_type]/ckpt
    • Logs for Tensorboard and a text log file will be saved in: ./train/[mask_type]/logs
    • [mask_type]: small_mask, large_mask
  • If you've followed all the training steps (same training scheme as our paper), the final weights would be the ones in: ./train/large_mask/ckpt

Testing

Directory structure

  • Download the test data from here and put it under a directory named data
  • Download the pretrained checkpoint from here and put it under a directory named ckpt
Zoom-to-Inpaint
├── ckpt
│    ├── checkpoint
│    ├── ckpt-1500.data-00000-of00001
│    └── ckpt-1500.index
└── data
     ├── div2k
     │    ├── image
     │    │    ├── 0001.png
     │    │    ├── ...
     │    │    └── 0100.png
     │    ├── mask
     │    │    ├── large
     │    │    │      └── ...
     │    │    └── small
     │    │          └── ...
     │    └── masked
     │        ├── large
     │        │      └── ...
     │        └── small
     │              └── ...
     ├── places_test
     │    └── ...
     └── places_val
         └── ...

Quick Start

  • Run: python main.py test
    • Result images will be saved in ./results.
  • To print metric values: python main.py test --eval

Flags

  • --img_dir=[path]: Directory containing images (PNG) to inpaint.
  • --mask_dir=[path]: Directory containing the corresponding inpainting masks (PNG).
  • --result_dir=[path]: Desired directory for saving inpainted results.
  • --eval: Add this flag if you wish to compute and print metric values.

Testing on provided data

python main.py test --img_dir='./data/[dataset]/image' --mask_dir='./data/[dataset]/mask/[mask_type]' --result_dir='./results'
  • [dataset]: div2k, places_val, places_test
  • [mask_type]: small, large
  • Result images will be saved in ./results.
  • Add --eval flag to evaluate on quality metrics.

Testing on your own data

  • Run python main.py test with appropriate flag values set to --img_dir, --mask_dir, and --result_dir
    • Files in --mask_dir directory should have the same file name as their corresponding images in --img_dir.
    • --eval flag for evaluating on performance metrics should only be used if full images (without holes) are given to --img_dir.

Testing with self-trained weights

  • Provide your checkpoint directory to --ckpt_dir (eg. --ckpt_dir=./train/large_mask/ckpt/ckpt-1500)

Additional notes

  • You can also set other hyperparameters in config.py by passing them as a flag or directly modifying the default values in the file.
    • Note that the location of the working directory can be changed with the --work_dir flag.

Disclaimer

This is not an officially supported Google product.

zoom-to-inpaint's People

Contributors

orlyliba avatar sooyekim avatar

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar  avatar  avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.