Salut tout le monde, je travaille en tant que programmeur de recherche dans l'équipe de vision par ordinateur du groupe Mail.ru. Pour le Jour de la Victoire de cette année, nous avons décidé de réaliser un projet de restauration de photographies militaires . Qu'est-ce que la restauration photo? Il se compose de trois étapes:

• on retrouve tous les défauts d'image: cassures, éraflures, trous;
  
• peindre les défauts trouvés en fonction des valeurs de pixels qui les entourent;
  
• colorise l'image.
  

Dans cet article, je passerai en revue chacune des étapes de la restauration en détail et vous dirai comment et où nous avons pris les données, quels réseaux nous avons appris, ce que nous avons fait, quels râteaux nous avons empruntés.

Recherche de défauts

Nous voulons trouver tous les pixels liés aux défauts dans la photo téléchargée. Tout d'abord, nous devons comprendre quel type de photographies des années de guerre les gens vont télécharger. Nous nous sommes tournés vers les organisateurs du projet Immortal Regiment, qui ont partagé des données avec nous. Après les avoir analysés, nous avons remarqué que les gens téléchargent souvent des portraits, seuls ou en groupe, qui présentent un nombre modéré ou important de défauts.

Il a ensuite fallu prélever un échantillon de formation. L'échantillon d'apprentissage pour la tâche de segmentation est une image et un masque sur lesquels tous les défauts sont marqués. Le moyen le plus simple est de donner des photos aux marqueurs aux marqueurs. Bien sûr, les gens sont capables de trouver des défauts, mais le problème est que le balisage est un processus très long.



Cela peut prendre entre une heure et une journée entière de travail pour marquer des pixels liés à des défauts sur une photo, il est donc difficile de collecter un échantillon de plus de 100 photos en quelques semaines. Par conséquent, nous avons essayé de compléter nos données en quelque sorte et avons écrit nous-mêmes les défauts: nous avons pris une photo propre, y avons appliqué des défauts artificiels et obtenu un masque nous indiquant quelles parties particulières de l'image étaient endommagées. La partie principale de notre échantillon de formation était de 79 photos marquées manuellement, dont 11 ont été transférées à l'échantillon de test.

L'approche la plus populaire pour le problème de segmentation: prenez Unet avec un encodeur pré-formé et minimisez la quantité $$$Bce$ ( entropie croisée binaire ) et $$$DICE$ ( Sørensen - Coefficient de dés ).

Quels problèmes se posent avec cette approche dans le problème de la segmentation des défauts?

• Même s'il nous semble qu'il y a beaucoup de défauts sur la photo, qu'elle est très sale et très déchirée par le temps, la zone occupée par les défauts est encore beaucoup plus petite que la partie intacte de l'image. Pour résoudre ce problème, vous pouvez augmenter le poids de la classe positive dans $$$Bce$ , et le poids optimal sera le rapport du nombre de pixels purs au nombre de pixels appartenant aux défauts.
  
• Le deuxième problème est que si nous utilisons Unet avec un encodeur pré-formé, par exemple Albunet-18, alors nous perdons beaucoup d'informations de position. La première couche d'Albunet-18 consiste en une convolution avec un noyau de 5 et une foulée égale à deux. Cela permet au réseau de fonctionner rapidement. Nous avons sacrifié le temps réseau pour une meilleure localisation des défauts: nous avons supprimé le pool max après la première couche, réduit la foulée à 1 et réduit le noyau de convolution à 3.
  
• Si nous travaillons avec de petites images, par exemple, en compressant une image en 256 x 256 ou 512 x 512, les petits défauts disparaîtront simplement en raison de l'interpolation. Par conséquent, vous devez travailler avec une grande image. Maintenant en production, nous segmentons les défauts dans les photographies de 1024 x 1024. Par conséquent, il était nécessaire de former le réseau neuronal en grande récolte d'images de grande taille. Et à cause de cela, il y a des problèmes avec la petite taille du lot sur une carte vidéo.
  
• Pendant la formation, nous avons environ 20 photos placées sur une carte. Pour cette raison, l'estimation de la moyenne et de la variance dans les couches BatchNorm est inexacte. BatchNorm sur place nous aide à résoudre ce problème, qui, d'une part, économise de la mémoire, et d'autre part, il a une version de BatchNorm synchronisé, qui synchronise les statistiques entre toutes les cartes. Maintenant, nous considérons la moyenne et la variance non pas par 20 images sur une carte, mais par 80 images de 4 cartes. Cela améliore la convergence du réseau.
  

En fin de compte, l'augmentation du poids $$$Bce$ En modifiant l'architecture et en utilisant BatchNorm sur place, nous avons commencé à rechercher des défauts sur la photo. Mais à moindre coût, vous pourriez faire encore un peu mieux en ajoutant une augmentation de la durée du test. Nous pouvons exécuter le réseau une fois dans l'image d'entrée, puis le mettre en miroir et exécuter à nouveau le réseau, cela peut nous aider à trouver de petits défauts.



En conséquence, notre réseau a convergé sur quatre GeForce 1080Ti en 18 heures. L'inférence prend 290 ms. Cela s'avère assez long, mais c'est le prix pour le fait que nous recherchons bien de petits défauts. Validation $$$DICE$ égal à 0,35, et $$$ROCAUC$ - 0,93.

Restauration de fragments

Unet nous a aidé à résoudre ce problème à nouveau. Nous lui avons donné l'image d'origine et un masque à l'entrée, sur lesquels nous marquons les espaces propres avec des unités, et les pixels que nous voulons peindre avec des zéros. Nous avons collecté les données comme suit: nous avons pris d'Internet un grand ensemble de données avec des images, par exemple, OpenImagesV4, et ajouté artificiellement des défauts de forme similaire à ceux trouvés dans la vie réelle. Et après cela, ils ont formé le réseau pour réparer les pièces manquantes.

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

Vous pouvez utiliser la convolution partielle au lieu de la convolution habituelle. Son idée est que lorsque nous réduisons une région d'une image avec un noyau, nous ne prenons pas en compte les valeurs de pixels liées aux défauts. Cela permet de rendre la peinture plus précise. Un exemple d'un article NVIDIA . Dans l'image centrale, ils ont utilisé Unet avec la convolution habituelle, et à droite - avec Convolution partielle:



Nous avons formé le réseau pendant 5 jours. Le dernier jour, nous avons gelé BatchNorm, ce qui a contribué à rendre les bords de la partie peinte de l'image moins visibles.

Le réseau traite une image de 512 x 512 en 50 ms. Le PSNR de validation est de 26,4. Cependant, les mesures ne peuvent pas être approuvées sans condition dans cette tâche. Par conséquent, nous avons d'abord exécuté plusieurs bons modèles sur nos données, anonymisé les résultats, puis voté pour ceux que nous aimions le plus. Nous avons donc choisi le modèle final.

J'ai mentionné que nous avions artificiellement ajouté des défauts pour nettoyer les images. Lors de la formation, vous devez surveiller attentivement la taille maximale des défauts superposés, car avec de très gros défauts que le réseau n'a jamais vus dans le processus d'apprentissage, il fantasmera énormément et donnera un résultat absolument inapplicable. Donc, si vous devez peindre de gros défauts, appliquez également de gros défauts pendant la formation.

Voici un exemple d'algorithme:

Coloration

Nous avons segmenté les défauts et les avons peints, la troisième étape est la reconstruction de la couleur. Permettez-moi de vous rappeler que parmi les photographies de l '"Immortal Regiment" il y a beaucoup de portraits simples ou de groupe. Et nous voulions que notre réseau fonctionne bien avec eux. Nous avons décidé de faire notre propre colorisation, car aucun des services que nous connaissons ne peint les portraits rapidement et bien.



GitHub a un référentiel populaire pour colorier des photos. En moyenne, il fait bien ce travail, mais il a plusieurs problèmes. Par exemple, il aime peindre des vêtements en bleu. Par conséquent, nous l'avons également rejeté.

Nous avons donc décidé de créer un réseau neuronal pour la colorisation. L'idée la plus évidente: prendre une image en noir et blanc et prévoir trois canaux, rouge, vert et bleu. Mais, d'une manière générale, nous pouvons simplifier notre travail. Nous pouvons travailler non pas avec la représentation RVB de la couleur, mais avec la représentation YCbCr. Le composant Y est la luminosité (luma). L'image en noir et blanc téléchargée est la chaîne Y, nous la réutiliserons. Il restait à prédire Cb et Cr: Cb est la différence de couleur bleue et de luminosité, et Cr est la différence de couleur rouge et de luminosité.



Pourquoi avons-nous choisi la vue YCbCr? L'œil humain est plus sensible aux changements de luminosité qu'aux changements de couleur. Par conséquent, nous réutilisons la composante Y (luminosité), à laquelle l'œil est initialement bien sensible, et prédisons le Cb et le Cr, dans lesquels nous pouvons faire un peu plus d'erreurs, car les gens remarquent moins de «fausses» couleurs. Cette fonctionnalité a commencé à être activement utilisée à l'aube de la télévision couleur, lorsque la bande passante du canal n'était pas suffisante pour transmettre toutes les couleurs en totalité. L'image a été transférée sur YCbCr, transférée sur le composant Y inchangée, et Cb et Cr ont été compressés deux fois.

Comment assembler la ligne de base

Vous pouvez à nouveau prendre Unet avec un encodeur pré-formé et minimiser la perte L1 entre le CbCr réel et le prédictif. Nous voulons colorer des portraits, donc en plus des photos d'OpenImages, nous devons ajouter des photos spécifiques à notre tâche.

Où puis-je obtenir des photographies en couleur de personnes en uniforme militaire? Il y a des gens sur Internet qui peignent de vieilles photographies comme passe-temps ou pour commander. Ils le font très soigneusement, essayant de se conformer pleinement à toutes les nuances. Colorant l'uniforme, les épaulettes, les médailles, ils se tournent vers des documents d'archives, de sorte que le résultat de leur travail peut être fait confiance. Au total, nous avons utilisé 200 photographies peintes à la main. La deuxième source de données utiles est le site de l' Armée rouge ouvrière et paysanne . L'un de ses créateurs a été photographié dans presque toutes les variantes possibles d'uniforme militaire pendant la Grande Guerre patriotique.



Dans certaines photographies, il a répété les poses de personnes à partir de photographies d'archives célèbres. C’est particulièrement bien qu’il ait tiré sur un fond blanc, cela nous a permis d’augmenter très bien les données, en ajoutant divers objets naturels à l’arrière-plan. Nous avons également utilisé des portraits modernes ordinaires de personnes, en les complétant d'insignes et d'autres attributs des vêtements de guerre.

Nous avons formé AlbuNet-50 - c'est Unet, dans lequel AlbuNet-50 est utilisé comme encodeur. Le réseau a commencé à donner des résultats adéquats: la peau est rose, les yeux sont gris-bleu, les bretelles sont jaunâtres. Mais le problème est qu'elle a peint les tableaux avec des taches. Cela est dû au fait que du point de vue de l'erreur L1, il est parfois plus rentable de ne rien faire que d'essayer de prédire une certaine couleur.


Nous comparons notre résultat avec une photo de Ground Truth - colorisation manuelle de l'artiste sous le pseudo Klimbim

Comment résoudre ce problème? Nous avons besoin d'un discriminateur: un réseau de neurones, auquel nous fournirons des images à l'entrée, et il dira à quel point cette image est réaliste. Ci-dessous, une photographie est peinte à la main, et la seconde par un réseau de neurones. Lequel pensez-vous?




En tant que discriminateur, nous utilisons le discriminateur de l'article Self-Attention GAN . Il s'agit d'un petit réseau convolutionnel, dans les dernières couches dont la soi-disant auto-attention est intégrée. Il vous permet de "faire plus attention" aux détails de l'image. Nous utilisons également la normalisation spectrale. L'explication et la motivation exactes se trouvent dans l'article. Nous avons formé un réseau avec une combinaison de perte L1 et de l'erreur renvoyée par le discriminateur. Maintenant, le réseau peint mieux les détails de l'image et l'arrière-plan est plus cohérent. Un autre exemple: à gauche est le résultat du réseau formé uniquement avec une perte L1, à droite - avec une perte L1 et une erreur de discrimination.



Sur quatre Geforce 1080Ti, la formation a duré deux jours. Le réseau a fonctionné en 30 ms dans l'image 512 x 512. Le MSE de validation était de 34,4. Comme dans le cas du problème de peinture, les mesures ne peuvent pas être entièrement approuvées. Par conséquent, nous avons sélectionné 6 modèles qui avaient les meilleures mesures pour la validation et avons voté aveuglément pour le meilleur modèle.

Après avoir déployé le modèle en production, nous avons poursuivi les expériences et sommes arrivés à la conclusion qu'il était préférable de minimiser non pas la perte L1 par pixel, mais la perte perceptuelle. Pour le calculer, vous devez exécuter la prédiction du réseau et la photo source via le réseau VGG-16, prendre les cartes d'attributs sur les couches inférieures et les comparer selon MSE. Cette approche peint plus de zones et aide à obtenir une image plus colorée.

Conclusions et conclusion

Unet est un modèle sympa. Dans le premier problème de segmentation, nous avons rencontré un problème de formation et de travail avec des images haute résolution, nous utilisons donc In-Place BatchNorm. Dans la deuxième tâche (Inpainting), au lieu de la convolution habituelle, nous avons utilisé la convolution partielle, cela a aidé à obtenir de meilleurs résultats. Dans le problème de colorisation pour Unet, nous avons ajouté un petit réseau de discriminateur qui a infligé une amende au générateur pour une image d'apparence irréaliste et utilisé une perte de perception.

La deuxième conclusion est que les accesseurs sont importants. Et pas seulement au stade du marquage des photos avant l'entraînement, mais aussi pour valider le résultat final, car en cas de problèmes de défauts de peinture ou de colorisation, il faut quand même valider le résultat avec l'aide d'une personne. Nous donnons à l'utilisateur trois photos: l'original avec les défauts supprimés, colorisé avec les défauts supprimés, et juste la photo colorisée au cas où l'algorithme de recherche et de peinture des défauts serait erroné.

Nous avons pris quelques photos du projet Album militaire et les avons traitées avec nos réseaux de neurones. Voici les résultats obtenus:



Et ici, vous pouvez les voir dans la résolution d'origine et à chaque étape du traitement.