En resumen, la tarea de este proyecto es colorear y restaurar fotografías antiguas. Voy a profundizar un poco más en los detalles, pero primero, ¡veamos las fotos! Por cierto, la mayoría de las imágenes de origen se tomaron del subreddit r / TheWayWeWere, agradezco a todos por sus tomas grandes de alta calidad.

Estos son solo algunos ejemplos, ¡y son bastante típicos!

Maria Anderson como la pequeña hada y su página Lyubov Ryabtsova en el ballet La bella durmiente en el Teatro Imperial, San Petersburgo, Rusia, 1890



Una mujer se relaja en su sala de estar (1920, Suecia)



Estudiantes de medicina posando cerca de un cadáver, hacia 1890



Surfista en Hawai, 1890



Caballo giratorio, 1898



El interior del bar Miller y Zapatero, 1899



París en la década de 1880



Vista aérea de Edimburgo en la década de 1920



Mujer de Texas en 1938



Personas en la estación de Waterloo ven televisión por primera vez, Londres, 1936



Lección de geografía en 1850



Fumadores chinos de opio en 1880



Tenga en cuenta que incluso las fotos realmente antiguas y / o de baja calidad aún se ven muy bien:

Deadwood, Dakota del Sur, 1877



Hermanos y hermanas en 1877 (Deadwood)



Portsmouth Square en San Francisco, 1851



Samurai, circa 1860



Por supuesto, el modelo no es perfecto. Esta mano roja me vuelve loco, pero de lo contrario funciona fantásticamente:

Séneca Iroquois Girl, 1908



También puede colorear dibujos en blanco y negro:

Detalles técnicos

Este es un modelo de aprendizaje profundo. En particular, combiné los siguientes enfoques:

• Auto-Atención GAN . Lo único es que el Unet pre-entrenado se usa como generador y simplemente lo cambié para la normalización espectral y, de hecho, el mecanismo de auto atención. Esta es una modificación bastante simple. Te diré que la diferencia es sorprendente en comparación con la versión anterior de Wasserstein GAN, que intenté hacer funcionar. Me gustó la teoría Wasserstein GAN, pero en la práctica no funciona. Pero me enamoré de la red Self-Attention GAN.
• Una estructura de aprendizaje como el crecimiento progresivo de una GAN (pero no exactamente la misma). La diferencia es que el número de capas permanece constante: simplemente cambié el tamaño de los datos de entrada y ajusté la velocidad de aprendizaje para que las transiciones entre los tamaños fueran exitosas. Parece que produce el mismo resultado final, pero aprende más rápido, es más estable y realiza mejor la generalización.
• Regla TTUR ( regla de actualización de dos escalas de tiempo). Aquí está bastante claro: solo una iteración uno a uno del generador / discriminador (crítica) y una mayor velocidad de aprendizaje del discriminador.
• La función de pérdida del generador consta de dos partes: una de ellas es la función principal de la Pérdida perceptiva (o Pérdida de características) basada en VGG16: simplemente empuja el modelo del generador para replicar la imagen de entrada. La segunda parte es la estimación de las pérdidas del discriminador (crítica). Para los curiosos: solo la función de pérdida perceptual no es suficiente para un buen resultado. Tiende a alentar simplemente un montón de marrón / verde / azul: ya sabes, engañando a la prueba, ¿en qué son realmente buenas las redes neuronales! El punto clave es que las GAN esencialmente aprenden la función de pérdida para usted, que en realidad es un gran paso hacia el ideal que estamos luchando en el aprendizaje automático. Y, por supuesto, los resultados mejorarán significativamente cuando la máquina misma aprenda lo que previamente codificó manualmente. Por supuesto, este es el caso aquí.

La belleza de este modelo es que es bastante bueno en una variedad de modificaciones de imagen. Lo que ve arriba son los resultados del modelo de coloración, pero este es solo un componente en la tubería que quiero desarrollar con el mismo modelo.

A continuación, intentaré perfeccionar las imágenes antiguas, y el siguiente elemento en la agenda es un modelo para mejorar la saturación y la riqueza (difuminado). Ahora ella está en las primeras etapas de entrenamiento. Este es básicamente el mismo modelo, pero con algunas configuraciones de contraste / brillo como una simulación de fotos desvaídas e imágenes tomadas con equipos viejos / pobres. Ya he recibido algunos resultados alentadores:

Detalles del proyecto

¿Cuál es la esencia de este proyecto? Solo quiero aplicar GAN para que las fotos antiguas se vean muy, muy bien. Y lo más importante, hará que el proyecto sea útil . Y sí, definitivamente estoy interesado en trabajar con el video, pero primero tengo que descubrir cómo tomar este modelo bajo control del consumo de memoria (esta es una verdadera bestia). Sería bueno si los modelos no aprendieran de dos a tres días en 1080Ti (desafortunadamente, típico de GAN). Aunque este es mi hijo y actualizaré y mejoraré activamente el código en el futuro inmediato, intentaré hacer que el programa sea lo más fácil de usar posible, aunque probablemente haya algunas dificultades con él.

Y juro que documentaré el código correctamente ... algún día. Es cierto que soy una de esas personas que cree en el "código de autodocumentación" (LOL).

Modelo de autolanzamiento

El proyecto se basa en la maravillosa biblioteca Fast.AI. Desafortunadamente, esta es una versión antigua, y aún debe actualizarse a una nueva (definitivamente está en la agenda). Entonces, los requisitos previos, en resumen:

• Antigua biblioteca Fast.AI. Después de haberme enterrado en el proyecto durante dos meses, extrañé un poco lo que le sucedió, porque el que ahora está marcado como "viejo" realmente no se parece al que tengo. Todo ha cambiado en los últimos dos meses más o menos. Por lo tanto, si nada funciona con otras versiones, lo bifurqué aquí . Nuevamente, la actualización a la última versión está en la agenda, me disculpo de antemano.
• Todas las dependencias Fast.AI : hay allí los convenientes archivos require.txt y environment.yml.
• Pytorch 0.4.1 (se requiere spectral_norm, por lo que necesita la última versión estable).
• JupyterLab .
• Tensorboard (es decir, instalar Tensorflow) y TensorboardX . Creo que esto no es estrictamente necesario, pero es mucho más fácil. ¡Para su comodidad, ya he proporcionado todos los ganchos / devoluciones de llamada necesarios en Tensorboard! Hay ejemplos de su uso. Cabe destacar que, de forma predeterminada, las imágenes durante el procesamiento se registran en el Tensorboard cada 200 iteraciones, por lo que obtiene una vista constante y conveniente de lo que hace el modelo.
• ImageNet : un excelente conjunto de datos para el entrenamiento.
• Potente tarjeta gráfica . Realmente me gustaría tener más memoria de 11 GB en mi GeForce 1080Ti. Si tienes algo más débil, será difícil. Unet y Critic son absurdamente geniales, pero cuanto más grandes sean, mejores serán los resultados.

Si desea comenzar el procesamiento de imágenes por su cuenta en este momento sin entrenar al modelo, puede descargar pesos ya preparados aquí . Luego abra ColorizationVisualization.ipynb en JupyterLab. Asegúrese de que haya una línea con un enlace a los pesos:

colorizer_path = Path('/path/to/colorizer_gen_192.h5') 

Luego, debe cargar el modelo de colorizador después de que se inicialice netG:

 load_model(netG, colorizer_path) 

Luego, coloque las imágenes en la carpeta / test_images /, desde donde inicia el programa. Puede visualizar los resultados en el Jupyter Notebook con las siguientes líneas:

 vis.plot_transformed_image("test_images/derp.jpg", netG, md.val_ds, tfms=x_tfms, sz=500) 

Ahorraría un tamaño de aproximadamente 500 px, más o menos, si ejecuta el programa en una GPU con mucha memoria (por ejemplo, GeForce 1080Ti 11 GB). Si hay menos memoria, debe reducir el tamaño de las imágenes o intentar ejecutarlas en la CPU. En realidad intenté hacer esto último, pero por alguna razón el modelo funcionó muy, absurdamente lento, y no encontré tiempo para investigar el problema. Los conocedores recomendaron construir Pytorch a partir de las fuentes, luego se produciría un gran aumento del rendimiento. Hmm ... En ese momento no fue antes de eso.

Información adicional

La visualización de las imágenes generadas a medida que aprende también se puede hacer en Jupyter: solo necesita establecerlo en verdadero al crear una instancia de este enlace de visualización:

GANVisualizationHook(TENSORBOARD_PATH, trainer, 'trainer', jupyter=True, visual_iters=100

Prefiero dejar falso y solo usar Tensorboard. Créeme, también quieres hacer eso. Además, si lo dejas funcionar por mucho tiempo, Jupyter comerá mucha memoria con esas imágenes.

Los pesos de los modelos también se guardan automáticamente durante las ejecuciones de entrenamiento de GANTrainer. Por defecto, se guardan cada 1000 iteraciones (esta es una operación costosa). Se almacenan en la carpeta raíz que especificó para el entrenamiento, y el nombre corresponde a save_base_name especificado en el programa de entrenamiento. Los pesos se almacenan por separado para cada tamaño de entrenamiento.

Recomendaría navegar el código de arriba a abajo, comenzando con el Jupyter Notebook. Tomo estas notas simplemente como una interfaz conveniente para la creación de prototipos y la visualización, todo lo demás irá a los archivos .py tan pronto como encuentre un lugar para ellos. Ya tengo ejemplos de visualización que puede activar convenientemente y ver: simplemente abra xVisualization en el Cuaderno, las imágenes de prueba incluidas en el proyecto se enumeran allí (están en test_images).

Si ve GAN Schedules, entonces esta es la cosa más fea del proyecto, solo mi versión de la implementación de GAN de aprendizaje progresivo, adecuada para el generador Unet.

Los pesos pre-entrenados para el generador de color también están aquí . El proyecto DeFade todavía está en funcionamiento, intentaré sacar buenos pesos en unos días.

Por lo general, durante el entrenamiento, verá los primeros buenos resultados a la mitad, es decir, con un tamaño de 192 px (si usa los ejemplos de entrenamiento provistos).

Estoy seguro de que me equivoqué en algún lugar, así que avíseme si es así.

Problemas conocidos

• Tienes que jugar un poco con el tamaño de la imagen para obtener el mejor resultado. El modelo claramente sufre alguna relación de aspecto y relación de aspecto al generar imágenes. Solía ​​ser mucho peor, pero la situación mejoró significativamente con el aumento de la iluminación / contraste y la introducción del aprendizaje progresivo. Quiero eliminar por completo este problema y centrarme en él, pero hasta ahora no me desespere si la imagen se ve demasiado saturada o con fallas extrañas. Lo más probable es que todo se normalice después de un pequeño cambio de tamaño. Como regla general, para imágenes sobresaturadas necesita aumentar el tamaño.
• Además de lo anterior: obtener las mejores imágenes realmente se reduce al arte de elegir los parámetros óptimos . Sí, los resultados se seleccionan manualmente. Estoy muy satisfecho con la calidad, y el modelo funciona de manera bastante confiable, pero no perfectamente. ¡El proyecto aún está en curso! Creo que la herramienta se puede utilizar como un "artista de IA", pero aún no está lista para el público en general. Simplemente no es el momento.
• Para complicar la situación: en la actualidad, el modelo está comiendo brutalmente la memoria , por lo que en mi tarjeta 1080Ti resulta procesar imágenes con un máximo de 500-600px. Apuesto a que hay muchas opciones de optimización aquí, pero aún no lo he hecho.
• Agregué cero relleno al generador de Unet para cualquier cosa que no se ajuste a los tamaños esperados (así es como puedo cargar una imagen de tamaño arbitrario). Fue un hack muy rápido, y conduce a estúpidos bordes inferiores y derechos en la salida de imágenes de prueba de tamaño arbitrario. Estoy seguro de que hay una mejor manera, pero aún no la he encontrado.
• Modelo ama la ropa azul. No estoy seguro de por qué, ¡la solución está en la búsqueda!

¿Quieres más?

Publicaré nuevos resultados en Twitter .

Adición del traductor.
De este último en Twitter:

Representantes de la propia nacionalidad en su piragua, 1880


( original )

La construcción del metro de Londres, 1860


( original )

Los barrios bajos de Baltimore, 1938



Gimnasio en el Titanic, 1912


( original )