Skip to content

Latest commit

 

History

History
104 lines (85 loc) · 4.8 KB

README.md

File metadata and controls

104 lines (85 loc) · 4.8 KB
Raw

PITI: Pretraining is All You Need for Image-to-Image Translation

Official PyTorch implementation

Pretraining is All You Need for Image-to-Image Translation
Tengfei Wang, Ting Zhang, Bo Zhang, Hao Ouyang, Dong Chen, Qifeng Chen, Fang Wen
2022

paper | project website | video | online demo

Introduction

We present a simple and universal framework that brings the power of the pretraining to various image-to-image translation tasks. You may try our online demo if interested.

Diverse samples synthesized by our approach.

Set up

Installation

git clone https://github.com/PITI-Synthesis/PITI.git
cd PITI

Environment

conda env create -f environment.yml

Quick Start

Pretrained Models

Please download our pre-trained models for both Base model and Upsample model, and put them in ./ckpt.

Model Task Dataset
Base-64x64 Mask-to-Image Trained on COCO.
Upsample-64-256 Mask-to-Image Trained on COCO.
Base-64x64 Sketch-to-Image Trained on COCO.
Upsample-64-256 Sketch-to-Image Trained on COCO.

If you fail to access to these links, you may alternatively find our pretrained models here.

Prepare Images

We put some example images in ./test_imgs, and you can quickly try them.

COCO

For COCO dataset, download the images and annotations from the COCO webpage.

For mask-to-image synthesis, we use the semantic maps in RGB format as inputs. To obtain such semantic maps, run ./preprocess/preprocess_mask.py (an example of the raw mask and the processed mask is given in preprocess/example). Note that we do not need instant masks like previous works.

For sketch-to-image synthesis, we use sketch maps extracted by HED as inputs. To obtain such sketch maps, run ./preprocess/preprocess_sketch.py.

Inference

Interactive Inference

Run the following script, and it would create an interactive GUI built by gradio. You can upload input masks or sketches and generate images.

pip install gradio
python inference.py

Batch Inference

Modify sample.sh according to the follwing instructions, and run:

bash sample.sh
Args Description
--model_path the path of ckpt for base model.
--sr_model_path the path of ckpt for upsample model.
--val_data_dir the path of a txt file that contains the paths for images.
--num_samples number of images that you want to sample.
--sample_c Strength of classifier-free guidance.
--mode The input type.

Training

Preparation

  1. Download and preprocess datasets. For COCO dataset, download the images and annotations from the COCO webpage. Run ./preprocess/preprocess_mask.py or ./preprocess/preprocess_sketch.py
  2. Download pretrained models by python preprocess/download.py .

Start Training

Taking mask-to-image synthesis as an example: (sketch-to-image is the same)

Finetune the Base Model

Modify mask_finetune_base.sh and run:

bash mask_finetune_base.sh

Finetune the Upsample Model

Modify mask_finetune_upsample.sh and run:

bash mask_finetune_upsample.sh

Citation

If you find this work useful for your research, please cite:

@article{wang2022pretraining,
 title = {Pretraining is All You Need for Image-to-Image Translation},
  author = {Wang, Tengfei and Zhang, Ting and Zhang, Bo and Ouyang, Hao and Chen, Dong and Chen, Qifeng and Wen, Fang},
  journal={arXiv:2205.12952},
  year = {2022},
}

Acknowledgement

Thanks for GLIDE for sharing their code.