RGBIT is a completely open-source and free platform for restoring color to black and white images.
We implemented a convolutional neural network (CNN) to colorize grayscale images using a U-Net architecture with the VGG-19 model. U-Net is a popular deep learning architecture known for its effectiveness in image segmentation tasks. VGG-19 is a large model with almost 150 million parameters that is pre-trained. It is traditionally used for feature detection and was adapted for colorizing in our project. Our model is trained using the MIT Places365 dataset, which contains 365,000 images of scenes (which we split into 328,500 train and 36,500 test images, a 90/10 split). Moreover, the model makes use of a custom Perceptual Loss function for a higher level chromatic evaluation of the CNN. Our results show that the model produces vibrant and realistically colored images. This project reinforces the potential of deep learning in creative image processing. Below is our VGG-19 U-Net architecture.
The results of perceptual loss showed our model and architecture is viable for creating naturally looking colorful photos, but doesn't correctly account for unique coloring and saturation. Colors returned are plausible and look natural to the human eye. The model can be used to color any grayscale image, but has best use-cases for naturally existing photos, such as old black and white photography or night vision goggles. Below are some example results from our model test dataset. The first image is the L channel, the second image is the truth coloring, and the third image is the predicted coloring.
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
Model results on real legacy black and white photographs:
Note: Local client and server setup requires a cloud MongoDB URI and generating RSA keys. The database is for storing results, and the RSA keys are for API key authentication. Feel free to skip around the local setup to find what you're looking for. Here covers using just the
ml/
directory andml/code/visualize.py
to colorize images from the terminal.
Clone the repository:
$ git clone https://github.com/johnsfarrell/rgbit
In most cases, you'll be using the ml
directory (for training or colorizing with pre-trained model). Please refer to ./setup.sh
for a list of commands for setting up the Python virtual environment.
$ ./ml/setup.sh
If you want to play with the pre-trained model, make sure to have ml/models/model.h5
downloaded.
If you are interested in modifying the model architecture and training a model yourself, see here.
Setup .env
files in client
and server
directories:
$ cp client/.env.example client/.env
$ cp server/.env.example server/.env
The client
and server
directories depend on MongoDB and RSA keys. Obtain a MongoDB URI and generate RSA keys. Update the .env
files with the appropriate values.
Once you have the model.h5
file in ml/models
, you can run the server and both .env
files are set up, use the go
script to run the application:
$ ./go
The go
script should let you know if you missed one of the steps above, and will promp you to install the necessary dependencies if you haven't already.
Alternatively, you can use Docker to run the application. Make sure you have Docker installed on your machine.
Follow the same steps as above to set up the .env
files in the client
and server
directories.
For server/.env
, use COLORIZE_FLASK_URL=http://ml:4004/api
instead of localhost
.
$ docker-compose up
You can also use ml/code/visualize.py
to colorize images without the client and server. Please make sure you have ml/models/model.h5
downloaded.
$ cd ../code && python3 visualize.py # -h for help
If you're interested in hosting a local, light-weight server for colorizing images, use ml/code/api.py
. Please make sure you have ml/models/model.h5
downloaded.
$ cd ml/code/api.py && python3 api.py
To train your own model, start by downloading the training data:
$ cd ml/data && ./download.sh
Feel free to modify the model architecture (model.py
), hyperparameters (hyperparameters.py
), and dataset preprocessing and augmentation (preprocess.py
). When you're ready, train the model by running:
$ cd ml/code && python3 main.py
Navigate to http://localhost:3000
to use the application.
The API is available at http://localhost:4004
.
The colorize endpoint is available at http://localhost:4004/colorize
.
Command | Action |
---|---|
npm install |
Installs the client dependencies |
npm run start |
Starts the client |
npm run build |
Builds the client |
npm run deploy |
Deploys the client to GitHub Pages (for production) |
Command | Action |
---|---|
npm install |
Installs the server dependencies |
npm run start |
Starts the server |
Command | Action |
---|---|
./data/download.sh |
Downloads the training data |
./setup.sh |
View setup commands |
source venv/bin/activate |
Activates the virtual environment |
pip install -r requirements.txt |
Installs the ML dependencies |
python3 code/api.py |
Starts the Flask server |
python3 main.py |
Trains the model |
python3 code/visualize.py -h |
Commands for visualizing test image(s) |
Trained on Places365
The server is deployed on AWS EC2. As an extra layer of security, the Flask ML server is hosted locally on the EC2 instance, and can only be accessed through the Express server. The Flask microservice also isolates the ML development from the REST development. As EC2 hosts a virtual machine (Linux for this project), we launched multiple servers and configure the ports through an EC2 security group. This makes it easy to redirect the Express server with Caddy to api.rgbit, while making the Flask server inaccessible to the internet. The server runtime is maintained by systemctl
. If you were interested in hosting your own API as such, check out my guide!
-
tfjs
client-only deployment - Refactor UI workflow
- Docker deployments
- Official favicon
- API endpoint for getting total images colorized
- EC2 instance optimiztaion or downgrade
- HEIC/HEIF image upload support
-
https://rgbit.io
only authorized API calls for managing API keys - Limit API requests by IP address
- Limit file upload size to API
- API endpoint for getting total users
- HEIC image support
- Downscale oversized files
Contributions from Tyler Gurth, John Farrell, Jania Vandevoorde, and Hunter Adrian.
Thank you to Brown University Professor Srinath Sridhar and Jiahua Chen for advice and insights.
Disclaimer: When using our live deployment or api, you agree to our terms and conditions. This includes storing user uploaded images. If you choose to create your own instance, we do not track your usage. License.