Codificadores automáticos de colorización usando Keras

Este artículo ofrece un caso de uso práctico de Autoencoders , es decir, la colorización de imágenes en escala de grises . Usaremos Keras para codificar el codificador automático.

Como todos sabemos, un AutoEncoder tiene dos operadores principales:

Codificador Esto transforma la entrada en un vector latente de baja dimensión. A medida que reduce la dimensión, se ve obligado a aprender las características más importantes de la entrada.
Decodificador: trata de reconstruir la entrada tanto como sea posible a partir del vector latente.

Durante el diseño de Autoencoder, es muy esencial elegir correctamente una dimensión latente. Como si fuera
más que la dimensión de entrada, Autoencoder tiende a memorizar la entrada. Implementaremos la parte del codificador usando CNN y usaremos Conv2DTranspose para la sección del decodificador del codificador automático.

Para simplificar las cosas, vamos a utilizar
Conjunto de datos CIFAR100
, que está fácilmente disponible en
Conjuntos de datos de Keras
El conjunto de datos contiene
50k color
imágenes de forma
32 * 32 * 3
para la formación, y
10k color
imágenes de la misma forma con fines de prueba.
Código: Importar todas las bibliotecas

import numpy as np
import matplotlib.pyplot as plt
import os
  
from keras.layers import Dense, Input, Conv2D, Conv2DTranspose, Flatten, Reshape
from keras.models import Model
from keras.callbacks import ReduceLROnPlateau, ModelCheckpoint
from keras.datasets import cifar100
from keras import backend as K

Como el conjunto de datos contiene solo imágenes en color, para el propósito de nuestra tarea debemos cambiarlo a escala de grises. Por lo tanto, definimos una función para eso.
Código: función para convertir imágenes RGB a escala de grises

def rgb_2_gray(image):
    return np.dot(image[..., :3], [0.299, 0.587, 0.114])

Código: cargar el conjunto de datos

(x_train, _), (x_test, _) = cifar100.load_data()

Para que el modelo aprenda de manera eficiente, es mejor convertir las imágenes en flotantes. También necesitamos normalizar los valores para que estén entre 0 y 1 . Esto se hace para que durante la retropropagación, los gradientes no se salgan de control.

Código: Normalizar los datos

x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
  
x_train_gray = x_train_gray.astype('float32') / 255.
x_test_gray = x_test_gray.astype('float32') / 255.

El rendimiento de los modelos de Deep Learning depende en gran medida del conjunto de hiperparámetros (incluido el número de capas, el número de filtros en cada capa, el tamaño del lote, etc.). Por lo tanto, una buena elección de los hiperparámetros es una habilidad esencial. Para obtener los mejores resultados, debemos probar y experimentar con un conjunto diferente de ellos. Aquí, estamos usando estos conjuntos de hiperparámetros,
Código: Hiperparámetros

input_shape = (rows, cols, 1)
batch_size = 32
kernel_size = 3
latent_dim = 256
layer_filters = [64, 128, 256]

Para la tarea de colorear, la entrada es una imagen en escala de grises. La imagen en escala de grises tiene solo 1 canal en comparación con las imágenes en color que tienen 3, a saber, rojo, verde y azul. Usamos Input from Keras library para tomar una entrada con la forma de (filas, columnas, 1) .
El codificador es una pila de 3 capas convolucionales con un número creciente de filtros, seguido de una capa densa con 256 unidades para generar vectores latentes.

Código: Codificador

inputs = Input(shape = input_shape)
x = inputs
for filters in layer_filters:
  x = Conv2D(filters = filters,
           kernel_size = kernel_size,
           strides = 2,
           activation ='relu',
           padding ='same')(x)
  
shape = K.int_shape(x)
x = Flatten()(x)
latent = Dense(latent_dim, name ='latent_vector')(x)
encoder = Model(inputs, latent, name ='encoder')

La sección del decodificador del Autoencoder intenta descomprimir el vector latente para la entrada. En nuestro caso, la entrada al Decodificador es una capa de forma (Ninguno, 256) . Sigue una pila de tres capas DeConvolutional con números de filtros decrecientes en cada capa. Nos aseguramos de que la última capa, en este caso, sea de forma (Ninguna, 32, 32, 3) . El número de canales debe ser 3 para poder comparar la reconstrucción con la verdad de campo de las imágenes durante la retropropagación.
No es obligatorio que el codificador y el decodificador sean una imagen especular de los dos.

Código: Decodificador

latent_inputs = Input(shape =(latent_dim, ), name ='decoder_input')
x = Dense(shape[1]*shape[2]*shape[3])(latent_inputs)
x = Reshape((shape[1], shape[2], shape[3]))(x)
# stack of Conv2DTranspose(256)-Conv2DTranspose(128)-
# Conv2DTranspose(64)
for filters in layer_filters[::-1]:
  x = Conv2DTranspose(filters = filters,
                      kernel_size = kernel_size,
                      strides = 2,
                      activation ='relu',
                      padding ='same')(x)
outputs = Conv2DTranspose(filters = channels,
                          kernel_size = kernel_size,
                          activation ='sigmoid',
                          padding ='same',
                          name ='decoder_output')(x)
decoder = Model(latent_inputs, outputs, name ='decoder')

Finalmente, definimos el modelo, denominado autoencoder, que toma una entrada y luego la pasa al codificador y luego la pasa por el decodificador.

autoencoder = Model(inputs, decoder(encoder(inputs)),
                    name ='autoencoder')

Ahora entrenamos el modelo de codificador automático dividiendo todos los datos en lotes de tamaño de lote = tamaño_de_lote , durante 30 épocas . El punto importante a tener en cuenta aquí es que, si verificamos la función de ajuste, encontramos que la entrada al modelo es el conjunto de datos de imágenes en escala de grises y la imagen en color correspondiente sirve como etiqueta. También sucede algo similar para el conjunto de validación.
Generalmente, para la tarea de clasificación, alimentamos las imágenes al modelo como entradas, y sus respectivas clases se dan como etiquetas. y durante el entrenamiento, comparamos el rendimiento del modelo midiendo qué tan bien clasifica las imágenes en sus respectivas clases dadas como etiquetas. Pero, para esta tarea, proporcionamos las imágenes en color como etiquetas, ya que queremos que el modelo proporcione las imágenes RGB como salidas cuando le proporcionamos una imagen en escala de grises.
También hemos utilizado devoluciones de llamada para reducir la tasa de aprendizaje si la pérdida de validación no disminuye mucho.

autoencoder.fit(x_train_gray,
                x_train,
                validation_data =(x_test_gray, x_test),
                epochs = 30,
                batch_size = batch_size,
                callbacks = callbacks)

Código: Resultados y análisis

Datos reales de las primeras 100 imágenes de prueba

Entrada en escala de grises al codificador automático

Salida coloreada del Autoencoder

El autocodificador tiene un rendimiento aceptable en el trabajo de colorización. Predijo correctamente que el cielo es azul, los chimpancés tienen diferentes tonos de marrón, las hojas son verdes, etc. es oscuro y no rojizo, etc.