myeldib / ColorFormer

ColorFormer: Image Colorization via Color Memory assisted Hybrid-attention Transformer

This is the implementation of ``ColorFormer: Image Colorization via Color Memory assisted Hybrid-attention Transformer'' (ECCV22)

Xiaozhong Ji, Boyuan Jiang, Donghao Luo, Guangpin Tao, Wenqing Chu, Zhifeng Xie, Chengjie Wang, Ying Tai

This repository is heavily based on BasicSR.

Requirements

• pytorch==1.9.0
• torchvision
• scikit-image
• einops
• timm

Please run this to install basicsr.

python3 setup.py develop

Inference

1. Download pretrained weights pretrain from google drive and put the folder under './'
2. Run

python3 inference/inference_colorformer.py --input /path/to/input --output /path/to/output --model_path pretrain/net_g_200000.pth

Train

1. Download imagenet training set from https://www.image-net.org/, then list all the image paths in a txt file.
2. Specify 'meta_info_file' in options/train/ECCV22/train_colorformer.yml
3. To collect semantic and color priors, run

python3 memory_build/inference_GLH.py --input_txt image_paths.txt
python3 memory_build/semantic_color_clustering.py -m 512 -k 64

1. For multi-gpu training, run

sh scripts/train.sh

We thank the authors of BasicSR as we train colorformer based on the awesome training pipeline.

Xintao Wang, Ke Yu, Kelvin C.K. Chan, Chao Dong and Chen Change Loy. BasicSR: Open Source Image and Video Restoration Toolbox. https://github.com/xinntao/BasicSR, 2020.