Bref, la tâche de ce projet est de coloriser et restaurer des photographies anciennes. Je vais approfondir un peu les détails, mais d'abord, voyons les photos! Soit dit en passant, la plupart des images source sont tirées du sous-r / TheWayWeWere, je remercie tout le monde pour ces gros plans de haute qualité.

Ce ne sont que quelques exemples, et ils sont assez typiques!

Maria Anderson comme la petite fée et sa page Lyubov Ryabtsova dans le ballet Sleeping Beauty à l'Imperial Theatre, Saint-Pétersbourg, Russie, 1890

Une femme se détend dans son salon (1920, Suède)

Étudiants en médecine posant près d'un cadavre, vers 1890

Surfeur à Hawaï, 1890

Cheval qui tourne, 1898

L'intérieur du bar Miller and Shoemaker, 1899

Paris dans les années 1880

Vue aérienne d'Édimbourg dans les années 1920

Texas femme en 1938

Les gens de la gare de Waterloo regardent la télévision pour la première fois, Londres, 1936

Leçon de géographie en 1850

Les fumeurs d'opium chinois en 1880

Veuillez noter que même les photos vraiment anciennes et / ou de mauvaise qualité ont toujours l'air plutôt cool:

Deadwood, Dakota du Sud, 1877

Frères et sœurs en 1877 (Deadwood)

Portsmouth Square à San Francisco, 1851

Samouraï, vers 1860

Bien sûr, le modèle n'est pas parfait. Cette main rouge me rend fou, mais sinon ça marche fantastiquement:

Fille sénéca iroquoise, 1908

Elle peut également coloriser des dessins en noir et blanc:

Détails techniques

Il s'agit d'un modèle d'apprentissage en profondeur. En particulier, j'ai combiné les approches suivantes:

Auto-attention GAN . La seule chose est que l' Unet pré-formé est utilisé comme générateur et je viens de le changer pour la normalisation spectrale et, en fait, le mécanisme d'auto-attention. Il s'agit d'une modification assez simple. Je vais vous dire que la différence est frappante par rapport à la version précédente de Wasserstein GAN, que j'ai essayé de faire fonctionner. J'ai aimé la théorie de Wasserstein GAN, mais en pratique cela ne fonctionne pas. Mais je suis juste tombé amoureux du réseau Self-Attention GAN.
Une structure d'apprentissage comme la croissance progressive d'un GAN (mais pas exactement la même). La différence est que le nombre de couches reste constant: je viens de changer la taille des données d'entrée et d'ajuster la vitesse d'apprentissage pour que les transitions entre les tailles soient réussies. Il semble qu'il produise le même résultat final, mais apprend plus vite, est plus stable et effectue mieux la généralisation.
Règle TTUR (règle de mise à jour à deux échelles de temps). Ici, c'est assez clair: juste une itération individuelle du générateur / discriminateur (critique) et une vitesse d'apprentissage du discriminateur plus élevée.
La fonction de perte du générateur se compose de deux parties: l'une d'entre elles est la fonction principale de la perte perceptuelle (ou perte de fonctionnalité) basée sur VGG16 - elle pousse simplement le modèle du générateur pour reproduire l'image d'entrée. La deuxième partie est l'estimation des pertes du discriminateur (critique). Pour les curieux: seule la fonction Perte Perte Perte ne suffit pas pour un bon résultat. Cela tend à encourager simplement un tas de marron / vert / bleu - vous savez, en trompant le test, quels sont les réseaux de neurones vraiment bons! Le point clé est que les GAN apprennent essentiellement la fonction de perte pour vous, ce qui est en fait un grand pas vers l'idéal que nous recherchons dans l'apprentissage automatique. Et bien sûr, les résultats s'amélioreront considérablement lorsque la machine elle-même apprend ce que vous avez précédemment encodé manuellement. Bien sûr, c'est le cas ici.

La beauté de ce modèle est qu'il est assez bon dans une variété de modifications d'image. Ce que vous voyez ci-dessus sont les résultats du modèle de coloration, mais ce n'est qu'un composant du pipeline que je souhaite développer avec le même modèle.

Ensuite, j'essaierai de perfectionner les anciennes images, et le prochain point à l'ordre du jour est un modèle pour améliorer la saturation et la richesse (defade). Maintenant, elle est aux premiers stades de la formation. Il s'agit essentiellement du même modèle, mais avec certains paramètres de contraste / luminosité comme une simulation de photos fanées et de photos prises avec un équipement ancien / médiocre. J'ai déjà reçu des résultats encourageants:

Détails du projet

Quelle est l'essence de ce projet? Je veux juste appliquer GAN pour que les vieilles photos soient très, très bonnes. Et plus important encore, cela rendra le projet utile . Et oui, je suis vraiment intéressé à travailler avec la vidéo, mais je dois d'abord comprendre comment prendre ce modèle sous contrôle de la consommation de mémoire (c'est une vraie bête). Ce serait bien si les modèles n'avaient pas appris de deux à trois jours sur 1080Ti (malheureusement, typique pour GAN). Bien que ce soit mon enfant et je vais activement mettre à jour et améliorer le code dans un avenir prévisible, mais je vais essayer de rendre le programme aussi convivial que possible, bien qu'il y aura probablement des difficultés avec.

Et je jure que je documenterai le code correctement ... un jour. Certes, je fais partie de ces gens qui croient au "code auto-documenté" (LOL).

Modèle d'auto-lancement

Le projet est construit sur la merveilleuse bibliothèque Fast.AI. Malheureusement, il s'agit d'une ancienne version, et il reste à mettre à jour vers une nouvelle (c'est définitivement à l'ordre du jour). Donc, les prérequis, en bref:

Ancienne bibliothèque Fast.AI. Après m'être enfoui dans le projet pendant deux mois, j'ai un peu raté ce qui lui est arrivé, car celui qui est désormais marqué comme «ancien» ne ressemble pas vraiment à celui que j'ai. Tout a changé au cours des deux derniers mois environ. Par conséquent, si rien ne fonctionne avec d'autres versions, je l'ai fourché ici . Encore une fois, la mise à jour vers la dernière version est à l'ordre du jour, je m'excuse à l'avance.
Toutes les dépendances Fast.AI : il y a des fichiers requirements.txt et environment.yml pratiques.
Pytorch 0.4.1 (spectral_norm est requis, vous avez donc besoin de la dernière version stable).
JupyterLab .
Tensorboard (c'est-à-dire l'installation de Tensorflow) et TensorboardX . Je pense que ce n'est pas strictement nécessaire, mais c'est beaucoup plus facile. Pour votre commodité, j'ai déjà fourni tous les crochets / rappels nécessaires dans Tensorboard! Il existe des exemples de leur utilisation. Il est à noter que par défaut, les images pendant le traitement sont enregistrées dans le Tensorboard toutes les 200 itérations, de sorte que vous aurez une vue constante et pratique de ce que fait le modèle.
ImageNet : un excellent ensemble de données pour la formation.
Carte graphique puissante . J'aimerais vraiment avoir plus de mémoire que 11 Go dans ma GeForce 1080Ti. Si vous avez quelque chose de plus faible, ce sera difficile. Unet et Critic sont absurdement excellents, mais plus ils sont gros, meilleurs sont les résultats.

Si vous souhaitez commencer le traitement d'images par vous-même sans entraîner le modèle, vous pouvez télécharger des poids prêts à l'emploi ici . Ouvrez ensuite ColorizationVisualization.ipynb dans JupyterLab. Assurez-vous qu'il y a une ligne avec un lien vers les poids:

colorizer_path = Path('/path/to/colorizer_gen_192.h5')

Ensuite, vous devez charger le modèle de coloriseur après l'initialisation de netG:

 load_model(netG, colorizer_path)

Ensuite, placez simplement les images dans le dossier / test_images /, d'où vous démarrez le programme. Vous pouvez visualiser les résultats dans le bloc-notes Jupyter avec les lignes suivantes:

 vis.plot_transformed_image("test_images/derp.jpg", netG, md.val_ds, tfms=x_tfms, sz=500)

J'économiserais une taille d'environ 500 pixels, plus ou moins, si vous exécutez le programme sur un GPU avec beaucoup de mémoire (par exemple, GeForce 1080Ti 11 Go). S'il y a moins de mémoire, vous devez réduire la taille des images ou essayer de fonctionner sur le CPU. En fait, j'ai essayé de faire ce dernier, mais pour une raison quelconque, le modèle a fonctionné très, absurdement lentement, et je n'ai pas trouvé le temps d'étudier le problème. Les connaisseurs ont recommandé de construire Pytorch à partir des sources, ce qui entraînerait une augmentation importante des performances. Hmm ... À ce moment, ce n'était pas avant cela.

Information additionnelle

La visualisation des images générées au fur et à mesure que vous apprenez peut également être effectuée dans Jupyter: il vous suffit de la définir sur true lors de la création d'une instance de ce crochet de visualisation:

GANVisualizationHook(TENSORBOARD_PATH, trainer, 'trainer', jupyter=True, visual_iters=100

Je préfère laisser faux et utiliser simplement Tensorboard. Croyez-moi, vous voulez aussi faire exactement cela. De plus, si vous le laissez fonctionner trop longtemps, Jupyter mangera beaucoup de mémoire avec de telles images.

Les poids du modèle sont également enregistrés automatiquement pendant les cycles d'entraînement GANTrainer. Par défaut, ils sont enregistrés toutes les 1000 itérations (c'est une opération coûteuse). Ils sont stockés dans le dossier racine que vous avez spécifié pour la formation et le nom correspond à save_base_name spécifié dans le programme de formation. Les poids sont stockés séparément pour chaque taille d'entraînement.

Je recommanderais de parcourir le code de haut en bas, en commençant par le Jupyter Notebook. Je prends ces notes simplement comme une interface pratique pour le prototypage et la visualisation; tout le reste ira aux fichiers .py dès que je leur trouverai une place. J'ai déjà des exemples de visualisation que vous pouvez facilement activer et voir: ouvrez simplement xVisualization dans le Notebook, les images de test incluses dans le projet y sont listées (elles sont dans test_images).

Si vous voyez les horaires GAN, c'est la chose la plus laide du projet, juste ma version de l'implémentation GAN à apprentissage progressif, adaptée au générateur Unet.

Les poids pré-formés pour le générateur de colorant sont également ici . Le projet DeFade est toujours en cours, je vais essayer de mettre de bons poids dans quelques jours.

Habituellement, pendant la formation, vous verrez les premiers bons résultats à mi-chemin, c'est-à-dire avec une taille de 192 pixels (si vous utilisez les exemples de formation fournis).

Je suis sûr que j'ai foiré quelque part, alors faites-le moi savoir si c'est le cas.

Problèmes connus

Il faut jouer un peu avec la taille de l'image pour obtenir le meilleur résultat. Le modèle souffre clairement d'un certain rapport d'aspect et rapport d'aspect lors de la génération d'images. Auparavant, c'était bien pire, mais la situation s'est considérablement améliorée avec l'augmentation de l'éclairage / du contraste et l'introduction de l'apprentissage progressif. Je veux éliminer complètement ce problème et me concentrer dessus, mais jusqu'à présent, ne désespérez pas si l'image semble trop saturée ou avec des problèmes étranges. Très probablement, tout deviendra normal après un petit redimensionnement. En règle générale, pour les images sursaturées, vous devez augmenter la taille.
En plus de ce qui précède: obtenir les meilleures images se résume vraiment à l' art de choisir les paramètres optimaux . Oui, les résultats sont sélectionnés manuellement. Je suis très satisfait de la qualité et le modèle fonctionne de manière assez fiable, mais pas parfaitement. Le projet est toujours en cours! Je pense que l'outil peut être utilisé comme un «artiste IA», mais il n'est pas encore prêt pour le grand public. Mais pas le temps.
Pour compliquer la situation: à l'heure actuelle, le modèle mange brutalement de la mémoire , donc sur ma carte 1080Ti, il s'avère qu'il traite des images avec un maximum de 500-600px. Je parie qu'il existe de nombreuses options d'optimisation ici, mais je ne l'ai pas encore fait.
J'ai ajouté un remplissage nul au générateur Unet pour tout ce qui ne correspond pas aux tailles attendues (c'est ainsi que je peux charger une image de taille arbitraire). C'était un hack très rapide, et cela conduit à des bordures droites et inférieures stupides à la sortie d'images de test de taille arbitraire. Je suis sûr qu'il existe un meilleur moyen, mais je ne l'ai pas encore trouvé.
Le mannequin aime les vêtements bleus. Je ne sais pas trop pourquoi; solution à la recherche!

Vous en voulez plus?

Je publierai de nouveaux résultats sur Twitter .

Ajout du traducteur.
De ce dernier sur Twitter:

Les représentants de la nationalité eux-mêmes à leur pirogue, 1880

( original )

La construction du métro de Londres, 1860

( original )

Les bidonvilles de Baltimore, 1938

Gym sur le Titanic, 1912

( original )

DeOldify: un programme pour colorier des images en noir et blanc