Salut tout le monde! Je suis ingénieur de recherche au sein de l'équipe de vision par ordinateur du groupe Mail.ru. Dans cet article, je vais raconter comment nous avons créé un projet de restauration de photos basé sur l'IA pour de vieilles photos militaires. Qu'est-ce que la «restauration photo»? Il se compose de trois étapes:

• on retrouve tous les défauts d'image: fractures, éraflures, trous;
  
• nous peignons les défauts découverts, sur la base des valeurs de pixels qui les entourent;
  
• on colorise l'image.
  

De plus, je décrirai chaque étape de la restauration de photos et vous dirai comment nous avons obtenu nos données, quels réseaux nous avons formés, ce que nous avons accompli et quelles erreurs nous avons commises.

Recherche de défauts

Nous voulons trouver tous les pixels liés aux défauts dans une photo téléchargée. Tout d'abord, nous devons déterminer quel type de photos les gens vont télécharger. Nous avons discuté avec les fondateurs du projet "Immortal Regiment", une organisation non commerciale stockant les anciennes photos de la Seconde Guerre mondiale, qui ont partagé leurs données avec nous. En l'analysant, nous avons remarqué que les gens téléchargent principalement des portraits individuels ou de groupe avec un nombre moyen à élevé de défauts.

Ensuite, nous avons dû rassembler un ensemble de formation. L'ensemble de formation pour une tâche de segmentation est une image et un masque où tous les défauts sont marqués. La façon la plus simple de le faire est de laisser les évaluateurs créer les masques de segmentation. Bien sûr, les gens savent très bien comment trouver les défauts, mais cela prendrait trop de temps.



Cela peut prendre une heure ou toute la journée de travail pour marquer les pixels défectueux sur une photo. Par conséquent, il n'est pas facile de collecter un ensemble de formation de plus de 100 images en quelques semaines. C'est pourquoi nous avons essayé d'augmenter nos données et de créer nos propres défauts: nous prenions une bonne photo, ajoutions des défauts en utilisant des marches aléatoires sur l'image et nous retrouvions avec un masque montrant les parties de l'image avec les défauts. Sans augmentations, nous avons 68 photos étiquetées manuellement dans l'ensemble de formation et 11 photos dans l'ensemble de validation.

L'approche de segmentation la plus populaire: prenez Unet avec un encodeur pré-formé et minimisez la somme de BCE ( entropie croisée binaire ) et DICE ( Sørensen - Coefficient de dés ).

Quels problèmes surviennent lorsque nous utilisons cette approche de segmentation pour notre tâche?

• Même s'il semble qu'il y ait des tonnes de défauts sur la photo, qu'elle soit très ancienne et minable, la zone avec des défauts est encore beaucoup plus petite que celle en bon état. Pour résoudre ce problème, nous pouvons augmenter le poids de classe positif dans BCE; un poids optimal serait le rapport des pixels propres aux pixels défectueux.
  
• Le deuxième problème est que si nous utilisons un Unet prêt à l'emploi avec un encodeur pré-formé (Albunet-18, par exemple), nous perdons beaucoup de données de position. La première couche d'Albunet-18 consiste en une convolution avec un noyau 5 et une foulée égale à deux. Il permet au filet de fonctionner rapidement. Nous avons échangé le temps de fonctionnement net pour avoir une meilleure localisation des défauts: nous avons supprimé le regroupement maximal après la première couche, diminué la foulée à 1 et diminué le noyau de convolution à 3.
  
• Si nous travaillons avec de petites images en les compressant, par exemple, en 256 x 256 ou 512 x 512 pixels, les petits défauts disparaîtront en raison de l'interpolation. Par conséquent, nous devons travailler avec des images plus grandes. Nous segmentons actuellement les défauts en photos 1024 x 1024 en production. C'est pourquoi nous avons dû former le filet aux cultures à grande image. Cependant, cela provoque des problèmes avec une petite taille de lot sur un seul GPU.
  
• Pendant la formation, nous pouvons adapter environ 20 images sur un GPU. Pour cette raison, nous nous retrouvons avec des valeurs moyennes et d'écart type inexactes dans les couches BatchNorm. Nous pouvons résoudre ce problème en utilisant BatchNorm sur place , qui, d'une part, économise de l'espace mémoire, et d'autre part, a une version BatchNorm synchronisée, qui synchronise les statistiques sur tous les GPU. Maintenant, nous calculons les valeurs moyenne et écart type non pas pour 20 images sur un seul GPU, mais pour 80 images à partir de 4 GPU. Cela améliore la convergence nette.
  

Enfin, en augmentant le poids de BCE, en changeant d'architecture et en utilisant BatchNorm sur place, nous avons amélioré la segmentation. Cependant, il ne coûterait pas trop cher de faire quelque chose d'encore mieux en ajoutant l'augmentation de la durée du test. Nous pouvons exécuter le net une fois sur une image d'entrée, puis la mettre en miroir et réexécuter le net pour trouver tous les petits défauts.



Le net converge en 18 heures sur quatre GeForce 1080Ti. L'inférence prend 290 ms. C'est assez long, mais c'est le prix de nos performances supérieures à celles par défaut. Validation DICE est égal à 0,35 et ROCAUC - 0,93.

Inpainting d'image

Même chose avec la tâche de segmentation que nous avons utilisée Unet. Pour faire de la peinture, nous téléchargions une image originale et un masque où nous avons marqué toute la zone propre avec des uns et avec des zéros - tous les pixels que nous voulons peindre. C'est ainsi que nous collections des données: pour toute photo d'un jeu de données d'images open-source, par exemple, OpenImagesV4, nous ajoutons les défauts similaires à ceux que nous voyons dans la vie réelle. Ensuite, nous avions formé le filet pour restaurer les pièces manquantes.

Comment pouvons-nous modifier Unet pour cette tâche?

Nous pouvons utiliser une convolution partielle au lieu d'une convolution originale. L'idée est que lorsque nous convolutons une zone avec un noyau, nous ne prenons pas en compte les valeurs des pixels défectueux. Cela rend la peinture plus précise. Nous vous montrons un exemple du récent article NVIDIA . Ils ont utilisé Unet avec une convolution bidimensionnelle par défaut dans l'image du milieu et une convolution partielle - dans l'image de droite.



Nous avons formé le filet pendant cinq jours. Le dernier jour, nous avons gelé BatchNorms pour rendre les bords de la partie peinte moins visibles.

Il faut 50 ms pour traiter une image 512 x 512. La validation PSNR est égale à 26,4. Cependant, vous ne pouvez pas totalement compter sur les métriques dans cette tâche. Pour choisir le meilleur modèle, nous avons exécuté plusieurs bons modèles sur des images d'évaluation, anonymisé les résultats, puis voté pour ceux que nous aimions le plus. C'est ainsi que nous avons choisi notre modèle final.

J'ai mentionné plus tôt que nous avons artificiellement ajouté des défauts aux images propres. Vous devez toujours suivre la taille maximale des défauts ajoutés pendant la formation; dans le cas où vous introduisez une image avec un très gros défaut dans le filet qui n'est jamais traité au stade de l'entraînement, le filet fonctionnera de manière sauvage et produira un résultat inapplicable. Par conséquent, si vous devez corriger de gros défauts, augmentez votre ensemble d'entraînement avec eux.

Voici l'exemple du fonctionnement de notre algorithme:

Colorisation

Nous avons segmenté les défauts et les avons peints; la troisième étape - la reconstruction des couleurs. Comme je l'ai déjà dit, il y a beaucoup de portraits individuels et de groupe parmi les photos du régiment immortel. Nous voulions que notre filet fonctionne bien avec eux. Nous avons décidé de proposer notre propre colorisation car aucun des services existants ne pouvait colorer les portraits rapidement et efficacement. Nous voulons que nos photos colorisées soient plus crédibles.



GitHub a un référentiel populaire pour la colorisation des photos. Il fait du bon travail mais a encore quelques problèmes. Par exemple, il a tendance à peindre les vêtements en bleu. C'est pourquoi nous l'avons également rejeté.

Nous avons donc décidé de créer un algorithme de colorisation de l'image. L'idée la plus évidente: prendre une image en noir et blanc et prévoir trois canaux: rouge, vert et bleu. Cependant, nous pouvons faciliter notre travail: travailler non pas avec la représentation des couleurs RVB, mais avec la représentation des couleurs YCbCr. La composante Y est la luminosité (luma). Une image téléchargée en noir et blanc est la chaîne Y, et nous allons la réutiliser. Maintenant, nous devons prédire Cb et Cr: Cb est la différence de couleur bleue et de luminosité et Cr - la différence de couleur rouge et de luminosité.



Pourquoi avons-nous choisi la représentation YCbCr? Un œil humain est plus sensible aux changements de luminosité qu'aux changements de couleur. C'est pourquoi nous réutilisons la composante Y (luminosité) avec laquelle l'œil humain est le plus sensible et prédisons le Cb et le Cr avec lesquels nous pourrions faire une erreur car nous ne pouvons pas très bien remarquer la fausseté des couleurs. Cette caractéristique spécifique a été largement utilisée à l'aube de la télévision couleur lorsque la capacité des canaux n'était pas suffisante pour transmettre toutes les couleurs. L'image a été transmise en YCbCr, inchangée à la composante Y, et Cb et Cr ont été réduits de moitié.

Comment créer une ligne de base

Nous pouvons prendre Unet avec un encodeur pré-formé et minimiser la perte L1 entre les valeurs CbCr existantes et celles prédites. Nous voulons colorer les portraits et, par conséquent, en plus des photos OpenImages, nous avons besoin de photos plus spécifiques à la tâche.

Où peut-on obtenir des photos colorisées de personnes vêtues d'un uniforme militaire? Il y a des gens sur Internet qui colorisent les vieilles photos comme passe-temps ou pour un prix. Ils le font très soigneusement, essayant d'être très précis. Lorsqu'ils colorent un uniforme, des épaulettes et des médailles, ils se réfèrent aux documents d'archives, de sorte que les résultats de leur travail sont dignes de confiance. Dans l'ensemble, nous avons utilisé 200 photos colorisées manuellement avec des personnes en uniforme militaire.

L'autre source de données utiles est le site Web de l' Armée rouge des travailleurs et des paysans . L'un de ses fondateurs s'est fait photographier dans presque tous les uniformes soviétiques de la Seconde Guerre mondiale disponibles.



Dans certaines photos, il a imité les poses de personnes à partir des célèbres photos d'archives. C'est une bonne chose que ses images aient un fond blanc: cela nous a permis d'augmenter très bien les données en ajoutant divers objets naturels en arrière-plan. Nous avons également utilisé des portraits réguliers, en les complétant d'insignes et d'autres attributs de guerre.

Nous avons formé AlbuNet-50 - c'est un Unet qui utilise ResNet-50 pré-formé comme encodeur. Le filet a commencé à donner des résultats satisfaisants: la peau était rose, les yeux - gris-bleu, les planches d'épaule - jaunâtres. Cependant, le problème est qu'il laisse certaines zones de la photo intactes. Cela était dû au fait que, selon l'erreur L1, trouver un tel optimum où il vaut mieux ne rien faire que d'essayer de prédire une certaine couleur.


Nous comparons notre résultat avec une photo de Ground Truth - une colorisation manuelle réalisée par Klimbim

Comment pouvons-nous résoudre ce problème? Nous avons besoin d'un discriminateur: un réseau neuronal qui recevrait une image et nous dirait si elle a l'air réaliste ou non. L'une des images ci-dessous est colorée manuellement et l'autre - par notre générateur, AlbuNet-50. Comment l'homme distingue-t-il les photos colorées manuellement et automatiquement? En regardant les détails. Pouvez-vous dire où se trouve la photo colorisée automatiquement par notre solution de base?




Nous utilisons le discriminateur du document Self-Attention GAN . C'est un petit filet de convolution avec soi-disant auto-attention intégré dans les couches supérieures. Cela nous permet de "faire plus attention" aux détails de l'image. Nous utilisons également la normalisation spectrale. Vous pouvez trouver plus d'informations dans le document susmentionné. Nous avons formé le filet avec une combinaison de perte L1 et de perte du discriminateur. Maintenant, le net colore mieux les détails de l'image et l'arrière-plan semble plus cohérent. Un autre exemple: à gauche, le travail par filet entraîné avec perte L1 uniquement; à droite - avec une combinaison de pertes de discriminateur L1.



Le processus de formation a pris deux jours sur quatre GeForce 1080Ti. Il faut 30 ms pour traiter une image 512 x 512. Validation MSE - 34.4. Tout comme avec l'inpainting, les métriques sur lesquelles vous ne voulez pas vous fier. C'est pourquoi nous avons choisi six modèles avec les meilleures métriques de validation et voté aveuglément pour le meilleur modèle.

Lorsque nous avons déjà créé un système de production et lancé un site Web, nous avons poursuivi nos expérimentations et conclu que nous devions mieux minimiser non pas la perte L1 par pixel, mais la perte perceptuelle. Pour le calculer, nous alimentons les prévisions du net et une photo au sol au VGG-16 net, prenons les cartes d'entités sur les couches inférieures et les comparons avec MSE. Cette approche peint plus de zones et donne des résultats plus colorés.

Récapitulatif

Unet est un modèle assez cool. Lors de la première tâche de segmentation, nous avons rencontré un problème pendant la formation et travaillons avec des images haute résolution et c'est pourquoi nous utilisons In-Place BatchNorm. Lors de notre deuxième tâche (Inpainting), nous avons utilisé la convolution partielle au lieu d'une convolution par défaut, et cela nous a permis d'obtenir de meilleurs résultats. Lorsque nous travaillions sur la colorisation, nous avons ajouté un petit filet de discrimination qui pénalisait le générateur pour des images irréalistes. Nous avons également utilisé une perte de perception.

Deuxième conclusion - les évaluateurs sont essentiels. Et pas seulement lors de la création des masques de segmentation mais aussi pour la validation du résultat final. En fin de compte, nous donnons à l'utilisateur trois photos: une image originale avec des défauts peints, une photo colorisée avec des défauts peints et une photo simplement colorisée au cas où l'algorithme de recherche de défauts et de peinture se serait trompé.

Nous avons pris quelques photos du projet War Album et les avons traitées sur ces neuronets. Voici les résultats que nous avons obtenus:



De plus, ici, vous pouvez regarder de plus près les images originales et toutes les étapes de traitement.