This project implements an image colorization model inspired by the paper Colorful Image Colorization by Zhang et al., using TensorFlow.
Like Zhang et al., we treat colorization as a classification problem, predicting discrete color classes instead of directly estimating color values per pixel. We preprocess images from RGB to LAB color space, which has a single grayscale lightness channel (L) and two color channels (A and B). We also quantize the whole AB space into 40 discrete color bins, allowing us to treat the task as a classification problem. For training, each pixel's color values are mapped to the nearest of these 40 "classes" and encoded as a one-hot vector. Our model is then trained using the below architecture and categorical cross-entropy loss to predict color classes for inputted grayscale images.
Examples of colorized images generated by the model:
Our Model Architecture:
Zhang, R., Isola, P., & Efros, A. A. (2016). Colorful Image Colorization.
Paper Link
- Upload the files that you would like to train the model on.
- Choose the amount of files to be trained.
- Choose which images you would like modified.
- Run the Jupyter Notebook provided.