Home

Awesome

Disentangled Image Colorization via Global Anchors

Paper | Project Page | Online Demo

:blush: This is the official PyTorch implementation of our colorization work DISCO, published in SIGGRAPH Asia 2022.

Conceptually, our disentangled colorization model consists of two components: (i) anchor color representation that predicts the global color anchors (i.e. location and probabilistic colors) as the color representation of the whole image; (ii) anchor-guided color generation that synthesizes the per-pixel colors deterministically by referring the specified anchor colors.

<div align="center"> <img src="asserts/network.png" width="95%"> </div>

:briefcase: Dependencies and Installation

1. Clone this repository

git clone https://github.com/MenghanXia/DisentangledColorization
cd DisentangledColorization

2. Configure the environment

<details> <summary><strong>Option I</strong>: step by step.</summary>
# create a new anaconda env
conda create -n DISCO python=3.8
source activate DISCO

# install pytortch
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch

# install other packages
pip install requirements.txt
</details> <details> <summary><strong>Option II</strong>: from script file.</summary>
# create a anaconda env with required packages installed
conda env create -f environment.yaml
</details>

:gift: Checkpoints

NameURLScriptModel DescriptionFID (COCO)
DISCOmodeltrain.shrecommended colorization model (kept updating)TBD
DISCO-c0.2modeltrain.shcolorization model with less aggressive color saturation10.47
DISCO-randmodeltrain.shcolorization model with higher robustness to anchor sites10.25

:zap: Quick Inference

Also, you can specify your own directories by runing the command below:

python ./main/colorizer/inference.py --checkpt [checkpoint path] --data [input dir] \
	--name [save name]

You are recommended to use the absolute paths as arguments, otherwise please note that running inference.py will redirect the current dir to ./main/colorizer. Note that, changing the random seed --seed may result in different colorization result because the clustering-based anchor location involves randomness.

:computer: Training

:triangular_ruler: Evaluation

We provide the python implementation of the colorization evaluation metrics HERE, and the corresponding running scripts are attached.

# fidelity metrics: PSNR, SSIM, LPIPS
sh run_fidelity.sh

# perceptual quality: FID, IS, colorfulness
sh run_perception.sh

Acknowledgement

Part of our codes are taken from from SpixelFCN, iDeepColor, and DETR. Thanks for their awesome works.

Citation

If any part of our paper and code is helpful to your work, please generously cite with:

@article{XiaHWW22,
	author   = {Menghan Xia and Wenbo Hu and Tien-Tsin Wong and Jue Wang},
	title    = {Disentangled Image Colorization via Global Anchors},
	journal  = {ACM Transactions on Graphics (TOG)},
	volume   = {41},
	number   = {6},
	pages    = {204:1--204:13},
	year = {2022}
}