Official PyTorch implementation of ICCV 2023 Paper "DDColor: Towards Photo-Realistic Image Colorization via Dual Decoders".
Xiaoyang Kang, Tao Yang, Wenqi Ouyang, Peiran Ren, Lingzhi Li, Xuansong Xie
DAMO Academy, Alibaba Group
🪄 DDColor can provide vivid and natural colorization for historical black and white old photos.
🎲 It can even colorize/recolor landscapes from anime games, transforming your animated scenery into a realistic real-life style! (Image source: Genshin Impact)
-
[2024-01-28] Support inferencing via Hugging Face! Thanks @Niels for the suggestion and example code and @Skwara for fixing bug.
-
[2024-01-18] Add Replicate demo and API! Thanks @Chenxi.
-
[2023-12-13] Release the DDColor-tiny pre-trained model!
-
[2023-09-07] Add the Model Zoo and release three pretrained models!
-
[2023-05-15] Code release for training and inference!
-
[2023-05-05] The online demo is available!
We provide online demos through ModelScope at and Replicate at .
Feel free to try them out!
In short: DDColor uses multi-scale visual features to optimize learnable color tokens (i.e. color queries) and achieves state-of-the-art performance on automatic image colorization.
- Python >= 3.7
- PyTorch >= 1.7
conda create -n ddcolor python=3.9
conda activate ddcolor
pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt
python3 setup.py develop # install basicsr
- Install modelscope:
pip install modelscope
- Run the following codes:
import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
img_colorization = pipeline(Tasks.image_colorization, model='damo/cv_ddcolor_image-colorization')
result = img_colorization('https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg')
cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
It will automatically download the DDColor models.
You can find the model file pytorch_model.pt
in the local path ~/.cache/modelscope/hub/damo.
- Download the pretrained model file by simply running:
from modelscope.hub.snapshot_download import snapshot_download
model_dir = snapshot_download('damo/cv_ddcolor_image-colorization', cache_dir='./modelscope')
print('model assets saved to %s'%model_dir)
then the weights will be modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt
.
Or, download the model from Hugging Face.
- Run
sh scripts/inference.sh
Now we can load model via Huggingface Hub like this:
from inference.colorization_pipeline_hf import DDColorHF
ddcolor_paper_tiny = DDColorHF.from_pretrained("piddnad/ddcolor_paper_tiny")
ddcolor_paper = DDColorHF.from_pretrained("piddnad/ddcolor_paper")
ddcolor_modelscope = DDColorHF.from_pretrained("piddnad/ddcolor_modelscope")
ddcolor_artistic = DDColorHF.from_pretrained("piddnad/ddcolor_artistic")
Check inference/colorization_pipeline_hf.py
for the details of the inference, or directly perform model inference by simply running:
python inference/colorization_pipeline_hf.py --model_name ddcolor_modelscope --input ./assets/test_images
# model_name: [ddcolor_paper | ddcolor_modelscope | ddcolor_artistic | ddcolor_paper_tiny]
- Install the gradio and other required libraries
!pip install gradio gradio_imageslider timm -q
- Run the demo
python gradio_app.py
We provide several different versions of pretrained models, please check out Model Zoo.
- Dataset preparation: download ImageNet dataset, or prepare any custom dataset of your own. Use the following script to get the dataset list file:
python data_list/get_meta_file.py
-
Download pretrained weights for ConvNeXt and InceptionV3 and put it into
pretrain
folder. -
Specify 'meta_info_file' and other options in
options/train/train_ddcolor.yml
. -
Run
sh scripts/train.sh
If our work is helpful for your research, please consider citing:
@inproceedings{kang2023ddcolor,
title={DDColor: Towards Photo-Realistic Image Colorization via Dual Decoders},
author={Kang, Xiaoyang and Yang, Tao and Ouyang, Wenqi and Ren, Peiran and Li, Lingzhi and Xie, Xuansong},
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
pages={328--338},
year={2023}
}
We thank the authors of BasicSR for 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.
Some codes are adapted from ColorFormer, BigColor, ConvNeXt, Mask2Former, and DETR. Thanks for their excellent work!