Restaurar color a partir de imágenes en escala de grises
Este proyecto de fin de semana empezó como una idea de compresión de imágenes. Una imagen en escala de grises conserva un canal en lugar de los tres canales de RGB, por lo que me preguntaba si una red neuronal podría aprender suficiente sobre un dominio concreto como para reconstruir colores plausibles a partir de esa información que falta.
Podría tener sentido en un caso muy específico. Si quisieras comprimir capítulos de una misma serie, por ejemplo, un modelo entrenado con esa serie podría aprender sus personajes, lugares y paleta de colores recurrentes. No recuperaría la imagen original de forma exacta, porque el fotograma en escala de grises da información sobre luminosidad, bordes y estructura, pero no dice si un coche era rojo o azul. Aun así, era algo que merecía la pena experimentar.
Empecé este proyecto con una red convolucional pequeña que recibía un canal en escala de grises y predecía tres canales RGB con MSE como función de pérdida. El loss bajaba rápido y después dejaba de mejorar más o menos en el mismo punto, independientemente de que usara una arquitectura más profunda o una ResNet. Las imágenes salían sobre todo beige o sepia.
Pensaba que sería un experimento pequeño, pero ese primer fallo dejó claro que el modelo no era el único problema. Cuando hay muchos colores plausibles para un mismo objeto, un modelo entrenado con MSE puede acercarse a una respuesta promedio, que termina siendo una imagen beige.
Reimplementar un paper
Empecé a buscar trabajos sobre colorización de imágenes y encontré Let there be Color!, una arquitectura de 2016 para colorización automática de imágenes. Su planteamiento tenía más sentido para el problema.
La red recibe el canal L del espacio de color LAB, que contiene la luminosidad de la imagen. Solo predice los canales a y b, que representan la crominancia. El canal L original se mantiene y se combina con la predicción al final.
Es una mejor forma de abordar el problema, porque la red no necesita recrear el contraste ni la geometría que ya se ven en la imagen en escala de grises. Puede centrarse en asignar un color plausible a cada región.
Características globales y locales
El paper también usa dos ramas después de las primeras capas convolucionales. Una mantiene la información espacial local necesaria para colorear una región concreta. La otra comprime la imagen en características globales y predice una categoría de escena de Places365.
La clasificación de escenas es una tarea auxiliar. Le da a la red contexto sobre la imagen completa, algo útil cuando la apariencia local por sí sola es ambigua. Una región azul puede ser cielo, agua o una pared, según la clase de la imagen. Las características globales se fusionan de nuevo con las locales antes de que las capas finales predigan los dos canales de crominancia.
Reimplementé la red de 2016 en PyTorch, incluidas sus ramas de características locales y globales, la tarea auxiliar de clasificación de escenas y las capas de colorización que predicen los canales a y b a partir de L.
También añadí checkpoints, soporte para reanudar, logging con Weights & Biases, subconjuntos del dataset y una interfaz pequeña con Makefile alrededor de Places365 para gestionar mejor los entrenamientos. Con esa preparación, reproduje el procedimiento de entrenamiento y los resultados de colorización del paper de 2016.
Entrada
Salida del modelo, Época 1Resultados
El invernadero es uno de los ejemplos que mejor funcionan. El modelo entiende la estructura general y produce verdes y marrones plausibles, aunque estén menos saturados que los originales. La escena de montañas y la cena también se reconocen, aunque los colores elegidos no coincidan exactamente.
El dibujo animado funciona mucho peor porque está completamente fuera de la distribución de entrenamiento. Sus colores siguen convenciones que un modelo entrenado con escenas naturales no tiene por qué conocer. El vídeo del parque funciona razonablemente bien fotograma a fotograma, pero tiene la misma tendencia a usar colores conservadores y desaturados.
Fine-tuning para un dominio específico
Quise comprobar cuánto podía cambiar el resultado en un dominio más específico. Hice un fine-tuning del checkpoint de Places365 con fotogramas de un vídeo de un bosque. Elegí aleatoriamente un tramo continuo de 20 segundos y lo excluí por completo del entrenamiento y la validación. El modelo mantuvo la arquitectura LAB original, con el clasificador de escenas adaptado a una única clase, forest.
Coloreé exactamente el mismo clip de bosque retenido con ambos modelos. El modelo fine-tuned redujo el error de color LAB un 77% y evitó el tono anaranjado que el modelo de Places365 aplicaba al bosque verde. Repetí la misma técnica en otro fine-tuning independiente con metraje de Barcelona. El modelo de Places365 hacía muchos más cambios de color, mientras que el modelo fine-tuned para esa escena los mantenía mucho más coherentes entre fotogramas. Estas comparativas sugieren que la idea funciona mucho mejor cuando el modelo se entrena o se hace fine-tuning para un dominio específico.
Clip de bosque retenido
Clip de Barcelona
Ese es el límite de la idea original de compresión. Una imagen en escala de grises no conserva la información de color exacta necesaria para recuperar el fotograma original. Un modelo bien entrenado puede hacer una predicción plausible, y eso puede ser útil en un dominio concreto, pero sigue siendo una predicción.
El proyecto incluye comandos de entrenamiento, checkpoints e inferencia para imágenes y vídeo en el repositorio. Reimplementar el paper fue la forma de aprender cómo se comportan esos inputs, outputs y limitaciones cuando el modelo tiene que entrenar con un dataset real durante varios días.