Colorizer Model

Table of Contents

Objectives
Iterations
- Release v0.0.1
Conclusion

1. Objectives

Given a grayscale image of a manga, comic, or drawing in general, the goal is to output a colorized version of it.
The dataset I gathered for this model contains 1079 images extracted from:

Bleach (Volume 1)
Dragon Ball Super (Volume 21)
Naruto (Volume 1)
One Piece (Volume 99)
Attack on Titan (Volume 1 and 2)

Of those 1079 images, 755 were used for the training (~ 70%), 215 for the validation (~ 20%), and 109 for the test set (~ 10%).

2. Iterations

Release v0.0.1

For the first release, I trained an encoder-decoder model from scratch with the following architecture:

MangaColorizer(
  (encoder): Sequential(
    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (5): ReLU(inplace=True)
  )
  (decoder): Sequential(
    (0): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (5): Tanh()
  )
)

The model was trained for 100 epochs with the following optimizer, using the MSE as the loss function: optimizer = optim.Adam(model.parameters(), lr=0.0001)

From the training history curve, we see that we can probably benefit from increasing the learning rate or training the model for more epochs, as it seems that the loss is still decreasing at the last epochs.
However, we also see that the rate of decrease is not very high, which could indicate that the model doesn't have enough capacity to learn more signal from the data.

Let's have a look at some examples from the test set:

It seems that the model learned to output the same drawing as in the input image, and to add a little bit of color to the image, but it is still nowhere close to a satisfactory result.

Performance on the test set

MSE: 0.008598181701702099

3. Conclusion

To conclude, training a model from scratch is probably not going to cut it for this task, especially given the lack of GPU resources (I use Kaggle's free GPU). The next step is to use some pre-trained models from hugging face and finetune with the manga data.