This weekend project started as an image compression idea. A grayscale image keeps one channel instead of the three channels of RGB. I wondered whether a neural network could learn enough about a narrow domain to reconstruct plausible colours from that missing information.

It could make sense in a very specific setting. If you wanted to compress episodes of one series, for example, a model trained on that series might learn its recurring characters, locations, and colour palette. It would not recover the original image exactly, because the grayscale frame tells you about lightness, edges, and structure, but not whether a car was red or blue, but it was something worth experimenting with.

I started this project with a small convolutional network that took one grayscale channel and predicted three RGB channels with MSE loss. The loss went down quickly and then stopped improving around the same point, regardless of whether I use a deeper architecture or a ResNet backbone. The resulting images were mostly beige or sepia.

I thought this would be a small experiment, but this first failure made it clear that the model was not the only problem. With many plausible colours for the same object, a model trained with MSE can move towards an average-looking answer, which turns out to be a beige image.

Reimplementing a paper

I started looking for work on image colorization and found Let there be Color!, a 2016 architecture for automatic image colorization. Its formulation made more sense for the task.

The network receives the L channel from the LAB colour space, which contains the lightness of the image. It predicts only the a and b channels, which represent chrominance. The original L channel is kept and combined with the prediction at the end.

This is a better way of approaching the problem, because the network does not need to recreate the contrast or geometry that is already visible in the grayscale image. It can focus only on assigning plausible colour to each region.

Global and local features

The paper also uses two branches after the first convolutional layers. One branch preserves the local spatial information needed to colour a particular region. The other compresses the image into global features and predicts a scene category from Places365.

The scene classification is an auxiliary task. It gives the network some context about the whole image, which is useful when local appearance alone is ambiguous. A blue patch can be sky, water, or a wall, depending on the image’s class. The global features are fused back with the local features before the final layers predict the two chrominance channels.

Architecture of the model, adapted from the original paper.

I reimplemented the 2016 network in PyTorch, including its local and global feature branches, the auxiliary scene classification task, and the colourization layers that predict the a and b channels from L.

I also added checkpoints, resume support, Weights & Biases logging, dataset subsets, and a small Makefile interface around Places365 to make the training runs easier to manage. With that setup, I reproduced the training procedure and colourization results from the 2016 paper.

↑ Batch loss

Batch lossMoving average

↑ Training

TrainingValidation

Results

The greenhouse is one of the better examples. The model understands the broad structure and produces plausible greens and browns, although they are less saturated than the original. The mountain scene and the dinner image are also recognisable, even when the chosen colours do not match exactly.

Grayscale input

Model output

The model runs on each cropped frame independently.

The cartoon is much worse because it is completely outside the training distribution. Its colours follow conventions that a model trained on natural scenes has no reason to know. The park video works reasonably well frame by frame, but it has the same tendency to use conservative, desaturated colours.

Fine-tuning for a specific domain

I wanted to see how much a narrower domain would change the result. I fine-tuned the Places365 checkpoint on frames from a forest video. I randomly selected one continuous 20-second segment and excluded it completely from training and validation. The model kept the original LAB architecture, with the scene classifier adapted to a single forest class.

I coloured the same held-out forest clip with both models. The fine-tuned model reduced LAB colour error by 77% and avoided the orange cast that the Places365 model gave to the green forest. I repeated the same experiment in a separate fine-tuning run on Barcelona footage. The Places365 model made many more colour shifts, while the model fine-tuned for that scene kept the colours much more coherent across frames. These comparisons suggest that the approach works much better when the model is trained or fine-tuned for a specific domain.

Held-out forest clip

Original video

Places365 model

Fine-tuned model

The 20-second forest clip was excluded from training and validation.

Barcelona clip

Original video

Places365 model

Fine-tuned model

A separate fine-tuning run trained on Barcelona footage.

That is the limit of the original compression idea. A grayscale image does not contain the exact colour information needed to restore the original frame. A well-trained model can make a plausible guess, and that can be useful in a narrow domain, but it is still a guess.

The project includes training commands, checkpoints, image inference, and video inference in the repository. Reimplementing the paper was the way I learned what those inputs, outputs, and limitations look like once the model has to train on a real dataset for several days.