Hallo allerseits! Ich bin ein Forschungsingenieur im Computer Vision Team der Mail.ru Group. In diesem Artikel werde ich eine Geschichte darüber erzählen, wie wir ein AI-basiertes Fotorestaurierungsprojekt für alte Militärfotos erstellt haben. Was ist „Fotorestaurierung“? Es besteht aus drei Schritten:

• wir finden alle Bildfehler: Brüche, Kratzer, Löcher;
  
• Wir malen die entdeckten Fehler basierend auf den Pixelwerten um sie herum.
  
• Wir färben das Bild.
  

Außerdem beschreibe ich jeden Schritt der Fotowiederherstellung und erkläre Ihnen, wie wir unsere Daten erhalten haben, welche Netze wir trainiert haben, was wir erreicht haben und welche Fehler wir gemacht haben.

Auf der Suche nach Mängeln

Wir möchten alle Pixel in Bezug auf Fehler in einem hochgeladenen Foto finden. Zuerst müssen wir herausfinden, welche Art von Bildern die Leute hochladen werden. Wir sprachen mit den Gründern des Projekts "Immortal Regiment", einer nichtkommerziellen Organisation, die die alten Fotos des Zweiten Weltkriegs speichert und ihre Daten mit uns teilte. Bei der Analyse haben wir festgestellt, dass die meisten Personen Einzel- oder Gruppenporträts mit einer moderaten bis großen Anzahl von Fehlern hochladen.

Dann mussten wir ein Trainingsset sammeln. Der Trainingssatz für eine Segmentierungsaufgabe besteht aus einem Bild und einer Maske, in der alle Fehler markiert sind. Der einfachste Weg, dies zu tun, besteht darin, die Prüfer die Segmentierungsmasken erstellen zu lassen. Natürlich wissen die Leute sehr gut, wie man Fehler findet, aber das würde zu lange dauern.



Es kann eine Stunde oder den ganzen Arbeitstag dauern, bis die fehlerhaften Pixel in einem Foto markiert sind. Daher ist es nicht einfach, in wenigen Wochen einen Trainingssatz mit mehr als 100 Bildern zu sammeln. Aus diesem Grund haben wir versucht, unsere Daten zu erweitern und unsere eigenen Fehler zu erstellen: Wir haben ein gutes Foto gemacht, Fehler mithilfe von zufälligen Schritten auf dem Bild hinzugefügt und am Ende eine Maske erhalten, die die Bildteile mit den Fehlern zeigt. Ohne Erweiterungen haben wir 68 manuell beschriftete Fotos im Trainingssatz und 11 Fotos im Validierungssatz.

Der beliebteste Segmentierungsansatz: Nehmen Sie Unet mit einem vorab trainierten Encoder und minimieren Sie die Summe aus BCE ( binäre Kreuzentropie ) und DICE ( Sørensen - Würfelkoeffizient ).

Welche Probleme treten auf, wenn wir diesen Segmentierungsansatz für unsere Aufgabe verwenden?

• Selbst wenn es so aussieht, als ob das Foto Unmengen von Fehlern enthält, dass es sehr alt und schäbig ist, ist der Bereich mit Fehlern immer noch viel kleiner als der unbeschädigte. Um dieses Problem zu lösen, können wir das positive Klassengewicht in BCE erhöhen. Ein optimales Gewicht wäre das Verhältnis von sauberen zu fehlerhaften Pixeln.
  
• Das zweite Problem besteht darin, dass bei Verwendung eines sofort einsatzbereiten Unet mit vorab trainiertem Encoder (z. B. Albunet-18) viele Positionsdaten verloren gehen. Die erste Schicht von Albunet-18 besteht aus einer Faltung mit einem Kern 5 und einem Schritt, der zwei entspricht. Dadurch kann das Netz schnell arbeiten. Wir haben die Nettobetriebszeit abgewogen, um eine bessere Fehlerlokalisierung zu erzielen: Wir haben das maximale Pooling nach der ersten Schicht entfernt, den Schritt auf 1 verringert und den Faltungskern auf 3 verringert.
  
• Wenn wir mit kleinen Bildern arbeiten, indem wir sie beispielsweise auf 256 x 256 oder 512 x 512 Pixel komprimieren, verschwinden kleine Fehler aufgrund der Interpolation. Daher müssen wir mit größeren Bildern arbeiten. Derzeit segmentieren wir Fehler in Fotos in der Größe 1024 x 1024 in der Produktion. Deshalb mussten wir das Netz auf Big Image Crops trainieren. Dies führt jedoch zu Problemen mit einer kleinen Stapelgröße auf einer einzelnen GPU.
  
• Während des Trainings können wir ungefähr 20 Bilder auf eine GPU passen. Aus diesem Grund erhalten wir ungenaue Mittel- und Standardabweichungswerte in BatchNorm-Layern. Wir können dieses Problem mit In-Place BatchNorm lösen, das einerseits Speicherplatz spart und andererseits über eine synchronisierte BatchNorm-Version verfügt, die Statistiken über alle GPUs hinweg synchronisiert. Jetzt berechnen wir den Mittelwert und die Standardabweichung nicht für 20 Bilder auf einer einzelnen GPU, sondern für 80 Bilder von 4 GPUs. Dies verbessert die Nettokonvergenz.
  

Nachdem wir das BCE-Gewicht erhöht, die Architektur geändert und In-Place-BatchNorm verwendet haben, haben wir die Segmentierung verbessert. Es würde jedoch nicht zu viel kosten, etwas noch besser zu machen, indem Sie Test Time Augmentation hinzufügen. Wir können das Netz einmal auf einem Eingabebild ausführen, es dann spiegeln und das Netz erneut ausführen, um alle kleinen Fehler zu finden.



Das Netz konvergiert in 18 Stunden auf vier GeForce 1080Ti. Inferenz dauert 290 ms. Es ist ziemlich lang, aber das ist der Preis für unsere überdurchschnittliche Leistung. Validierungswürfel sind gleich 0,35 und ROCAUC - 0,93.

Bildinpainting

Gleiches gilt für die Segmentierungsaufgabe, die wir für Unet verwendet haben. Zum Malen haben wir ein Originalbild und eine Maske hochgeladen, in der wir den gesamten sauberen Bereich mit Einsen und Nullen markiert haben - alle Pixel, die wir malen möchten. So haben wir Daten gesammelt: Für jedes Foto aus einem Open-Source-Bilddatensatz, z. B. OpenImagesV4, fügen wir die Fehler hinzu, die denen im wirklichen Leben ähneln. Dann hatten wir das Netz trainiert, um die fehlenden Teile wiederherzustellen.

Wie können wir Unet für diese Aufgabe ändern?

Wir können anstelle einer ursprünglichen teilweise Faltung verwenden. Die Idee ist, dass wir beim Falten eines Bereichs mit einem Kernel die Fehlerpixelwerte nicht berücksichtigen. Dies macht das Inpainting präziser. Wir zeigen Ihnen ein Beispiel aus dem aktuellen NVIDIA-Papier . Sie verwendeten Unet mit einer standardmäßigen zweidimensionalen Faltung im mittleren Bild und einer teilweisen Faltung - im Bild rechts.



Wir haben das Netz fünf Tage lang trainiert. Am letzten Tag haben wir BatchNorms eingefroren, um die Ränder des lackierten Teils weniger sichtbar zu machen.

Die Verarbeitung eines 512 x 512-Bildes dauert netto 50 ms. Validierung PSNR entspricht 26,4. Sie können sich jedoch nicht vollständig auf die Metriken in dieser Aufgabe verlassen. Um das beste Modell auszuwählen, führen wir mehrere gute Modelle für Bewertungsbilder aus, anonymisieren die Ergebnisse und stimmen dann für diejenigen ab, die uns am besten gefallen haben. So haben wir unser endgültiges Modell ausgewählt.

Ich habe bereits erwähnt, dass wir den sauberen Bildern künstlich einige Fehler hinzugefügt haben. Sie sollten während des Trainings immer die maximale Größe der hinzugefügten Fehler verfolgen. In einem Fall, in dem Sie ein Bild mit einem sehr großen Fehler in das Netz einspeisen, das in der Trainingsphase noch nie behandelt wurde, läuft das Netz aus und führt zu einem nicht zutreffenden Ergebnis. Wenn Sie daher große Fehler beheben müssen, erweitern Sie Ihr Trainingsset um diese.

Hier ist das Beispiel, wie unser Algorithmus funktioniert:

Färbung

Wir haben die Mängel segmentiert und lackiert; der dritte Schritt - Farbrekonstruktion. Wie ich bereits sagte, gibt es viele Einzel- und Gruppenporträts unter den Fotos des Unsterblichen Regiments. Wir wollten, dass unser Netz gut mit ihnen zusammenarbeitet. Wir haben uns für eine eigene Farbgebung entschieden, da keiner der vorhandenen Dienste die Porträts schnell und effizient färben konnte. Wir möchten, dass unsere kolorierten Fotos glaubwürdiger sind.



GitHub verfügt über ein beliebtes Repository für die Farbfärbung von Fotos. Es macht einen guten Job, hat aber immer noch einige Probleme. Zum Beispiel neigt es dazu, Kleidung blau zu malen. Deshalb haben wir es auch abgelehnt.

Deshalb haben wir uns entschlossen, einen Algorithmus für die Bildfärbung zu erstellen. Die naheliegendste Idee: Nehmen Sie ein Schwarzweißbild auf und sagen Sie drei Kanäle voraus: Rot, Grün und Blau. Wir können unsere Arbeit jedoch erleichtern: Arbeiten Sie nicht mit der RGB-Farbdarstellung, sondern mit der YCbCr-Farbdarstellung. Die Y-Komponente ist die Helligkeit (Luma). Ein hochgeladenes Schwarzweißbild ist ein Y-Kanal, und wir werden es wiederverwenden. Jetzt müssen wir Cb und Cr vorhersagen: Cb ist der Unterschied zwischen blauer Farbe und Helligkeit und Cr - der Unterschied zwischen roter Farbe und Helligkeit.



Warum haben wir uns für die YCbCr-Darstellung entschieden? Ein menschliches Auge reagiert empfindlicher auf Helligkeitsänderungen als auf Farbänderungen. Aus diesem Grund verwenden wir die Y-Komponente (Helligkeit), für die ein menschliches Auge am empfindlichsten ist, wieder und sagen Cb und Cr voraus, mit denen wir möglicherweise einen Fehler machen, da wir Farbfehler nicht sehr gut bemerken können. Diese spezielle Eigenschaft wurde zu Beginn des Farbfernsehens häufig verwendet, als die Kanalkapazität nicht ausreichte, um alle Farben zu übertragen. Das Bild wurde in YCbCr unverändert zur Y-Komponente übertragen, und Cb und Cr wurden um die Hälfte reduziert.

So erstellen Sie eine Basislinie

Wir können Unet mit einem vorab trainierten Encoder verwenden und den L1-Verlust zwischen den vorhandenen und den vorhergesagten CbCr-Werten minimieren. Wir möchten Porträts ausmalen und benötigen daher neben OpenImages-Fotos mehr aufgabenspezifische Fotos.

Wo können wir kolorierte Fotos von Menschen in Militäruniform bekommen? Es gibt Leute im Internet, die alte Fotos als Hobby oder zu einem Preis kolorieren. Sie tun es sehr sorgfältig und versuchen sehr genau zu sein. Wenn sie eine Uniform, Schulterklappen und Medaillen ausmalen, beziehen sie sich auf die Archivmaterialien, sodass die Ergebnisse ihrer Arbeit vertrauenswürdig sind. Insgesamt haben wir 200 manuell kolorierte Bilder mit Personen in Militäruniform verwendet.

Die andere nützliche Datenquelle ist die Website der Roten Armee der Arbeiter und Bauern . Einer seiner Gründer ließ sich in nahezu jeder sowjetischen Uniform des Zweiten Weltkriegs fotografieren.



In einigen Bildern ahmte er die Posen von Menschen aus den berühmten Archivfotos nach. Es ist gut, dass seine Bilder einen weißen Hintergrund haben: Dadurch konnten wir die Daten sehr gut erweitern, indem wir verschiedene natürliche Objekte in den Hintergrund einfügten. Wir haben auch einige regelmäßige Porträts verwendet, die durch Abzeichen und andere Attribute aus der Kriegszeit ergänzt wurden.

Wir haben AlbuNet-50 trainiert - es ist ein Unet, das vorab geschultes ResNet-50 als Encoder verwendet. Das Netz lieferte adäquate Ergebnisse: Die Haut war rosa, die Augen - grau-blau, die Schulterklappen - gelblich. Das Problem war jedoch, dass einige Bereiche auf dem Foto unberührt bleiben. Dies wurde durch die Tatsache verursacht, dass laut Fehler L1 ein solches Optimum gefunden wird, bei dem es besser ist, nichts zu tun, als zu versuchen, eine Farbe vorherzusagen.


Wir vergleichen unser Ergebnis mit einem Foto von Ground Truth - einer manuellen Kolorierung von Klimbim

Wie können wir dieses Problem lösen? Wir brauchen einen Diskriminator: ein neuronales Netzwerk, das ein Bild empfängt und uns sagt, ob es realistisch aussieht oder nicht. Eines der Bilder unten ist manuell eingefärbt und das andere - von unserem Generator AlbuNet-50. Wie unterscheidet der Mensch manuell und automatisch farbige Fotos? Durch das Betrachten von Details. Können Sie anhand unserer Basislösung erkennen, wo sich das automatisch eingefärbte Foto befindet?




Wir verwenden den Diskriminator aus dem Self-Attention GAN- Papier. Es ist ein kleines Faltungsnetz mit sogenannter Selbstaufmerksamkeit in den obersten Schichten. Dadurch können wir den Bilddetails "mehr Aufmerksamkeit schenken". Wir verwenden auch die spektrale Normalisierung. Weitere Informationen finden Sie im oben genannten Artikel. Wir haben das Netz mit einer Kombination aus L1-Verlust und Verlust durch den Diskriminator trainiert. Jetzt färbt das Netz die Bilddetails besser ein und der Hintergrund sieht konsistenter aus. Ein weiteres Beispiel: Auf der linken Seite ist die Arbeit von net nur mit L1-Verlust trainiert; rechts - mit einer Kombination von L1-Diskriminatorverlusten.



Der Trainingsprozess mit vier GeForce 1080Ti dauerte zwei Tage. Die Verarbeitung eines 512 x 512-Bildes dauert netto 30 ms. Validierung MSE - 34.4. Genau wie beim Inpainting möchten Sie sich nicht auf Metriken verlassen. Aus diesem Grund haben wir sechs Modelle mit den besten Validierungsmetriken ausgewählt und blind für das beste Modell gestimmt.

Als wir bereits ein Produktionssystem erstellt und eine Website gestartet haben, haben wir weiter experimentiert und sind zu dem Schluss gekommen, dass wir nicht den L1-Verlust pro Pixel, sondern den Wahrnehmungsverlust besser minimieren können. Um dies zu berechnen, geben wir die Netzvorhersagen und ein Boden-Wahrheitsfoto an das VGG-16-Netz weiter, nehmen die Feature-Maps auf den unteren Ebenen und vergleichen sie mit MSE. Dieser Ansatz malt mehr Bereiche und liefert farbenfrohere Ergebnisse.

Rückblick

Unet ist ein ziemlich cooles Modell. Bei der ersten Segmentierungsaufgabe hatten wir während des Trainings ein Problem und arbeiteten mit hochauflösenden Bildern. Deshalb verwenden wir In-Place BatchNorm. Bei unserer zweiten Aufgabe (Inpainting) haben wir Partial Convolution anstelle einer Standardaufgabe verwendet, um bessere Ergebnisse zu erzielen. Bei der Arbeit an der Kolorierung haben wir ein kleines Diskriminatornetz hinzugefügt, das den Generator für unrealistische Bilder bestraft. Wir haben auch einen Wahrnehmungsverlust verwendet.

Zweite Schlussfolgerung - Gutachter sind unerlässlich. Und das nicht nur während der Erstellung von Segmentierungsmasken, sondern auch für die endgültige Ergebnisvalidierung. Am Ende geben wir dem Benutzer drei Fotos: ein Originalbild mit unlackierten Fehlern, ein koloriertes Foto mit unlackierten Fehlern und ein einfach koloriertes Foto für den Fall, dass der Algorithmus für die Fehlersuche und -lackierung falsch war.

Wir haben einige Bilder aus dem War Album-Projekt aufgenommen und sie über diese Neuronets verarbeitet. Hier sind die Ergebnisse, die wir erhalten haben:



Darüber hinaus können Sie hier die Originalbilder und alle Verarbeitungsstufen genauer betrachten.