A UNet architecture, utilizing transfer learning by using a pretrained ResNet-34 as an encoder.
The model takes a 1x224x224 L tensor as input and outputs 2x224x224 ab channels. The decoder has been trained from scratch. The encoder (ResNet-34) was initially frozen for the decoder to adapt to the task, then it was progressively unfrozen layer by layer. Initial layers were not unfrozen, only deeper layers were fine-tuned. Read various research papers. It took 20+ hours of training on Google Colab and Kaggle T4 GPUs to train the model.
There are no dedicated datasets for image colorisation, hence I curated my own dataset and used it to train the model. The COCO 2017 dataset was filtered to remove grayscale images, heavily filtered images, and other artifacts not suitable for training a natural colorization model. Also the images were center-cropped and resized to 224x224. The dataset can be found
here. This repository contains the model weights and the UNet architecture to load the weights into.
1hf_hub_download(
2 repo_id="ayushshah/imagecolorization",
3 filename="model.py",
4 local_dir=".",
5 local_dir_use_symlinks=False
6)
7
8weights_path = hf_hub_download(
9 repo_id=REPO_ID,
10 filename="model.safetensors"
11)
Make sure the input image(s) are of the size 224x224. Convert them to LAB color space. You can use
kornia.
Isolate the L channel and make sure it is in the range [0, 1]. L channel is originally in the range [0, 100].
1from model import UNet
2from safetensors.torch import load_file
3
4model = UNet().to(DEVICE)
5state_dict = load_file(weights_path)
6model.load_state_dict(state_dict)
7model.eval()
8
9with torch.no_grad():
10 ab_pred = model(L_normalized)
The outputs are in the range [-1, 1]. You can convert the ab channels to their original range using a linear scaling function. Afterwards you can concatenate the original L and the ab channels to get the LAB image.
1ab = (ab+1) * 255.0 / 2 - 128.0
2ab = torch.clamp(ab, -128, 127)
3lab = torch.cat((L, ab), dim=1)
