Hallo allerseits, ich arbeite als Forschungsprogrammierer im Computer Vision Team der Mail.ru Group. Für den diesjährigen Tag des Sieges haben wir beschlossen, ein Projekt zur Restaurierung von Militärfotografien durchzuführen . Was ist Fotorestaurierung? Es besteht aus drei Stufen:

• wir finden alle Bildfehler: Brüche, Kratzer, Löcher;
  
• Übermalen Sie die gefundenen Fehler anhand der Pixelwerte um sie herum.
  
• kolorieren Sie das Bild.
  

In diesem Artikel werde ich jede Phase der Wiederherstellung im Detail durchgehen und Ihnen erklären, wie und wo wir Daten aufgenommen haben, welche Netzwerke wir gelernt haben, was wir getan haben, auf welche Rechen wir getreten sind.

Fehlersuche

Wir möchten alle Pixel im Zusammenhang mit Fehlern im hochgeladenen Foto finden. Zunächst müssen wir verstehen, welche Art von Fotos der Kriegsjahre die Menschen hochladen werden. Wir wandten uns an die Organisatoren des Immortal Regiment-Projekts, die Daten mit uns teilten. Nach der Analyse haben wir festgestellt, dass Personen häufig Einzel- oder Gruppenporträts hochladen, die eine mäßige oder große Anzahl von Fehlern aufweisen.

Dann war es notwendig, eine Trainingsprobe zu sammeln. Das Trainingsmuster für die Segmentierungsaufgabe ist ein Bild und eine Maske, auf der alle Fehler markiert sind. Am einfachsten ist es, den Markern Fotos zu geben. Natürlich sind die Leute gut darin, Fehler zu finden, aber das Problem ist, dass das Markup ein sehr langer Prozess ist.



Das Markieren von Pixeln im Zusammenhang mit Fehlern in einem Foto kann von einer Stunde bis zu einem ganzen Arbeitstag dauern. Daher ist es schwierig, in wenigen Wochen eine Probe von mehr als 100 Fotos zu sammeln. Deshalb haben wir versucht, unsere Daten irgendwie zu ergänzen und selbst Fehler zu schreiben: Wir haben ein sauberes Foto gemacht, künstliche Fehler darauf angebracht und eine Maske bekommen, die uns zeigt, welche Teile des Bildes beschädigt wurden. Der Hauptteil unserer Trainingsprobe bestand aus 79 manuell markierten Fotos, von denen 11 auf die Testprobe übertragen wurden.

Der beliebteste Ansatz für das Segmentierungsproblem: Nehmen Sie Unet mit einem vorab trainierten Encoder und minimieren Sie die Menge $$$Bce$( binäre Kreuzentropie ) und $$$DICE$( Sørensen - Würfelkoeffizient ).

Welche Probleme treten bei diesem Ansatz beim Problem der Segmentierung von Fehlern auf?

• Auch wenn es uns so vorkommt, als ob das Foto viele Mängel aufweist, dass es sehr schmutzig und zeitlich sehr zerrissen ist, ist der von Fehlern eingenommene Bereich immer noch viel kleiner als der unbeschädigte Teil des Bildes. Um dieses Problem zu lösen, können Sie das Gewicht der positiven Klasse in erhöhen $$$Bce$und das optimale Gewicht ist das Verhältnis der Anzahl aller reinen Pixel zur Anzahl der zu den Defekten gehörenden Pixel.
  
• Das zweite Problem ist, dass wir viele Positionsinformationen verlieren, wenn wir Unet sofort mit einem vorab trainierten Encoder verwenden, beispielsweise Albunet-18. Die erste Schicht von Albunet-18 besteht aus einer Faltung mit einem Kern von 5 und einem Schritt von zwei. Dadurch kann das Netzwerk schnell arbeiten. Wir haben die Betriebszeit des Netzwerks für eine bessere Lokalisierung von Fehlern geopfert: Wir haben das maximale Pooling nach der ersten Schicht entfernt, den Schritt auf 1 reduziert und den Faltungskern auf 3 reduziert.
  
• Wenn wir beispielsweise mit kleinen Bildern arbeiten und ein Bild auf 256 x 256 oder 512 x 512 komprimieren, verschwinden kleine Fehler aufgrund der Interpolation einfach. Daher müssen Sie mit einem großen Bild arbeiten. Jetzt in der Produktion segmentieren wir Defekte in Fotografien von 1024 x 1024. Daher war es notwendig, das neuronale Netzwerk in einer großen Anzahl großer Bilder zu trainieren. Aus diesem Grund gibt es Probleme mit der geringen Größe des Stapels auf einer Grafikkarte.
  
• Während des Trainings haben wir ungefähr 20 Bilder auf einer Karte. Aus diesem Grund ist die Schätzung des Mittelwerts und der Varianz in den BatchNorm-Ebenen ungenau. In-Place-BatchNorm hilft uns bei der Lösung dieses Problems, das zum einen Speicher spart und zum anderen eine Version von Synchronized BatchNorm enthält, mit der Statistiken zwischen allen Karten synchronisiert werden. Jetzt betrachten wir den Durchschnitt und die Varianz nicht mit 20 Bildern auf einer Karte, sondern mit 80 Bildern von 4 Karten. Dies verbessert die Netzwerkkonvergenz.
  

Am Ende Gewichtszunahme $$$Bce$Durch Ändern der Architektur und Verwenden von In-Place BatchNorm haben wir begonnen, nach Fehlern im Foto zu suchen. Aber billig könnten Sie es noch ein bisschen besser machen, indem Sie Test Time Augmentation hinzufügen. Wir können das Netzwerk einmal im Eingabebild ausführen, es dann spiegeln und das Netzwerk erneut ausführen. Dies kann uns helfen, kleine Fehler zu finden.



Infolgedessen konvergierte unser Netzwerk in 18 Stunden auf vier GeForce 1080Ti. Inferenz dauert 290 ms. Es stellt sich lange genug heraus, aber dies ist der Preis dafür, dass wir gut nach kleinen Mängeln suchen. Validierung $$$DICE$gleich 0,35 und $$$ROCAUC$- 0,93.

Fragmentwiederherstellung

Unet hat uns geholfen, dieses Problem wieder zu lösen. Zur Eingabe gaben wir ihm das Originalbild und eine Maske, auf der wir saubere Räume mit Einheiten markieren, sowie die Pixel, die wir mit Nullen übermalen möchten. Wir haben die Daten wie folgt gesammelt: Wir haben aus dem Internet einen großen Datensatz mit Bildern, beispielsweise OpenImagesV4, und künstlich hinzugefügten Fehlern entnommen, deren Form denen des realen Lebens ähnelt. Danach trainierten sie das Netzwerk, um die fehlenden Teile zu reparieren.

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

Sie können die Teilfaltung anstelle der üblichen Faltung verwenden. Ihre Idee ist, dass wir beim Reduzieren eines Bildbereichs mit einem Kernel die Pixelfehler, die sich auf Fehler beziehen, nicht berücksichtigen. Dies hilft, das Bild genauer zu machen. Ein Beispiel aus einem NVIDIA-Artikel . Im zentralen Bild verwendeten sie Unet mit der üblichen Faltung und rechts - mit Teilfaltung:



Wir haben das Netzwerk 5 Tage lang trainiert. Am letzten Tag haben wir BatchNorm eingefroren. Dies hat dazu beigetragen, dass die Ränder des gemalten Teils des Bildes weniger auffällig sind.

Das Netzwerk verarbeitet ein Bild von 512 x 512 in 50 ms. Das Validierungs-PSNR beträgt 26,4. Metriken können in dieser Aufgabe jedoch nicht unbedingt vertraut werden. Aus diesem Grund haben wir zunächst mehrere gute Modelle für unsere Daten erstellt, die Ergebnisse anonymisiert und dann für diejenigen gestimmt, die uns besser gefallen haben. Also haben wir das endgültige Modell gewählt.

Ich erwähnte, dass wir sauberen Bildern künstlich Fehler hinzugefügt haben. Während des Trainings müssen Sie die maximale Größe der überlagerten Fehler sorgfältig überwachen, da das Netzwerk bei sehr großen Fehlern, die das Netzwerk im Lernprozess noch nie gesehen hat, wild phantasiert und ein absolut nicht anwendbares Ergebnis liefert. Wenn Sie also große Fehler übermalen müssen, wenden Sie während des Trainings auch große Fehler an.

Hier ist ein Beispiel für den Algorithmus:

Färben

Wir haben die Fehler segmentiert und bemalt, der dritte Schritt ist die Rekonstruktion der Farbe. Ich möchte Sie daran erinnern, dass es auf den Fotografien des "Unsterblichen Regiments" viele Einzel- oder Gruppenporträts gibt. Und wir wollten, dass unser Netzwerk gut mit ihnen zusammenarbeitet. Wir haben uns für eine eigene Kolorierung entschieden, da keiner der uns bekannten Dienste schnell und gut Porträts malt.



GitHub hat ein beliebtes Repository zum Ausmalen von Fotos. Im Durchschnitt macht er diesen Job gut, aber er hat mehrere Probleme. Zum Beispiel malt er gerne Kleidung in Blau. Deshalb haben wir es auch abgelehnt.

Also beschlossen wir, ein neuronales Netzwerk für die Kolorierung zu erstellen. Die naheliegendste Idee: Nehmen Sie ein Schwarzweißbild auf und sagen Sie drei Kanäle voraus: Rot, Grün und Blau. Aber im Allgemeinen können wir unsere Arbeit vereinfachen. Wir können nicht mit der RGB-Darstellung der Farbe arbeiten, sondern mit der YCbCr-Darstellung. Komponente Y ist die Helligkeit (Luma). Das heruntergeladene Schwarzweißbild ist der Y-Kanal, wir werden ihn wiederverwenden. Es blieb Cb und Cr vorherzusagen: Cb ist der Unterschied in blauer Farbe und Helligkeit, und Cr ist der Unterschied in roter Farbe und Helligkeit.



Warum haben wir die YCbCr-Ansicht gewählt? Das menschliche Auge ist anfälliger für Helligkeitsänderungen als für Farbänderungen. Daher verwenden wir die Y-Komponente (Helligkeit), für die das Auge anfangs gut empfänglich ist, wieder und sagen Cb und Cr voraus, bei denen wir etwas mehr Fehler machen können, da die Menschen weniger „falsche“ Farben bemerken. Diese Funktion wurde zu Beginn des Farbfernsehens aktiv genutzt, als die Kanalbandbreite nicht ausreichte, um alle Farben vollständig zu übertragen. Das Bild wurde auf YCbCr übertragen, unverändert auf die Y-Komponente übertragen und Cb und Cr wurden zweimal komprimiert.

So montieren Sie die Grundlinie

Sie können Unet erneut mit einem vorab trainierten Encoder verwenden und den L1-Verlust zwischen dem realen CbCr und dem vorhergesagten minimieren. Wir möchten Porträts einfärben, daher müssen wir zusätzlich zu den Fotos von OpenImages Fotos hinzufügen, die für unsere Aufgabe spezifisch sind.

Wo kann ich Farbfotos von Menschen in Militäruniform bekommen? Es gibt Leute im Internet, die alte Fotos als Hobby oder auf Bestellung malen. Sie tun dies sehr sorgfältig und versuchen, alle Nuancen vollständig einzuhalten. Sie färben die Uniform, Schulterklappen und Medaillen und wenden sich Archivmaterial zu, damit das Ergebnis ihrer Arbeit vertrauenswürdig ist. Insgesamt haben wir 200 handgemalte Fotos verwendet. Die zweite nützliche Datenquelle ist der Standort der Roten Armee der Arbeiter und Bauern . Einer seiner Schöpfer wurde während des Großen Vaterländischen Krieges in fast allen möglichen Varianten einer Militäruniform fotografiert.



Auf einigen Fotografien wiederholte er die Posen von Menschen aus berühmten Archivfotos. Es ist besonders gut, dass er auf einem weißen Hintergrund aufgenommen hat. Dadurch konnten wir die Daten sehr gut erweitern und dem Hintergrund verschiedene natürliche Objekte hinzufügen. Wir haben auch gewöhnliche moderne Porträts von Menschen verwendet, die durch Abzeichen und andere Attribute von Kriegskleidung ergänzt wurden.

Wir haben AlbuNet-50 trainiert - dies ist Unet, in dem AlbuNet-50 als Encoder verwendet wird. Das Netzwerk lieferte adäquate Ergebnisse: Die Haut ist rosa, die Augen sind grau-blau, die Schultergurte sind gelblich. Das Problem ist aber, dass sie die Bilder mit Flecken bemalt hat. Dies liegt an der Tatsache, dass es aus Sicht des L1-Fehlers manchmal rentabler ist, nichts zu tun, als zu versuchen, eine Farbe vorherzusagen.


Wir vergleichen unser Ergebnis mit einem Foto von Ground Truth - manuelle Kolorierung des Künstlers unter dem Spitznamen Klimbim

Wie kann man dieses Problem lösen? Wir brauchen einen Diskriminator: ein neuronales Netzwerk, an das wir Bilder für die Eingabe liefern und das sagt, wie realistisch dieses Bild aussieht. Unten ist ein Foto handgemalt und das zweite von einem neuronalen Netzwerk. Welches denkst du?




Als Diskriminator verwenden wir den Diskriminator aus dem Artikel Self-Attention GAN . Dies ist ein kleines Faltungsnetzwerk, in dessen letzten Schichten die sogenannte Selbstaufmerksamkeit eingebaut ist. Es ermöglicht Ihnen, mehr auf die Bilddetails zu achten. Wir verwenden auch die spektrale Normalisierung. Die genaue Erklärung und Motivation finden Sie im Artikel. Wir haben ein Netzwerk mit einer Kombination aus L1-Verlust und dem vom Diskriminator zurückgegebenen Fehler trainiert. Jetzt malt das Netzwerk die Details des Bildes besser und der Hintergrund ist konsistenter. Ein weiteres Beispiel: Links ist das Ergebnis des Netzwerks, das nur mit L1-Verlust trainiert wurde, rechts - mit L1-Verlust und einem Diskriminatorfehler.



Bei vier Geforce 1080Ti dauerte das Training zwei Tage. Das Netzwerk hat in 30 ms im Bild 512 x 512 gearbeitet. Die Validierungs-MSE betrug 34,4. Wie beim Inpainting-Problem können Metriken nicht vollständig als vertrauenswürdig eingestuft werden. Daher haben wir 6 Modelle ausgewählt, die die besten Metriken für die Validierung hatten, und blind für das beste Modell gestimmt.

Nachdem wir das Modell in der Produktion eingeführt hatten, setzten wir die Experimente fort und kamen zu dem Schluss, dass es besser ist, nicht den L1-Verlust pro Pixel, sondern den Wahrnehmungsverlust zu minimieren. Um dies zu berechnen, müssen Sie die Netzwerkvorhersage und das Quellfoto über das VGG-16-Netzwerk ausführen, die Attributkarten auf den unteren Ebenen verwenden und sie gemäß MSE vergleichen. Dieser Ansatz malt mehr Bereiche und hilft, ein farbenfroheres Bild zu erhalten.

Schlussfolgerungen und Schlussfolgerung

Unet ist ein cooles Modell. Beim ersten Segmentierungsproblem ist beim Training und Arbeiten mit hochauflösenden Bildern ein Problem aufgetreten. Daher verwenden wir In-Place BatchNorm. Bei der zweiten Aufgabe (Inpainting) haben wir anstelle der üblichen Faltung die Teilfaltung verwendet, um bessere Ergebnisse zu erzielen. Im Kolorierungsproblem für Unet haben wir ein kleines Diskriminatornetzwerk hinzugefügt, das den Generator für ein unrealistisch aussehendes Bild bestraft und Wahrnehmungsverlust verwendet.

Die zweite Schlussfolgerung ist, dass Accessoren wichtig sind. Und das nicht nur beim Markieren von Bildern vor dem Training, sondern auch zur Validierung des Endergebnisses, denn bei Problemen mit Lackfehlern oder Kolorierungen müssen Sie das Ergebnis noch mit Hilfe einer Person validieren. Wir geben dem Benutzer drei Fotos: das Original mit den entfernten Fehlern, das mit den entfernten Fehlern eingefärbt ist, und nur das eingefärbte Foto, falls der Algorithmus zum Auffinden und Malen von Fehlern falsch ist.

Wir haben einige Fotos des Military Album- Projekts gemacht und sie mit unseren neuronalen Netzen verarbeitet. Hier sind die erzielten Ergebnisse:



Und hier können Sie sie in der Originalauflösung und in jeder Phase der Verarbeitung sehen.