[CVPR 2020] Instance-aware Image Colorization

[Paper] [Project Website] [Google Colab]

Image colorization is inherently an ill-posed problem with multi-modal uncertainty. Previous methods leverage the deep neural network to map input grayscale images to plausible color outputs directly. Although these learning-based methods have shown impressive performance, they usually fail on the input images that contain multiple objects. The leading cause is that existing models perform learning and colorization on the entire image. In the absence of a clear figure-ground separation, these models cannot effectively locate and learn meaningful object-level semantics. In this paper, we propose a method for achieving instance-aware colorization. Our network architecture leverages an off-the-shelf object detector to obtain cropped object images and uses an instance colorization network to extract object-level features. We use a similar network to extract the full-image features and apply a fusion module to full object-level and image-level features to predict the final colors. Both colorization networks and fusion modules are learned from a large-scale dataset. Experimental results show that our work outperforms existing methods on different quality metrics and achieves state-of-the-art performance on image colorization.

Instance-aware Image Colorization
Jheng-Wei Su, Hung-Kuo Chu, and Jia-Bin Huang
In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.

Prerequisites

• Python 3.11+
• PyTorch 2.x + torchvision 0.17+ (CPU, CUDA, or DirectML)
• Detectron2 (optional; used only for bbox generation)
• OpenCV-Python, Pillow, scikit-image, tqdm
• See requirements.txt for a lightweight, modern list.

Getting Started

1. Use Python 3.11+ (no venv required; global installs are fine).
2. Install core deps (CPU):

pip install -r requirements.txt

3. For CUDA builds, install the matching wheel first, e.g.:

pip install --extra-index-url https://download.pytorch.org/whl/cu121 torch torchvision
pip install -r requirements.txt

4. (Optional) Detectron2 for bbox prediction. Install per your CUDA/OS from https://detectron2.readthedocs.io/en/latest/tutorials/install.html then run pip install opencv-python if missing.

Pretrained Model

1. Download it from google drive.

sh scripts/download_model.sh

2. Now the pretrained models would place in checkpoints.

Instance Prediction

Please follow the command below to predict all the bounding boxes for the images in example folder.

python inference_bbox.py --test_img_dir example

All the prediction results are saved in example_bbox folder.

Colorize Images (legacy)

python test_fusion.py --name test_fusion --sample_p 1.0 --model fusion --fineSize 256 --test_img_dir example --results_img_dir results

All the colorized results would save in results folder.

• Note: all the images would convert into L channel to colorize in test_fusion.py's L51

Modern Inference (single script)

python colorize.py input.jpg --style siggraph17 --output results_modern --use-bbox

• Styles: --style siggraph17 (large) or --style eccv16 (small).
• Models auto-download if missing (checked in global torch cache first, then checkpoints/base).
• Works on CPU/CUDA/DirectML (--device cpu|cuda|directml).
• Uses fusion path when bbox npz files exist (default: <input>_bbox/<name>.npz).
• Falls back to SIGGRAPH-only colorization when boxes are absent.
• Add --run-detectron2 to auto-run Detectron2 when available; otherwise, it will skip boxes.
• Deterministic and float16 options: --deterministic, --dtype float16 (GPU only).

ONNX Export

Export the SIGGRAPH/ECCV generator to ONNX and optionally validate with ONNX Runtime (CPU/DML):

python export_onnx.py --style siggraph17 --output checkpoints/base/siggraph.onnx --dynamic --validate
python export_onnx.py --style eccv16 --output checkpoints/base/eccv16.onnx --validate

Training the Model

Please follow this tutorial to train the colorization model.

License

This work is licensed under MIT License. See LICENSE for details.

Citation

If you find our code/models useful, please consider citing our paper:

@inproceedings{Su-CVPR-2020,
  author = {Su, Jheng-Wei and Chu, Hung-Kuo and Huang, Jia-Bin},
  title = {Instance-aware Image Colorization},
  booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year = {2020}
}

Acknowledgments

Our code borrows heavily from the amazing colorization-pytorch repository.