Intelligence artificielle

Restauration de photos basée sur l'IA – Fedor Kitashov

Restauration de photos basée sur l'IA - Fedor Kitashov


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 une histoire sur la façon dont nous avons créé un Projet de restauration de photos basé sur l'IA pour les vieilles photos militaires. Qu'est-ce que la «restauration photo»? Il se compose de trois étapes:

  • on retrouve tous les défauts de l'image: fractures, éraflures, trous;
  • nous repeignons les défauts découverts, en fonction des valeurs de pixel les entourant;
  • nous colorisons l'image.

De plus, je vais décrire chaque étape de la restauration de photo et vous expliquer 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.

Nous voulons trouver tous les pixels liés aux défauts dans une photo téléchargée. Premièrement, nous devons déterminer le type de photos que les gens téléchargeront. Nous avons parlé aux 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. Après analyse, nous avons constaté que les gens téléchargent principalement des portraits individuels ou en groupe présentant un nombre de défauts modéré à élevé.

Ensuite, nous avons dû collecter un ensemble de formation. L'ensemble d'apprentissage 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 des 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 d’une photo. Par conséquent, il n’est pas facile de rassembler plus de 100 images d’entraînement 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, nous ajoutions des défauts à l’aide de marches aléatoires sur l’image et nous retrouvions un masque montrant les parties de l’image contenant les défauts. Sans augmentations, nous avons 68 photos étiquetées manuellement dans le kit de formation et 11 photos dans le kit de validation.

L’approche de segmentation la plus répandue consiste à prendre Unet avec un encodeur pré-entraîné et à minimiser la somme de BCE (entropie croisée binaire) et DICE (Coefficient de Sørensen – Dice).

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

  • Même s’il semble y avoir des tonnes de défauts sur la photo, qu’elle est très vieille et usée, la zone de défauts est toujours beaucoup plus petite que celle qui n’a pas été endommagée. Pour résoudre ce problème, nous pouvons augmenter le poids de classe positif dans BCE; un poids optimal serait le rapport entre les pixels propres et les pixels défectueux.
  • Le deuxième problème est que si nous utilisons un Unet prêt à l'emploi avec un encodeur pré-entraîné (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. Cela permet au réseau de fonctionner rapidement. Nous avons échangé le temps net d’opération pour obtenir une meilleure localisation des défauts: nous avons supprimé le pooling maximum après la première couche, réduit le pas à 1 et réduit le noyau de convolution à 3.
  • Si nous travaillons avec de petites images en les compressant, par exemple, à 256 x 256 ou 512 x 512 pixels, les petits défauts disparaîtront par interpolation. Par conséquent, nous devons travailler avec des images plus grandes. Nous segmentons actuellement les défauts des photos au format 1024 x 1024 en production. C’est pourquoi nous avons dû former le filet aux grandes cultures d’image. Cependant, cela pose 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 de déviation moyenne et standard inexactes dans les couches BatchNorm. Nous pouvons résoudre ce problème en utilisant BatchNorm sur placequi économise de l’espace mémoire et dispose d’une version Synchronized BatchNorm synchronisant les statistiques sur tous les GPU. Nous calculons maintenant les valeurs de déviation moyenne et standard 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 modifiant l’architecture et en utilisant BatchNorm sur place, nous avons amélioré la segmentation. Toutefois, il ne coûterait pas trop cher de faire encore mieux en ajoutant une augmentation de la durée du test. Nous pouvons exécuter le réseau une fois sur une image d'entrée, puis l'inverser et le réexécuter pour rechercher 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 à payer pour notre performance meilleure que celle par défaut. Validation DICE est égal à 0,35 et ROCAUC - 0,93.

Même chose avec la tâche de segmentation que nous avons utilisée Unet. Pour faire de la peinture, nous importions une image originale et un masque dans lesquels nous marquions toute la zone de nettoyage avec des uns et des zéros - tous les pixels que nous voulions peindre. Voici comment nous collections les données: pour toute photo d'un jeu de données d'image open source, par exemple, OpenImagesV4, nous ajoutons les défauts similaires à ceux que nous voyons dans la réalité. Nous avons ensuite formé le réseau pour restaurer les parties 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 convertissons une zone avec un noyau, nous ne tenons pas compte des valeurs de pixels défectueux. Cela rend l'inpainting plus précis. Nous vous montrons un exemple tiré du récent document NVIDIA. Ils ont utilisé Unet avec une convolution à 2 dimensions par défaut dans l'image du milieu et une convolution partielle - dans l'image à droite.

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

Il faut environ 50 ms pour traiter une image de 512 x 512. Le PSNR de validation est égal à 26,4. Cependant, vous ne pouvez pas totalement vous fier aux métriques dans cette tâche. Pour choisir le meilleur modèle, nous avons utilisé plusieurs bons modèles sur des images d'évaluation, anonymisé les résultats, puis voté pour ceux qui nous ont le plus plu. C’est ainsi que nous avons choisi notre modèle final.

J'ai déjà mentionné que nous avions ajouté artificiellement des défauts aux images nettes. 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 grand défaut sur le net, ce qui n’a jamais été traité au stade de la formation, le net se déchaînera et produira un résultat inapplicable. Par conséquent, si vous devez réparer des défauts importants, complétez votre jeu d’entraînement avec eux.

Voici l'exemple de notre algorithme:

Nous avons segmenté les défauts et les avons peints; la troisième étape - reconstruction de couleur. Comme je l'ai déjà dit, il y a beaucoup de portraits individuels et de groupe parmi les photos du Régiment Immortal. Nous voulions que notre réseau fonctionne bien avec eux. Nous avons décidé de créer 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 possède un référentiel populaire pour la colorisation de 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ée.

Nous avons donc décidé de créer un algorithme pour la colorisation des images. L'idée la plus évidente est de prendre une image en noir et blanc et de prédire trois canaux: le rouge, le vert et le bleu. Cependant, nous pouvons faciliter notre travail: ne travaillez 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 en noir et blanc téléchargée est le canal Y et nous allons la réutiliser. Nous devons maintenant prévoir Cb et Cr: Cb est la différence de couleur et de luminosité bleues 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 la raison pour laquelle nous réutilisons le composant Y (luminosité) auquel un œil humain est le plus sensible et nous prévoyons les Cb et les Cr avec lesquels nous pourrions nous tromper, car nous ne remarquons pas très bien la couleur faussée. Cette caractéristique spécifique était 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é.

Nous pouvons prendre Unet avec un encodeur pré-entraîné et minimiser la perte L1 entre les valeurs CbCr existantes et celles prédites. Nous souhaitons colorier des portraits. Par conséquent, outre les photos OpenImages, nous avons besoin de photos plus spécifiques à une tâche.

Où pouvons-nous obtenir des photos colorisées de personnes vêtues d'un uniforme militaire? Il y a des gens sur Internet qui colorient de vieilles photos en passe-temps ou pour un prix. Ils le font très soigneusement, en essayant d'être très précis. Lorsqu'ils colorient un uniforme, des épaulettes et des médailles, ils se réfèrent aux archives pour que les résultats de leur travail soient fiables. Au total, nous avons utilisé 200 images colorisées manuellement avec des personnes en uniforme militaire.

L’autre source de données utile est le site Web de l’Armée rouge des travailleurs et des paysans. L'un de ses fondateurs a fait prendre sa photo dans pratiquement tous les uniformes soviétiques de la Seconde Guerre mondiale disponibles.

Dans certaines images, il a imité les poses de personnes de la célèbre photo 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 le ResNet-50 pré-entraîné comme encodeur. Le filet a commencé à donner des résultats adéquats: la peau était rose, les yeux gris-bleu, les épaulettes jaunâtres. Cependant, le problème était qu’il laissait certaines zones intactes sur la photo. Cela est dû au fait que, selon l’erreur L1, nous trouvons cet optimum là où il vaut mieux ne rien faire que d’essayer de prédire une couleur.

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

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

Réponse: la photo de gauche est colorée manuellement, à droite - automatiquement.

Nous utilisons le discriminateur du papier Self-Attention GAN. C’est un petit réseau de convolution avec une soi-disant Self-Attention construite 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 l'article susmentionné. Nous avons entraîné le filet avec une combinaison de perte de L1 et de perte du discriminateur. Maintenant, le réseau colorise mieux les détails de l'image et l'arrière-plan est plus cohérent. Un autre exemple: à gauche, le travail effectué par le réseau est formé uniquement avec une perte de L1; à droite - avec une combinaison de pertes de discriminateur L1.

Le processus de formation a pris deux jours sur quatre GeForce 1080Ti. Il faut au réseau 30 ms pour traiter une image de 512 x 512. Validation MSE - 34.4. Comme pour l’inpainting, vous ne souhaitez pas utiliser les métriques. C’est pourquoi nous avons sélectionné six modèles avec les meilleurs paramètres de validation et avons voté à l’aveugle 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ériences et conclu qu'il valait mieux minimiser les pertes non pas par pixel, mais les pertes perceptuelles. Pour le calculer, nous transmettons les prédictions de filet et une photo de terrain au VGG-16, prenons les cartes de caractéristiques sur les couches inférieures et les comparons à MSE. Cette approche peint plus de zones et donne des résultats plus colorés.

Unet est un modèle plutôt cool. Lors de la première tâche de segmentation, nous avons rencontré un problème lors de la formation et travaillons avec des images haute résolution. C’est la raison pour laquelle nous utilisons BatchNorm sur place. Lors de notre deuxième tâche (Inpainting), nous avons utilisé la convolution partielle au lieu d’une convolution par défaut, ce qui nous a permis d’obtenir de meilleurs résultats. Lorsque nous travaillons sur la colorisation, nous avons ajouté un petit réseau discriminateur qui pénalise le générateur d’images irréalistes. Nous avons également utilisé une perte de perception.

Deuxième conclusion - les évaluateurs sont essentiels. Et non seulement lors de la création des masques de segmentation, mais également lors de la validation du résultat final. Au final, nous donnons trois photos à l’utilisateur: une image originale avec des défauts non peints, une photo colorisée avec des défauts non peints et une image tout simplement colorée au cas où l’algorithme de recherche de défaut et d’inpainting se tromperait.

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

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

Afficher plus

SupportIvy

SupportIvy.com : Un lieu pour partager le savoir et mieux comprendre le monde. Meilleure plate-forme de support gratuit pour vous, Documentation &Tutoriels par les experts.

Articles similaires

Laisser un commentaire

Votre adresse de messagerie ne sera pas publiée. Les champs obligatoires sont indiqués avec *

Bouton retour en haut de la page
Fermer