Hola a todos! Soy ingeniero de investigación en el equipo de visión por computadora del Grupo Mail.ru. En este artículo, voy a contar una historia de cómo hemos creado un proyecto de restauración de fotos basado en IA para viejas fotos militares. ¿Qué es la "restauración de fotos"? Consta de tres pasos:

• encontramos todos los defectos de la imagen: fracturas, rasguños, agujeros;
  
• pintamos los defectos descubiertos, en función de los valores de píxeles a su alrededor;
  
• Coloramos la imagen.
  

Además, describiré cada paso de la restauración de fotos y le diré cómo obtuvimos nuestros datos, qué redes capacitamos, qué logramos y qué errores cometimos.

Buscando defectos

Queremos encontrar todos los píxeles relacionados con defectos en una foto cargada. Primero, tenemos que averiguar qué tipo de imágenes subirán las personas. Hablamos con los fundadores del proyecto "Regimiento inmortal", una organización no comercial que almacena las fotos heredadas de la Segunda Guerra Mundial, quienes compartieron sus datos con nosotros. Al analizarlo, notamos que las personas cargan principalmente retratos individuales o grupales con una cantidad moderada a gran cantidad de defectos.

Luego tuvimos que recoger un conjunto de entrenamiento. El conjunto de entrenamiento para una tarea de segmentación es una imagen y una máscara donde se marcan todos los defectos. La forma más fácil de hacerlo es dejar que los asesores creen las máscaras de segmentación. Por supuesto, la gente sabe muy bien cómo encontrar defectos, pero eso llevaría demasiado tiempo.



Puede tomar una hora o todo el día laboral marcar los píxeles defectuosos en una foto. Por lo tanto, no es fácil recopilar un conjunto de entrenamiento de más de 100 imágenes en unas pocas semanas. Es por eso que tratamos de aumentar nuestros datos y crear nuestros propios defectos: tomamos una buena foto, agregamos defectos usando recorridos aleatorios en la imagen y terminamos con una máscara que muestra las partes de la imagen con los defectos. Sin aumentos, tenemos 68 fotos etiquetadas manualmente en el conjunto de entrenamiento y 11 fotos en el conjunto de validación.

El enfoque de segmentación más popular: tome Unet con codificador pre-entrenado y minimice la suma de BCE ( entropía cruzada binaria ) y DICE ( coeficiente de Sørensen - Dice ).

¿Qué problemas surgen cuando usamos este enfoque de segmentación para nuestra tarea?

• Incluso si parece que hay toneladas de defectos en la foto, que es muy vieja y en mal estado, el área con defectos es mucho más pequeña que la que no está dañada. Para resolver este problema, podemos aumentar el peso positivo de la clase en BCE; un peso óptimo sería la proporción de píxeles limpios a los defectuosos.
  
• El segundo problema es que si usamos un Unet listo para usar con codificador pre-entrenado (Albunet-18, por ejemplo), perdemos muchos datos posicionales. La primera capa de Albunet-18 consiste en una convolución con un núcleo 5 y una zancada que equivale a dos. Permite que la red funcione rápido. Cambiamos el tiempo de operación neto para tener una mejor localización de defectos: eliminamos la agrupación máxima después de la primera capa, disminuimos el paso a 1 y disminuimos el núcleo de convolución a 3.
  
• Si trabajamos con imágenes pequeñas comprimiéndolas, por ejemplo, a 256 x 256 o 512 x 512 píxeles, los pequeños defectos desaparecerán debido a la interpolación. Por lo tanto, necesitamos trabajar con imágenes más grandes. Actualmente estamos segmentando defectos en fotos de 1024 x 1024 en producción. Es por eso que tuvimos que entrenar la red en cultivos de gran imagen. Sin embargo, esto causa problemas con un tamaño de lote pequeño en una sola GPU.
  
• Durante el entrenamiento, podemos colocar unas 20 imágenes en una GPU. Debido a eso, terminamos con valores de desviación estándar y media inexactos en las capas BatchNorm. Podemos resolver este problema utilizando In-place BatchNorm , que, por un lado, ahorra espacio en la memoria y, por otro lado, tiene una versión BatchNorm sincronizada, que sincroniza las estadísticas en todas las GPU. Ahora calculamos los valores de la media y la desviación estándar no para 20 imágenes en una sola GPU, sino para 80 imágenes de 4 GPU. Esto mejora la convergencia neta.
  

Finalmente, al aumentar el peso de BCE, cambiar la arquitectura y usar In-place BatchNorm, mejoramos la segmentación. Sin embargo, no costaría demasiado hacer algo aún mejor agregando Test Time Augmentation. Podemos ejecutar la red una vez en una imagen de entrada, luego reflejarla y volver a ejecutar la red para encontrar todos los pequeños defectos.



La red converge en 18 horas en cuatro GeForce 1080Ti. La inferencia toma 290 ms. Es bastante largo, pero ese es el precio de nuestro rendimiento mejor que el predeterminado. Validación DICE es igual a 0,35 y ROCAUC - 0,93.

Inpainting de imagen

Lo mismo con la tarea de segmentación que usamos Unet. Para pintar, cargaríamos una imagen original y una máscara donde marcamos toda el área limpia con unos y con ceros, todos los píxeles que queremos pintar. Así es como estábamos recopilando datos: para cualquier foto de un conjunto de datos de imagen de código abierto, por ejemplo, OpenImagesV4, agregamos los defectos similares a los que vemos en la vida real. Luego habíamos entrenado la red para restaurar las partes faltantes.

¿Cómo podemos modificar Unet para esta tarea?

Podemos usar convolución parcial en lugar de una original. La idea es que cuando envolvemos un área con algo de kernel, no tomamos en cuenta los valores de píxeles defectuosos. Esto hace que la pintura sea más precisa. Le mostramos un ejemplo del reciente documento de NVIDIA . Usaron Unet con una convolución bidimensional predeterminada en la imagen del medio y una convolución parcial, en la imagen de la derecha.



Entrenamos la red durante cinco días. El último día, congelamos BatchNorms para que los bordes de la parte pintada sean menos visibles.

La red tarda 50 ms en procesar una imagen de 512 x 512. Validación PSNR es igual a 26.4. Sin embargo, no puede confiar totalmente en las métricas en esta tarea. Para elegir el mejor modelo, ejecutamos varios buenos modelos en imágenes de valoración, anonimizamos los resultados y luego votamos por los que más nos gustaron. Así es como elegimos nuestro modelo final.

Mencioné anteriormente que agregamos artificialmente algunos defectos a las imágenes limpias. Siempre debe realizar un seguimiento del tamaño máximo de defectos añadidos durante el entrenamiento; en el caso de que alimente una imagen con un defecto muy grande en la red que nunca se trata en la etapa de entrenamiento, la red se volverá loca y producirá un resultado inaplicable. Por lo tanto, si necesita reparar defectos grandes, aumente su conjunto de entrenamiento con ellos.

Aquí está el ejemplo de cómo funciona nuestro algoritmo:

Coloración

Segmentamos los defectos y los pintamos; El tercer paso: la reconstrucción del color. Como dije antes, hay muchos retratos individuales y grupales entre las fotos del Regimiento Inmortal. Queríamos que nuestra red funcionara bien con ellos. Decidimos crear nuestra propia coloración ya que ninguno de los servicios existentes podía colorear los retratos de forma rápida y eficiente. Queremos que nuestras fotos coloreadas sean más creíbles.



GitHub tiene un repositorio popular para la coloración de fotos. Hace un buen trabajo pero aún tiene algunos problemas. Por ejemplo, tiende a pintar la ropa de azul. Por eso lo rechazamos también.

Entonces, decidimos crear un algoritmo para la coloración de la imagen. La idea más obvia: tomar una imagen en blanco y negro y predecir tres canales: rojo, verde y azul. Sin embargo, podemos facilitar nuestro trabajo: no trabaje con la representación de color RGB, sino con la representación de color YCbCr. El componente Y es brillo (luma). Una imagen cargada en blanco y negro es el canal Y, y la vamos a reutilizar. Ahora tenemos que predecir Cb y Cr: Cb es la diferencia de color azul y brillo y Cr, la diferencia de color rojo y brillo.



¿Por qué elegimos la representación de YCbCr? Un ojo humano es más sensible a los cambios de brillo que a los cambios de color. Es por eso que reutilizamos el componente Y (brillo) al cual el ojo humano es más sensible y predecimos Cb y Cr con los que podríamos cometer un error, ya que no podemos notar muy bien la falsedad del color. Esta característica específica fue ampliamente utilizada en los albores de la televisión en color cuando la capacidad del canal no era suficiente para transmitir todos los colores. La imagen se transmitió en YCbCr, sin cambios al componente Y, y Cb y Cr se redujeron a la mitad.

Cómo crear una línea base

Podemos tomar Unet con un codificador previamente entrenado y minimizar la pérdida de L1 entre los valores de CbCr existentes y los pronosticados. Queremos colorear los retratos y, por lo tanto, además de las fotos de OpenImages, necesitamos más fotos específicas de la tarea.

¿Dónde podemos obtener fotos coloreadas de personas vestidas con un uniforme militar? Hay personas en Internet que colorean fotos antiguas como un pasatiempo o por un precio. Lo hacen con mucho cuidado, tratando de ser muy precisos. Cuando colorean un uniforme, hombreras y medallas, se refieren a los materiales de archivo, por lo que los resultados de su trabajo son confiables. Con todo, utilizamos 200 imágenes coloreadas manualmente con personas con uniforme militar.

La otra fuente de datos útil es el sitio web del Ejército Rojo de los Trabajadores y Campesinos . A uno de sus fundadores se le tomó una foto en casi todos los uniformes soviéticos de la Segunda Guerra Mundial disponibles.



En algunas fotos, imitaba las poses de personas de las famosas fotos de archivo. Es bueno que sus imágenes tengan fondo blanco: nos permitió aumentar muy bien los datos al agregar varios objetos naturales en el fondo. También utilizamos algunos retratos regulares, completándolos con insignias y otros atributos de tiempos de guerra.

Entrenamos a AlbuNet-50: es una Unet que utiliza ResNet-50 previamente entrenado como codificador. La red comenzó a dar resultados adecuados: la piel era rosada, los ojos - gris azulado, las hombreras - amarillentas. Sin embargo, el problema fue que deja algunas áreas en la foto intactas. Esto fue causado por el hecho de que, de acuerdo con el error, L1 encuentra un valor óptimo en el que es mejor no hacer nada que intentar predecir algo de color.


Estamos comparando nuestro resultado con una foto de Ground Truth, una coloración manual realizada por Klimbim

¿Cómo podemos resolver este problema? Necesitamos un discriminador: una red neuronal que reciba una imagen y nos diga si se ve realista o no. Una de las imágenes a continuación está coloreada manualmente y la otra, por nuestro generador, AlbuNet-50. ¿Cómo distingue el ser humano las fotos coloreadas manual y automáticamente? Al mirar los detalles. ¿Puedes decir dónde está la foto coloreada automáticamente por nuestra solución de referencia?




Usamos el discriminador del documento GAN de auto atención . Es una pequeña red de convolución con la llamada Auto-Atención construida en las capas superiores. Nos permite "prestar más atención" a los detalles de la imagen. También utilizamos la normalización espectral. Puede encontrar más información en el documento mencionado anteriormente. Hemos entrenado la red con una combinación de pérdida L1 y una pérdida del discriminador. Ahora la red colorea mejor los detalles de la imagen y el fondo se ve más consistente. Un ejemplo más: a la izquierda está el trabajo por red entrenado con pérdida L1 solamente; a la derecha, con una combinación de pérdidas discriminatorias L1.



El proceso de capacitación tomó dos días en cuatro GeForce 1080Ti. La red tarda 30 ms en procesar una imagen de 512 x 512. Validación MSE - 34.4. Al igual que con la pintura, las métricas que no desea confiar en las métricas. Es por eso que elegimos seis modelos con las mejores métricas de validación y votamos ciegamente por el mejor modelo.

Cuando ya creamos un sistema de producción y lanzamos un sitio web, continuamos experimentando 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, alimentamos las predicciones netas y una foto de verdad real a la red VGG-16, tomamos los mapas de características en las capas inferiores y los comparamos con MSE. Este enfoque pinta más áreas y da resultados más coloridos.

Recapitulación

Unet es una modelo genial. En la primera tarea de segmentación, enfrentamos un problema durante la capacitación y trabajamos con imágenes de alta resolución y es por eso que usamos In-Place BatchNorm. En nuestra segunda tarea (Inpainting) usamos Convolución parcial en lugar de una predeterminada, y nos permitió obtener mejores resultados. Al trabajar en la coloración, agregamos una pequeña red discriminadora que penalizaba al generador por imágenes poco realistas. También usamos una pérdida perceptual.

Segunda conclusión: los evaluadores son esenciales. Y no solo durante la etapa de creación de máscaras de segmentación sino también para la validación del resultado final. Al final, le damos al usuario tres fotos: una imagen original con defectos pintados, una foto coloreada con defectos pintados y una simplemente coloreada en caso de que el algoritmo para la búsqueda de defectos y la pintura se equivoque.

Tomamos algunas fotos del proyecto War Album y las procesamos en estas neuronas. Aquí están los resultados que obtuvimos:



Además, aquí puede ver más de cerca las imágenes originales y todas las etapas de procesamiento.