RGBIT.io
RGBIT is a completely open-source and free platform for restoring color to black and white images.
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.
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.
Model results on real legacy black and white photographs:
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 andml/code/visualize.pyto 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.
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+
Trained on Places365
Deployed with AWS EC2, Route 53, and 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!
Log in or sign up for Devpost to join the conversation.