Colourisation with a GAN

The images were generated with the Flux-dev model and then edited by the author of the task.

Introduction

The subject of this task is image quality improvement (image enhancement), a collective term covering activities such as:

• increasing the resolution (image super-resolution),
• denoising (image denoising),
• sharpening (image deblurring),
• brightening (low-light image enhancement),
• colourisation (image colorization),
• and many others.

The aim of these techniques is to obtain a high-quality image from a low-quality photograph. This process often requires filling in missing information in the image on the basis of its context, which is why deep learning methods are commonly used.

Thanks to such techniques, we can in a sense "look" into the past from a new perspective. An example is the photograph of Warsaw below, from the early interwar period – on the left in the original, and on the right in a version colourised by Mariusz Zając.

In this task we will focus on colourising black-and-white photographs showing human faces.

Colourising photographs

A colour photograph is described by three RGB channels (red, green, blue). It can easily be converted into a monochrome image described by a single channel. However, the conversion to greyscale does not use a simple arithmetic mean of the RGB values, because that does not correspond to the way the human eye perceives the brightness of individual colours (e.g. we perceive blue light as darker than green light). For this reason, the pixels of a greyscale image are computed according to the formula:

$$x^{(gray)} = 0.299 \cdot x^{(red)} + 0.587 \cdot x^{(green)} + 0.114 \cdot x^{(blue)}.$$

It is worth noting that the conversion from the RGB space to greyscale is unambiguous and easy to perform. The reverse operation – recovering the colour information – is not unambiguous, because many different combinations of RGB values can lead to the same greyscale value. This means that image colourisation is an ill-posed task, which implies that many correct colour versions can be generated for a single black-and-white photograph.

For example, a black-and-white photograph of a car is difficult to colourise unambiguously – the vehicle could have been almost any colour. On the other hand, models trained on suitable data are able to exploit statistical regularities – such as the fact that grass is usually green, the sky blue, and a tiger orange and black – in order to generate the most probable colour versions.

Formally, we define the colourisation task as learning a model $\mathcal{M}_\theta$ with parameters $\theta$, which takes a greyscale image $x^{(gray)}$ as input and generates its colour version $\hat{x}^{(rgb)}$ as output:

$$\hat{x}^{(rgb)} = \mathcal{M}_\theta(x^{(gray)}).$$

We would like the model's predictions $\hat{x}^{(rgb)}$ to resemble the original colour images $x^{(rgb)}$ as closely as possible. To this end, we seek values of the parameters $\theta$ that minimise some measure of difference (e.g. the mean squared error) between the predictions and the real images:

$$\theta^* = \arg \min_{\theta} \sum_{i = 1}^N d(x^{(rgb)}_i, \mathcal{M}_\theta(x^{(gray)}_i)),$$

where $N$ denotes the number of examples in the training set.

Generative adversarial networks

Generative adversarial networks (GANs) were one of the first approaches to address the problem of generating data using deep learning techniques. Since the generation process does not require any input to the network, it is not possible to propose a loss function in a simple way, because we do not know what output to expect from the network. For example, the network may generate a plausible image of a face, but when we compare it with another random photograph of a face, the penalty will be very large, even though the network's output is correct.

For this reason, the GAN architecture consists of two networks: a generator and a discriminator.

• The role of the generator is to generate new, random images.
• The role of the discriminator is to predict whether the input image comes from the dataset or was generated by the generator.

The generator tries to fool the discriminator (it maximises the discriminator's classification error), whereas the discriminator tries to guess whether the image is artificial or not (it minimises its classification error). The two networks are trained alternately until convergence is reached.

In this task we will not be training a GAN; instead, we will have at our disposal a previously trained generator model which, from noise, can generate a random image of a face with a resolution of $256 \times 256$. The architecture of this network is called StyleGAN and is presented in the diagram below. It is divided into two modules: a mapping network $f$ and a generating network $g$.

The network $f$ is an MLP architecture that transforms a random vector $z$, with $512$ elements whose values are drawn from the normal distribution $\mathcal{N}(0, \mathbb{I})$, into a $512$-dimensional vector $w$, which is used to condition the network $g$. The vector $w$ can be interpreted as a latent representation of the generated image.

The network $g$ is a convolutional network with a constant $4 \times 4$ input, which progressively increases its resolution up to $256 \times 256$. In each block, the network $g$ is conditioned on the vector $w$, which ensures the diversity of the generated images. In addition, some noise is also fed into the network, which increases the diversity even further.

Task

In this task, given a trained generator capable of generating face images with a resolution of $256 \times 256$, you must propose a method for colourising black-and-white face photographs of the same resolution. The aim here is to extract the knowledge that lies dormant in the generator's weights and to use it in a completely different task, previously unknown to the generator.

Data

No training data are available in this task, only validation data for a preliminary assessment of the approach you propose. The validation data contain 500 black-and-white photographs paired with their colour counterparts.

Scoring criterion

To score the quality of your solution, a metric consisting of two sub-metrics will be used: PSNR and LPIPS.

PSNR is inversely proportional to the mean squared error between the generated sample $\hat{x}^{(rgb)}$ and the real colour image $x^{(rgb)}$. It is given by the formula

$$PSNR = 10 \cdot \log_{10}\left(\frac{\text{range}^2}{MSE(\hat{x}^{(rgb)}, x^{(rgb)})}\right),$$

where $\text{range}$ is the range of possible values that the images $\hat{x}^{(rgb)}$ and $x^{(rgb)}$ can take. We want to maximise PSNR. The scored range of values of this metric is $(22.0, 26.0)$ - points are awarded proportionally.

LPIPS is the mean $L1$ distance between the feature maps after selected layers of the AlexNet network, obtained for the inputs $\hat{x}^{(rgb)}$ and $x^{(rgb)}$. We want to minimise LPIPS. The scored range of values of this metric is $(0.15, 0.11)$ - points are awarded proportionally.

The final point value is the weighted mean of the points obtained from the PSNR and LPIPS metrics, with weights $0.25$ (PSNR) and $0.75$ (LPIPS).

Constraints

• Your solution will be tested on the Contest Platform without internet access and in an environment with a GPU.
• The evaluation of your final solution on the Contest Platform must not take longer than 5 minutes with a GPU.
• List of permitted libraries: torch, numpy, torchvision, pillow.

Submission files

You must submit only this notebook, completed with your solution (see the class YourModel).

Evaluation

During grading, the FINAL_EVALUATION_MODE flag will be set to True.

You can score between 0 and 100 points for this task. The number of points you receive will be calculated on the (secret) test set on the Contest Platform according to the formula given above, rounded to an integer. If your solution does not meet the criteria above or does not run correctly, you will receive 0 points for the task.

Starter code

The cell below contains the functions that score and visualise your solutions, as well as the validation dataset.

The cell below contains code for generating faces with the StyleGAN model. The generation is presented in two versions. The first version is more detailed and consists of the following steps:

• first, we generate $z$ from the normal distribution
• next, through the get_latent method, we use the model $f(z)$ to generate the vectors $w$
• finally, we call the style_to_image method, which uses the network $g$ to generate face images conditioned on the vector $w$.

Alternatively, we can generate images with the forward method, which combines steps 2 and 3.

Your solution

Implement your solution in the cell below. Make sure that the fit method prepares the model that colourises the images, while the predict method uses the model to colourise the input data x_gray.