This is an attempt to reimplement the CHROMAGAN paper using pytorch.
So the paper looks forward to tackle a simple probelem and that is to colourize black and white photos without any human intervention. The paper uses a Generative Adversarial Network to handle this. The Model training is based on the fact that for each given grayscale image L, the model tries to learn a mapping G:L -->(a,b) such that the image I = (L,a,b) is a plausible colour image.
- the easiest part of this project is to implement the model architecture
As we can see, the architecture is denoted by several colour schemes and each colour scheme has different importance.
- Yellow branch is a simple VGG network with no top layer
- Small red boxes away from the main branches are modules consisting of Conv-BatchN-Relu blocks
- Grey blocks are fc layers which extract the semantic meaning by classifying the classes.
- Red blocks are again connected back to the main branch and is concatenated with the purple block.
- the blue blocks are upsample conv blocks.
These blocks make the colourization_block/generator block(please check the paper for more details)
The other block is the discriminator block which is based on a simple Markovian Discriminator (PatchGan used in Image translations tasks).
- Next up is with loss functions and the training. Accoriding to the author of the paper, they used multiple loss functions to learn the two chrominance value - Colour Error Loss ( L2 Norm Loss ) - Class distribution loss (KL divergence ) - WGAN loss (adverserial Wassertien GAN loss)
The original paper was trained on Imagenet. I trained it on a few images(50k)
Results of their paper and every minute detail can be found in pvitoria repo
The result of this repository are given below
Please do cite the original paper if you like the repo and please do check out the official keras implementation of the ChromaGan. LINK to the repo