Hola a todos, trabajo como programador de investigación en el equipo de visión por computadora de Mail.ru Group. Para el Día de la Victoria de este año, decidimos hacer un proyecto para la restauración de fotografías militares . ¿Qué es la restauración de fotos? Consta de tres etapas:

• encontramos todos los defectos de imagen: roturas, rasguños, agujeros;
  
• pintar sobre los defectos encontrados en función de los valores de píxeles a su alrededor;
  
• colorear la imagen
  

En este artículo, pasaré por cada una de las etapas de restauración en detalle y le diré cómo y dónde tomamos datos, qué redes aprendimos, qué hicimos, qué rastrillos pisamos.

Búsqueda de defectos

Queremos encontrar todos los píxeles relacionados con defectos en la foto cargada. Primero, necesitamos entender qué tipo de fotografías de los años de guerra subirán las personas. Nos dirigimos a los organizadores del proyecto Regimiento Inmortal, quienes compartieron datos con nosotros. Después de analizarlos, notamos que las personas a menudo cargan retratos, individuales o grupales, que tienen una cantidad moderada o grande de defectos.

Luego fue necesario recolectar una muestra de entrenamiento. La muestra de entrenamiento para la tarea de segmentación es una imagen y una máscara en la que se marcan todos los defectos. La forma más fácil es dar fotos a los marcadores a los marcadores. Por supuesto, las personas son buenas para encontrar defectos, pero el problema es que el marcado es un proceso muy largo.



Puede tomar de una hora a un día hábil completo para marcar píxeles relacionados con defectos en una foto, por lo que es difícil recopilar una muestra de más de 100 fotos en unas pocas semanas. Por lo tanto, tratamos de complementar de alguna manera nuestros datos y escribimos defectos nosotros mismos: tomamos una foto limpia, le aplicamos defectos artificiales y obtuvimos una máscara que nos muestra qué partes particulares de la imagen estaban dañadas. La parte principal de nuestra muestra de entrenamiento fue 79 fotos marcadas manualmente, de las cuales 11 fueron transferidas a la muestra de prueba.

El enfoque más popular para el problema de segmentación: tome Unet con un codificador previamente entrenado y minimice la cantidad $$$Bce$( entropía cruzada binaria ) y $$$DICE$( Sørensen - Coeficiente de dados ).

¿Qué problemas surgen con este enfoque en el problema de la segmentación de defectos?

• Incluso si nos parece que hay muchos defectos en la foto, que está muy sucia y muy raída por el tiempo, el área ocupada por los defectos es aún mucho más pequeña que la parte no dañada de la imagen. Para resolver este problema, puede aumentar el peso de la clase positiva en $$$Bce$, y el peso óptimo será la relación entre el número de píxeles puros y el número de píxeles que pertenecen a los defectos.
  
• El segundo problema es que si usamos Unet fuera de la caja con un codificador pre-entrenado, por ejemplo Albunet-18, entonces perdemos mucha información posicional. La primera capa de Albunet-18 consiste en una convolución con un núcleo de 5 y una zancada igual a dos. Esto permite que la red funcione rápidamente. Sacrificamos el tiempo de la red en aras de una mejor localización de los defectos: eliminamos la agrupación máxima después de la primera capa, redujimos el paso a 1 y redujimos el núcleo de convolución a 3.
  
• Si trabajamos con imágenes pequeñas, por ejemplo, comprimiendo una imagen a 256 x 256 o 512 x 512, los pequeños defectos simplemente desaparecerán debido a la interpolación. Por lo tanto, debe trabajar con una imagen grande. Ahora en producción estamos segmentando defectos en fotografías de 1024 x 1024. Por lo tanto, era necesario entrenar la red neuronal en grandes cultivos de imágenes grandes. Y debido a esto, hay problemas con el pequeño tamaño del lote en una tarjeta de video.
  
• Durante el entrenamiento, tenemos alrededor de 20 imágenes colocadas en una tarjeta. Debido a esto, la estimación de la media y la varianza en las capas BatchNorm es inexacta. In-place BatchNorm nos ayuda a resolver este problema, que, en primer lugar, ahorra memoria y, en segundo lugar, tiene una versión de Synchronized BatchNorm, que sincroniza las estadísticas entre todas las tarjetas. Ahora consideramos el promedio y la varianza no por 20 imágenes en una tarjeta, sino por 80 imágenes de 4 tarjetas. Esto mejora la convergencia de la red.
  

Al final, aumentar de peso. $$$Bce$Al cambiar la arquitectura y usar In-place BatchNorm, comenzamos a buscar defectos en la foto. Pero a un precio económico, podría hacerlo incluso un poco mejor agregando Test Time Augmentation. Podemos ejecutar la red una vez en la imagen de entrada, luego reflejarla y ejecutar la red nuevamente, esto puede ayudarnos a encontrar pequeños defectos.



Como resultado, nuestra red convergió en cuatro GeForce 1080Ti en 18 horas. La inferencia toma 290 ms. Resulta que es lo suficientemente largo, pero es una tarifa por el hecho de que estamos buscando defectos pequeños. Validación $$$DICE$igual a 0,35 y $$$ROCAUC$- 0,93.

Restauración de fragmentos

Unet nos ayudó a resolver este problema nuevamente. A la entrada le dimos la imagen original y una máscara en la que marcamos espacios limpios con unidades y esos píxeles que queremos pintar con ceros. Recopilamos los datos de la siguiente manera: tomamos de Internet un gran conjunto de datos con imágenes, por ejemplo, OpenImagesV4, y defectos añadidos artificialmente que son similares en forma a los que se encuentran en la vida real. Y después de eso, entrenaron a la red para reparar las partes faltantes.

¿Cómo podemos modificar Unet para esta tarea?

Puede usar la convolución parcial en lugar de la convolución habitual. Su idea es que cuando colapsamos una región de una imagen con un núcleo, no tenemos en cuenta los valores de píxeles relacionados con los defectos. Esto ayuda a que la pintura sea más precisa. Un ejemplo de un artículo de NVIDIA . En la imagen central, usaron Unet con la convolución habitual, y a la derecha, con convolución parcial:



Entrenamos la red durante 5 días. El último día, congelamos BatchNorm, esto ayudó a que los bordes de la parte pintada de la imagen fueran menos visibles.

La red procesa una imagen de 512 x 512 en 50 ms. La validación PSNR es 26.4. Sin embargo, no se puede confiar incondicionalmente en esta tarea. Por lo tanto, primero ejecutamos varios modelos buenos en nuestros datos, anonimizamos los resultados y luego votamos por los que más nos gustaron. Entonces elegimos el modelo final.

Mencioné que agregamos artificialmente defectos para limpiar imágenes. Al entrenar, debe monitorear cuidadosamente el tamaño máximo de los defectos superpuestos, porque con defectos muy grandes que la red nunca ha visto en el proceso de aprendizaje, fantaseará y dará un resultado absolutamente inaplicable. Entonces, si necesita pintar sobre defectos grandes, también aplique defectos grandes durante el entrenamiento.

Aquí hay un ejemplo del algoritmo:

Colorear

Segmentamos los defectos y los pintamos, el tercer paso es la reconstrucción del color. Permíteme recordarte que entre las fotografías del "Regimiento Inmortal" hay muchos retratos individuales o grupales. Y queríamos que nuestra red funcionara bien con ellos. Decidimos hacer nuestra propia coloración, porque ninguno de los servicios que conocemos pinta retratos de forma rápida y adecuada.



GitHub tiene un repositorio popular para colorear fotos. En promedio, hace bien este trabajo, pero tiene varios problemas. Por ejemplo, le encanta pintar la ropa de azul. Por lo tanto, también lo rechazamos.

Entonces, decidimos hacer una red neuronal para la coloración. La idea más obvia: tomar una imagen en blanco y negro y predecir tres canales, rojo, verde y azul. Pero, en términos generales, podemos simplificar nuestro trabajo. Podemos trabajar no con la representación RGB del color, sino con la representación YCbCr. El componente Y es brillo (luma). La imagen en blanco y negro descargada es el canal Y, la reutilizaremos. Quedaba por predecir Cb y Cr: Cb es la diferencia en color azul y brillo, y Cr es la diferencia en color rojo y brillo.



¿Por qué elegimos la vista YCbCr? El ojo humano es más susceptible a los cambios de brillo que a los cambios de color. Por lo tanto, reutilizamos el componente Y (brillo), algo a lo que el ojo es inicialmente receptivo, y predecimos Cb y Cr, en el que podemos cometer un poco más de error, ya que las personas notan menos colores "falsos". Esta característica comenzó a utilizarse activamente en los albores de la televisión en color, cuando el ancho de banda del canal no era suficiente para transmitir todos los colores en su totalidad. La imagen se transfirió a YCbCr, se transfirió al componente Y sin cambios, y Cb y Cr se comprimieron dos veces.

Cómo ensamblar la línea base

Puede volver a tomar Unet con un codificador previamente entrenado y minimizar la pérdida de L1 entre el CbCr real y el previsto. Queremos colorear los retratos, por lo que, además de las fotos de OpenImages, necesitamos agregar fotos específicas para nuestra tarea.

¿Dónde puedo obtener fotografías en color de personas con uniforme militar? Hay personas en Internet que pintan fotografías antiguas como un pasatiempo o por encargo. Lo hacen con mucho cuidado, tratando de cumplir con todos los matices. Coloreando el uniforme, charreteras, medallas, recurren a materiales de archivo, por lo que se puede confiar en el resultado de su trabajo. En total, utilizamos 200 fotografías pintadas a mano. La segunda fuente de datos útil es el sitio del Ejército Rojo de Trabajadores y Campesinos . Uno de sus creadores fue fotografiado en casi todas las variantes posibles de un uniforme militar durante la Gran Guerra Patria.



En algunas fotografías, repitió las poses de personas de famosas fotografías de archivo. Es especialmente bueno que haya disparado sobre un fondo blanco, esto nos permitió aumentar muy bien los datos, agregando varios objetos naturales al fondo. También utilizamos retratos modernos de personas comunes, que los complementan con insignias y otros atributos de la ropa de guerra.

Entrenamos a AlbuNet-50: este es Unet, en el que se utiliza AlbuNet-50 como codificador. La red comenzó a dar resultados adecuados: la piel es rosada, los ojos son gris azulados, las correas de los hombros son amarillentas. Pero el problema es que ella pintó las imágenes con manchas. Esto se debe al hecho de que, desde el punto de vista del error L1, a veces es más rentable no hacer nada que intentar predecir algo de color.


Estamos comparando nuestro resultado con una foto de Ground Truth - coloración manual del artista bajo el apodo de Klimbim

¿Cómo resolver este problema? Necesitamos un discriminador: una red neuronal, a la que proporcionaremos imágenes a la entrada, y dirá cuán realista se ve esta imagen. A continuación, una fotografía está pintada a mano y la segunda por una red neuronal. ¿Cuál te parece?




Como discriminador, utilizamos el discriminador del artículo Self-Attention GAN . Esta es una pequeña red de convolución, en las últimas capas de las cuales está incorporada la llamada Auto-Atención. Le permite "prestar más atención" a los detalles de la imagen. También utilizamos la normalización espectral. La explicación exacta y la motivación se pueden encontrar en el artículo. Entrenamos una red con una combinación de pérdida L1 y el error devuelto por el discriminador. Ahora la red pinta mejor los detalles de la imagen y el fondo es más consistente. Otro ejemplo: a la izquierda es el resultado de la red entrenada solo con pérdida L1, a la derecha, con pérdida L1 y un error discriminador.



En cuatro Geforce 1080Ti, el entrenamiento tomó dos días. La red funcionó en 30 ms en la imagen de 512 x 512. La validación MSE fue de 34.4. Como en el problema de la pintura, no se puede confiar plenamente en las métricas. Por lo tanto, seleccionamos 6 modelos que tenían las mejores métricas para la validación y votamos ciegamente por el mejor modelo.

Después de implementar el modelo en producción, continuamos los experimentos y llegamos a la conclusión de que es mejor minimizar no la pérdida L1 por píxel, sino la pérdida perceptiva. Para calcularlo, debe ejecutar la predicción de red y la foto de origen a través de la red VGG-16, tomar los mapas de atributos en las capas inferiores y compararlos de acuerdo con MSE. Este enfoque pinta más áreas y ayuda a obtener una imagen más colorida.

Conclusiones y conclusiones

Unet es una modelo genial. En el primer problema de segmentación, encontramos un problema al entrenar y trabajar con imágenes de alta resolución, por lo que usamos In-Place BatchNorm. En la segunda tarea (Inpainting), en lugar de la convolución habitual, utilizamos Convolución parcial, esto ayudó a lograr mejores resultados. En el problema de colorización de Unet, agregamos una pequeña red discriminadora que multó al generador por una imagen de aspecto poco realista y usamos la pérdida de percepción.

La segunda conclusión es que los accesores son importantes. Y no solo en la etapa de marcar las imágenes antes del entrenamiento, sino también para validar el resultado final, porque en problemas de defectos de pintura o coloración, aún debe validar el resultado con la ayuda de una persona. Le damos al usuario tres fotos: la original con los defectos eliminados, coloreada con los defectos eliminados, y solo la foto coloreada en caso de que el algoritmo para encontrar y pintar defectos sea incorrecto.

Tomamos algunas fotos del proyecto Military Album y las procesamos con nuestras redes neuronales. Aquí están los resultados obtenidos:



Y aquí puede verlos en la resolución original y en cada etapa del procesamiento.