Colorization for Image Compression

Colorization for Image Compression proposes a method to colorize gray-scale images for use in a compression setting. We propose a new architecture and a method for training our new architecture
with the end goal of achieving superior performance on the compression task.

To run this code, please build the provided version of caffe on your machine. This should be followed by 'make pycaffe'. Once pycaffe has compiled, move the file path_to_caffe/python/caffe/proto/caffe_pb2.py to the directory path_to_caffe/python/colorization/ .

This repository contains the code and the model that was used for training on the Imagenet dataset.

The code builds on the Caffe framework and adds layers to the caffe framework. The following new layers are introduced
- Modified Euclidean loss layer for GPU to allow for training with Branched objective.
- Custom Data Layer, Image Data Layer and Transformations layer (for colorization). ( To allow for reading 1 Luma and 2 Chroma files for Training ).
- Python script for generating the Resnet-50 Architecture with the branching module

The directory examples contains a subdirectory "colorization". To train the model on the imagenet dataset please follow the following steps.

Create the training dataset by running the matlab script "generate_imagenet_dataset.m" provided under path_to_caffe/matlab/colorization/
You should set paths in the script to point to your dataset and the destintation directory where you want to store the newly created dataset
This script generates the following
Y channel images
Cb and Cr channel images
train_paths.txt and test_paths.txt file containing paths to these files for train and validation sets
It also computes a single average value for the intensity channel and the Cb and Cr channels and stores them in a file averages.txt in the directory path_to_colorization_caffe/examples/colorization/

Run the python script generate_training_arch.py which creates the train.prototxt file in the path_to_colorization_caffe/examples/colorization/
-Before running this script please update appropriate paths at the begining of the script to point to the paths being used.

Open solver.prototxt in "examples/colorization/" and update the path to where the snapshots would be stored
-update the path that points to the training script

Create a sub-directory called LOG_TRAIN under the directory path_to_colorization_caffe/examples/colorization/

now enter the directory path_to_colorization_caffe/examples/colorization/ in the terminal and run ./train_color.in.sh

Please cite Colorization for Compression and Caffe in your publications if it helps your research: