Skip to content

johnsfarrell/rgbit

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

banner

research paper Colab Postman pages-build-deployment License: MIT

RGBIT is a completely open-source and free platform for restoring color to black and white images.


Resources and Contents

Model Summary

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.

Architecture

Example Results

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.

Places365_val_00015300 Places365_val_00016263 Places365_val_00013593
Places365_val_00017696 Places365_val_00011056 Places365_val_00006786
Places365_val_00001475 Places365_val_00001356 Places365_val_00033979
Places365_val_00008392 Places365_val_00007964 Places365_val_00007777

Model results on real legacy black and white photographs:

legacy-banner

Local Setup

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 and ml/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.

Docker Setup

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

Clientless and Serverless Colorization

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

Flask Colorization

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

Training

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

Usage

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.

Client Commands

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)

Server Commands

Command Action
npm install Installs the server dependencies
npm run start Starts the server

ML Commands

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)

Built on MERN+

MongoDB Express.js React NodeJS Flask Chakra

Trained on Places365

Keras TensorFlow scikit-learn Matplotlib NumPy

Deployed with AWS EC2, Route 53, and Pages

AWS Github Pages

More on API Deployment

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!

Changelog & Todos πŸ“

  • 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

Acknowledgements

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.