DeOldify: a free/open photo-retoucher based on machine learning

Jason Antic's DeOldify is a Self-Attention Generative Adversarial Network-based machine learning system that colorizes and restores old images. It's only in the early stages but it's already producing really impressive results, and the pipeline includes a "defade" model that is "just training the same model to reconstruct images that augmented with ridiculous contrast/brightness adjustments, as a simulation of fading photos and photos taken with old/bad equipment."

This is a deep learning based model. More specifically, what I've done is combined the following approaches:

* Self-Attention Generative Adversarial Network (https://arxiv.org/abs/1805.08318). Except the generator is a pretrained Unet, and I've just modified it to have the spectral normalization and self-attention. It's a pretty straightforward translation. I'll tell you what though - it made all the difference when I switched to this after trying desperately to get a Wasserstein GAN version to work. I liked the theory of Wasserstein GANs but it just didn't pan out in practice. But I'm in love with Self-Attention GANs.

* Training structure inspired by (but not the same as) Progressive Growing of GANs (https://arxiv.org/abs/1710.10196). The difference here is the number of layers remains constant - I just changed the size of the input progressively and adjusted learning rates to make sure that the transitions between sizes happened successfully. It seems to have the same basic end result - training is faster, more stable, and generalizes better.

* Two Time-Scale Update Rule (https://arxiv.org/abs/1706.08500). This is also very straightforward - it's just one to one generator/critic iterations and higher critic learning rate.

* Generator Loss is two parts: One is a basic Perceptual Loss (or Feature Loss) based on VGG16 - this just biases the generator model to replicate the input image. The second is the loss score from the critic. For the curious - Perceptual Loss isn't sufficient by itself to produce good results. It tends to just encourage a bunch of brown/green/blue - you know, cheating to the test, basically, which neural networks are really good at doing! Key thing to realize here is that GANs essentially are learning the loss function for you - which is really one big step closer to toward the ideal that we're shooting for in machine learning. And of course you generally get much better results when you get the machine to learn something you were previously hand coding. That's certainly the case here.

DeOldify [Jantic/Github]

(via Four Short Links)