Teaching a Neural Network to See Colour

My final year project: building a hybrid CNN-Transformer model to colourise greyscale images from scratch

For my final year project at TAR UMT, I built a deep learning model that takes a greyscale image as input and outputs a plausible colourised version of it. Sounds deceptively simple. It wasn't.

This post is a rough account of how I approached it, the design decisions I made, and honestly, where things went sideways too.

Why colourisation?

A lot of historical photographs and archival footage only exist in black and white. Colourising them by hand in Photoshop is a legitimate profession — and a painfully slow one. Getting the skin tones on a single portrait right can take up to 20 layers of carefully blended pink, green, and blue. For an entire archive of images, that just doesn't scale.

The deep learning angle felt like a natural fit. But as I dug into the literature, I realised why it's still an open problem: colourisation is inherently ill-posed. Given a pixel's greyscale value alone, there are infinitely many valid colours it could be. An apple could be red, green, or yellow. A building at dusk could be lit orange, grey, or purple. The model doesn't need to find the correct colour — it needs to find a plausible one.

That's actually a harder objective than it sounds, because naive approaches like L2 regression just average all plausible solutions together and produce washed-out, desaturated outputs. Which is exactly the failure mode I spent the whole project trying to avoid.

The architecture: CNN encoder + Transformer + decoder

The core insight driving my design was that the two dominant architectures each solve different parts of the problem:

CNNs are great at picking up local texture and edge information — the kind of low-level detail that tells you "this is skin" or "this is grass." But because their receptive fields are local by design, they really struggle with large uniform regions. Things like a clear sky or a body of water end up desaturated and muddy because the network can't reason about the region as a whole.

Transformers on the other hand are built for global context. Self-attention lets every patch in the image attend to every other patch, so long-range colour consistency becomes tractable. The downside is they don't naturally extract the kind of fine spatial features that CNNs are so good at.

So I combined them. The full pipeline looks like this:

Grayscale image (L channel)
    → CNN Encoder          [local feature extraction]
    → Transformer Encoder  [global context modelling]
    → Colour Decoder        [upsampling + skip connections → AB channels]

The CNN encoder

Four conv blocks with progressively increasing channels (64 → 128 → 256 → 256), using stride-2 convolutions to downsample. BatchNorm and ReLU throughout. Nothing exotic — this part just needed to be a solid feature extractor.

self.conv2 = nn.Sequential(
    nn.Conv2d(64, 128, 3, stride=2, padding=1),
    nn.BatchNorm2d(128),
    nn.ReLU(inplace=True),
)

The encoder returns the final feature map and the intermediate activations from each block — those get used later as skip connections.

The Transformer encoder

I used PyTorch's built-in TransformerEncoderLayer with 8 attention heads and 4 layers. Before passing the CNN features in, I reshape them into a sequence of patches (patch size 4×4) and project each patch down to the model dimension (256). Then learnable positional embeddings go on top.

x = x.unfold(2, self.patch_size, self.patch_size)\
     .unfold(3, self.patch_size, self.patch_size)
x = x.contiguous().view(B, C, -1, self.patch_size * self.patch_size)
x = x.permute(0, 2, 1, 3).flatten(2)
x = self.patch_projection(x)
x = x + pos_emb
x = self.transformer(x)

This is essentially a lightweight ViT-style encoder grafted onto the top of the CNN. The Transformer gets to see the full spatial context of the image in one shot, which is exactly what you want when trying to colour a wide open sky consistently.

The colour decoder

The decoder upsamples back to the original resolution using transposed convolutions, and uses skip connections from the CNN encoder at each scale — similar to a U-Net. The final layer outputs 2 channels (the A and B components of LAB colour space), activated with Tanh to keep the values in [-1, 1].

Why LAB colour space?

If you're doing colourisation, working in LAB is the obvious choice. The L channel carries all the luminance (lightness) information, while A and B carry the colour. Since the input is already greyscale (i.e., we have L), we just need to predict A and B.

This decoupling means the model never has to touch luminance at all — we keep the original L channel from the input and just paste the predicted AB channels on top to reconstruct the full colour image. Much cleaner than trying to predict RGB directly.

Training

I trained on the Places365 dataset, running on Kaggle with multi-GPU support via DataParallel. Loss function was L1 (MAE between predicted and ground truth AB channels), optimiser was Adam at 3e-4. Ran for around 55 epochs.

One practical thing I had to handle was that the Transformer processes a fixed sequence length, but at different image sizes the number of patches changes. I handled this with interpolation on the positional embeddings when the number of patches exceeded the fixed budget — a bit hacky, but it works.

Results

Quantitative:

The SSIM number I'm happiest with — 0.9277 means the structural content of the image is very well preserved. Edges, object boundaries, and overall scene composition come through clearly. The PSNR at 23.80 dB is solidly above the "visually acceptable" threshold of 20 dB.

How does it stack up against existing work?

Benchmarking against historical models on Places205 (the closest comparable dataset with published numbers):

Sitting just below DeOldify and above Colorful Colorization. Not top of the table, but for a from-scratch architecture built for a final year project and trained on Kaggle compute, I'll take it.

Results after many epochs of training:

Where it struggles

Being honest: there are real failure modes.

Large uniform regions (skies, oceans) are the main one. The Transformer was supposed to help with this, and it does improve things, but the model still occasionally produces unrealistic or inconsistent colours in wide open skies. My suspicion is that the patch size and the training data distribution both play a role here.

Colour bleeding is another visible issue. In complex scenes — trees in front of buildings, faces against busy backgrounds — the model sometimes lets colours from one region bleed into adjacent objects. This points to a lack of semantic awareness: the model doesn't really "know" what objects are in the scene, so it can't respect their boundaries properly.

Reflective surfaces are basically unsolved. Water reflections in particular produce very inconsistent outputs.

These aren't surprising failures for a model trained purely on photometric reconstruction loss. To really fix them, you'd probably want to incorporate semantic segmentation features, or use a perceptual / adversarial loss to push the model toward more globally coherent predictions.

What I'd do differently

A few things I'd revisit if I were starting over:

• Add a GAN discriminator — L1 loss is safe but conservative. An adversarial loss pushes the model to commit to more saturated, realistic colours rather than averaging toward grey.
• Larger Transformer — I kept it small (4 layers, 8 heads) for compute reasons. More capacity here would probably help global consistency.
• Semantic guidance — feeding in class labels or segmentation masks as conditioning would give the model a fighting chance at knowing what colour something should be, rather than relying purely on texture cues.
• Better positional embeddings — the interpolation trick for variable-length sequences is workable but not elegant. A proper RoPE or sinusoidal scheme would be cleaner.

Final thoughts

This project taught me more about the gap between "reasonable idea" and "working implementation" than anything else I've done academically. The architecture makes sense on paper, and the results are respectable — but there's a clear ceiling on what you can achieve with pure reconstruction loss and limited compute.

If you've worked on colourisation or similar image-to-image tasks and have thoughts on the architecture or the failure modes, I'd genuinely be curious to hear them. This is the kind of problem where there's still a lot of headroom.

Built with PyTorch. Trained on Kaggle. Evaluated on Places365.