Image Colorization with Convolutional Neural Networks
Introduction¶
In this post, we're going to build a machine learning model to automatically turn grayscale images into colored images. We'll build the model from scratch (using PyTorch), and we'll learn the tools and techniques we need along the way. In fact, this entire post is an iPython notebook (published here) which you can run on your computer.
At the end of the day, we'll be able to colorize old images or videos (by doing each frame of a video one at a time). Our results will look something like this:
The model used on the clip above is slightly more complex than the model we'll build today, but only slightly. For the full code of that model, or for a more detailed technical report on colorization, you are welcome to check out the full project here on GitHub.
Now, let's dive into colorization.
Overview¶
In image colorization, our goal is to produce a colored image given a grayscale input image. This problem is challenging because it is multimodal -- a single grayscale image may correspond to many plausible colored images. As a result, traditional models often relied on significant user input alongside a grayscale image.
Recently, deep neural networks have shown remarkable success in automatic image colorization -- going from grayscale to color with no additional human input. This success may in part be due to their ability to capture and use semantic information (i.e. what the image actually is) in colorization, although we are not yet sure what exacly makes these types of models perform so well.
Before explaining the model, we will first lay out our problem more precisely.
The Problem¶
We aim to infer a full-colored image, which has 3 values per pixel (lightness, saturation, and hue), from a grayscale image, which has only 1 value per pixel (lightness only). For simplicity, we will only work with images of size 256 x 256, so our inputs are of size 256 x 256 x 1 (the lightness channel) and our outputs are of size 256 x 256 x 2 (the other two channels).
Rather than work with images in the RGB format, as people usually do, we will work with them in the LAB colorspace (Lightness, A, and B) . This colorspace contains exactly the same information as RGB, but it will make it easier for us to separate out the lightness channel from the other two (which we call A and B). We'll make a helper function to do this conversion later on.
We'll try to predict the color values of the input image directly (that is, we do regression). There are other fancier ways of doing colorization with classification (see here), but we'll stick with regression for now as it's simple and works fairly well.
The Data¶
Colorization data is everywhere, as we can extract the grayscale channel from any colored image. For this project, we'll use a subset of the MIT Places dataset of places, landscapes, and buildings.
# Download and unzip (2.2GB)
!wget http://data.csail.mit.edu/places/places205/testSetPlaces205_resize.tar.gz
!tar -xzf testSetPlaces205_resize.tar.gz
# Move data into training and validation directories
import os
os.makedirs('images/train/class/', exist_ok=True) # 40,000 images
os.makedirs('images/val/class/', exist_ok=True) # 1,000 images
for i, file in enumerate(os.listdir('testSet_resize')):
if i < 1000: # first 1000 will be val
os.rename('testSet_resize/' + file, 'images/val/class/' + file)
else: # others will be val
os.rename('testSet_resize/' + file, 'images/train/class/' + file)
# Make sure the images are there
from IPython.display import Image, display
display(Image(filename='images/val/class/84b3ccd8209a4db1835988d28adfed4c.jpg'))
The Tools¶
We'll build and train our model with PyTorch. We'll also use torchvision, a helpful set of tools for working with images and videos in PyTorch, and scikit-learn for converting between RGB and LAB colorspces.
# Download and import libraries
!pip install torch torchvision matplotlib numpy scikit-image pillow==4.1.1
# For plotting
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
# For conversion
from skimage.color import lab2rgb, rgb2lab, rgb2gray
from skimage import io
# For everything
import torch
import torch.nn as nn
import torch.nn.functional as F
# For our model
import torchvision.models as models
from torchvision import datasets, transforms
# For utilities
import os, shutil, time
# Check if GPU is available
use_gpu = torch.cuda.is_available()
The Model¶
Our model is a convolutional neural network. We first apply a number of convolutional layers to extract features from our image, and then we apply deconvolutional layers to upscale (increase the spacial resolution) of our features.
Specifically, the beginning of our model will be ResNet-18, an image classification network with 18 layers and residual connections. We will modify the first layer of the network so that it accepts grayscale input rather than colored input, and we will cut it off after the 6th set of layers.
Now, we'll define our model in code. We'll start with the second half of the net, the upsampling layers:
class ColorizationNet(nn.Module):
def __init__(self, input_size=128):
super(ColorizationNet, self).__init__()
MIDLEVEL_FEATURE_SIZE = 128
## First half: ResNet
resnet = models.resnet18(num_classes=365)
# Change first conv layer to accept single-channel (grayscale) input
resnet.conv1.weight = nn.Parameter(resnet.conv1.weight.sum(dim=1).unsqueeze(1))
# Extract midlevel features from ResNet-gray
self.midlevel_resnet = nn.Sequential(*list(resnet.children())[0:6])
## Second half: Upsampling
self.upsample = nn.Sequential(
nn.Conv2d(MIDLEVEL_FEATURE_SIZE, 128, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Upsample(scale_factor=2),
nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Upsample(scale_factor=2),
nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 2, kernel_size=3, stride=1, padding=1),
nn.Upsample(scale_factor=2)
)
def forward(self, input):
# Pass input through ResNet-gray to extract features
midlevel_features = self.midlevel_resnet(input)
# Upsample to get colors
output = self.upsample(midlevel_features)
return output
Now let's create our model.
model = ColorizationNet()
Training¶
Loss Function¶
Since we are doing regression, we'll use a mean squared error loss function: we minimize the squared distance between the color value we try to predict, and the true (ground-truth) color value.
criterion = nn.MSELoss()
This loss function is slightly problematic for colorization due to the multi-modality of the problem. For example, if a gray dress could be red or blue, and our model picks the wrong color, it will be harshly penalized. As a result, our model will usually choose desaturated colors that are less likely to be "very wrong" than bright, vibrant colors. There has been significant research (see Zhang et al.) on this issue, but we will stick to this loss function for today.
Optimizer¶
We will optimize our loss function (criterion) with the Adam optimizer.
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=0.0)
Loading the data¶
We'll use torchtext to load the data. Since we want images in the LAB space, we first have to define a custom dataloader to convert the images.
class GrayscaleImageFolder(datasets.ImageFolder):
'''Custom images folder, which converts images to grayscale before loading'''
def __getitem__(self, index):
path, target = self.imgs[index]
img = self.loader(path)
if self.transform is not None:
img_original = self.transform(img)
img_original = np.asarray(img_original)
img_lab = rgb2lab(img_original)
img_lab = (img_lab + 128) / 255
img_ab = img_lab[:, :, 1:3]
img_ab = torch.from_numpy(img_ab.transpose((2, 0, 1))).float()
img_original = rgb2gray(img_original)
img_original = torch.from_numpy(img_original).unsqueeze(0).float()
if self.target_transform is not None:
target = self.target_transform(target)
return img_original, img_ab, target
Next we define transforms for our training and validation data.
# Training
train_transforms = transforms.Compose([transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip()])
train_imagefolder = GrayscaleImageFolder('images/train', train_transforms)
train_loader = torch.utils.data.DataLoader(train_imagefolder, batch_size=64, shuffle=True)
# Validation
val_transforms = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224)])
val_imagefolder = GrayscaleImageFolder('images/val' , val_transforms)
val_loader = torch.utils.data.DataLoader(val_imagefolder, batch_size=64, shuffle=False)
Helper functions¶
Before we train, we define helper functions for tracking the training loss and converting images back to RGB.
class AverageMeter(object):
'''A handy class from the PyTorch ImageNet tutorial'''
def __init__(self):
self.reset()
def reset(self):
self.val, self.avg, self.sum, self.count = 0, 0, 0, 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def to_rgb(grayscale_input, ab_input, save_path=None, save_name=None):
'''Show/save rgb image from grayscale and ab channels
Input save_path in the form {'grayscale': '/path/', 'colorized': '/path/'}'''
plt.clf() # clear matplotlib
color_image = torch.cat((grayscale_input, ab_input), 0).numpy() # combine channels
color_image = color_image.transpose((1, 2, 0)) # rescale for matplotlib
color_image[:, :, 0:1] = color_image[:, :, 0:1] * 100
color_image[:, :, 1:3] = color_image[:, :, 1:3] * 255 - 128
color_image = lab2rgb(color_image.astype(np.float64))
grayscale_input = grayscale_input.squeeze().numpy()
if save_path is not None and save_name is not None:
plt.imsave(arr=grayscale_input, fname='{}{}'.format(save_path['grayscale'], save_name), cmap='gray')
plt.imsave(arr=color_image, fname='{}{}'.format(save_path['colorized'], save_name))
Validation¶
In validation, we simply run model without backpropagation using torch.no_grad().
def validate(val_loader, model, criterion, save_images, epoch):
model.eval()
# Prepare value counters and timers
batch_time, data_time, losses = AverageMeter(), AverageMeter(), AverageMeter()
end = time.time()
already_saved_images = False
for i, (input_gray, input_ab, target) in enumerate(val_loader):
data_time.update(time.time() - end)
# Use GPU
if use_gpu: input_gray, input_ab, target = input_gray.cuda(), input_ab.cuda(), target.cuda()
# Run model and record loss
output_ab = model(input_gray) # throw away class predictions
loss = criterion(output_ab, input_ab)
losses.update(loss.item(), input_gray.size(0))
# Save images to file
if save_images and not already_saved_images:
already_saved_images = True
for j in range(min(len(output_ab), 10)): # save at most 5 images
save_path = {'grayscale': 'outputs/gray/', 'colorized': 'outputs/color/'}
save_name = 'img-{}-epoch-{}.jpg'.format(i * val_loader.batch_size + j, epoch)
to_rgb(input_gray[j].cpu(), ab_input=output_ab[j].detach().cpu(), save_path=save_path, save_name=save_name)
# Record time to do forward passes and save images
batch_time.update(time.time() - end)
end = time.time()
# Print model accuracy -- in the code below, val refers to both value and validation
if i % 25 == 0:
print('Validate: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
i, len(val_loader), batch_time=batch_time, loss=losses))
print('Finished validation.')
return losses.avg
Training¶
In training, we run model and backpropagate using loss.backward(). We first define a function that trains for one epoch:
def train(train_loader, model, criterion, optimizer, epoch):
print('Starting training epoch {}'.format(epoch))
model.train()
# Prepare value counters and timers
batch_time, data_time, losses = AverageMeter(), AverageMeter(), AverageMeter()
end = time.time()
for i, (input_gray, input_ab, target) in enumerate(train_loader):
# Use GPU if available
if use_gpu: input_gray, input_ab, target = input_gray.cuda(), input_ab.cuda(), target.cuda()
# Record time to load data (above)
data_time.update(time.time() - end)
# Run forward pass
output_ab = model(input_gray)
loss = criterion(output_ab, input_ab)
losses.update(loss.item(), input_gray.size(0))
# Compute gradient and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Record time to do forward and backward passes
batch_time.update(time.time() - end)
end = time.time()
# Print model accuracy -- in the code below, val refers to value, not validation
if i % 25 == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
epoch, i, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses))
print('Finished training epoch {}'.format(epoch))
Next, we define a training loop and we train for 100 epochs:
# Move model and loss function to GPU
if use_gpu:
criterion = criterion.cuda()
model = model.cuda()
# Make folders and set parameters
os.makedirs('outputs/color', exist_ok=True)
os.makedirs('outputs/gray', exist_ok=True)
os.makedirs('checkpoints', exist_ok=True)
save_images = True
best_losses = 1e10
epochs = 100
# Train model
for epoch in range(epochs):
# Train for one epoch, then validate
train(train_loader, model, criterion, optimizer, epoch)
with torch.no_grad():
losses = validate(val_loader, model, criterion, save_images, epoch)
# Save checkpoint and replace old best model if current model is better
if losses < best_losses:
best_losses = losses
torch.save(model.state_dict(), 'checkpoints/model-epoch-{}-losses-{:.3f}.pth'.format(epoch+1,losses))
Pretrained model¶
If you would like to run with a pretrained model rather than trained one from scratch, I've trained one for you. It's trained on a relatively small amount of data for a small amount of time, but it works. You can download and play around with it from the link below:
# Download pretrained model
!wget https://www.dropbox.com/s/kz76e7gv2ivmu8p/model-epoch-93.pth
#https://www.dropbox.com/s/9j9rvaw2fo1osyj/model-epoch-67.pth
# Load model
pretrained = torch.load('model-epoch-93.pth', map_location=lambda storage, loc: storage)
model.load_state_dict(pretrained)
# Validate
save_images = True
with torch.no_grad():
validate(val_loader, model, criterion, save_images, 0)
Results¶
Time for the fun part: let's see our results!
# Show images
import matplotlib.image as mpimg
image_pairs = [('outputs/color/img-2-epoch-0.jpg', 'outputs/gray/img-2-epoch-0.jpg'),
('outputs/color/img-7-epoch-0.jpg', 'outputs/gray/img-7-epoch-0.jpg')]
for c, g in image_pairs:
color = mpimg.imread(c)
gray = mpimg.imread(g)
f, axarr = plt.subplots(1, 2)
f.set_size_inches(15, 15)
axarr[0].imshow(gray, cmap='gray')
axarr[1].imshow(color)
axarr[0].axis('off'), axarr[1].axis('off')
plt.show()
Conclusion¶
In this post, we built a (simple) automatic image colorizer from scratch in PyTorch!
If you found these results exciting, there are lots of places to go from here:
- Incorporating user feedback into colorization: Zhang et al., 2017
- Using colorization for unsupervised learning: Larsson et al. , 2017
- Add it to a camera app, like Google just did
- Colorize a scene from an old film, as I experimented with
Also check back here for more posts coming soon -- next up is language modeling!