Skip to content
No description, website, or topics provided.
Branch: master
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Type Name Latest commit message Commit time
Failed to load latest commit information.

Colorization with Attention



The images are (from left to right): Ground-truth, generated by No Attention Model, generated by Weighted-loss Attention Model, generated by End-to-end Attention Model.

Our Report


Download a dataset

MIT CVCL Opencountry dataset at:

ImageNet at:

Use Pretrained Model

  1. Download our pre-trained models(Three different models) here, unzip it and put the model/ folder into the top level of the repo.

Note: The pretrained model is only trained using a subset of ImageNet (~50,000 animal images) so it does not give guarantee to performance on other type of images.

  1. Test (Follow the instruction in to set up parameters)

Train your own model

  1. Transform your training data to text_record file(The ratio of splitting training/testing can be set up in
python3 tools/
  1. Calculate the training data prior-probs(reference to tools/ Note: Currently, each bin has to appear at least once in the training images in order to avoid breaking the training procedure. If training images do not cover all the colors, you can use the existing "prior_probs.npy" in the resource folder by skipping this step.
python3 tools/
  1. Write your own train-configure file and put it in conf/ (refer to conf/train.cfg for an example)

  2. Train


python3 tools/ -c conf/train.cfg


python3 tools/ -c conf/train_gpu.cfg
  1. Test (Follow the instruction in to set up parameters)


Richard Zhang, Phillip Isola, Alexei A. Efros. Colorful Image Colorization, ECCV2016.

Saumya Jetley, Nicholas A. Lord, Namhoon Lee and Philip H. S. Torr, learn to pay attention, ICLR 2018


The code in this repo is built on top of the work at:

Tree Structure

├── conf (configure files)

├── data (dataset)

├── (Run quantitative measurement of errors of the prediction)

├── evaluation

│   ├── (evaluate colorization quantitatively)

├── (handling data processing for colorization)

├── (demo for testing correctness)

├── models (store trained models)

├── (no_attention/weighted-loss attention model)

├── (end-to-end attention model)

├── (densenet no_attention model to be fixed)

├── (customized tf layers)

├── resources

│   ├── prior_probs.npy (empiricial probability used for class rebalancing extracted from imagenet)

│   └── pts_in_hull.npy (values of each bin on ab color space)

├── (teacher CNN vgg-net for classification)

├── (graph and session for training)

├── (Can be used to extract attention map)

├── tools

│   ├── (create training list for colorization)

│   ├── (create empiricial probability used for class rebalancing from current dataset)

│   └── (wrapper for training)

└── (helper functions)

You can’t perform that action at this time.