Archivo del Autor: MARIO GUTIERREZ DELGADO

Colorización de imágenes usando Generative Adversarial Networks

Deja un comentario

Introducción

A lo largo de este trabajo vamos a intentar llevar a cabo la colorización de imágenes en blanco y negro de manera que se parezcan lo máximo posible a la realidad. Para la realización de esta investigación, hemos decidido utilizar un dataset llamado PETS, que contiene 7390 imágenes de mascotas en diversos ambientes. El motivo que nos llevó a utilizar este dataset es, simplemente, que no hemos encontrado ningún trabajo previo que lo utilizase con esta finalidad. Además, observamos que contenía una gran variedad de imágenes diferentes entre sí, por lo que es apto para el objetivo de esta investigación.

Para la colorización de las imágenes en blanco y negro, hemos utilizado como inspiración el código propuesto en el artículo de Medium llamado Colorizing black & white images with U-Net and Conditional GAN – a tutorial.

Estado del arte

La colorización de imágenes usando sistemas automáticos no es un proceso sencillo por la gran subjetividad de colores que se pueden apreciar en una imagen y el rango de sensibilidad que pueda medir en cada píxel. Con las Generative Adversarial Networks se puede aproximar al problema desde un punto de vista novedoso con grandes resultados.

La arquitectura pix2pix (Isola et al, 2017) de las GAN permite el aprendizaje directo entre la imagen de entrada y la de salida lo que agiliza el proceso y mejora los resultados obtenidos, esto no solo se ha visto reflejado en la colorización de imágenes sino en todos aquellos procesos de aprendizaje automático relativos a imágenes o vídeos.

Kamyar Nazeri, Eric Ng y Mehran Ebrahimi (Nazeri, Ng & Ebrahimi, 2018) crearon las bases de una Deep Convolutional GAN para el proceso de colorear imágenes automáticamente usando como entrada imágenes en escalas de grises de los conjuntos de entrenamientos públicos CIFAR-10 y Places365. La mayor ventaja de este modelo se basa en la función de coste para limitar el entrenamiento y que la parte generadora no cree colores irreales.

Rahul Reddy Pasham, Sameer Md y Ashwini K. (Ashwini, Pasham & Sameer, 2022) también utilizaron GANs en el problema de colorear imágenes y en su caso destacaron las capas de activaciones LeakyReLu y tanh y la capacidad de colorear imágenes a pesar de ser datos no etiquetados.

El mayor reto de las GANS en el proyecto de colorización de imágenes actualmente es la generación precisa de colores.

Desarrollo

LAB

En el espacio LAB (Shariatnia, 2020), al igual que en RGB, cada píxel se representa con tres valores. El primer canal, L, codifica la luminosidad (Lightness) del píxel. El segundo, A, representa la mezcla entre verde y rojo, y el tercero, B, indica la mezcla entre amarillo y azul.

Se utiliza este espacio de color en lugar de RGB porque la imagen que se le proporciona al modelo para entrenar debe estar en blanco y negro. Al usar LAB, se puede enviar directamente el canal correspondiente a la luminosidad, sin necesidad de convertir la imagen a escala de grises. De este modo, el modelo solo necesita predecir los canales A y B. Esto reduce el número de combinaciones posibles: de 256³ (alrededor de 17 millones) en el espacio RGB a 256² (65.536) al predecir únicamente dos canales de color.

GAN

Las GANs (Generative Adversarial Networks) son una red neuronal compuesta por dos modelos que compiten entre sí: un generador, que aprende a crear datos falsos que imitan los datos reales, y un discriminador, que intenta distinguir entre datos reales y generados. 

Ambos se entrenan simultáneamente en un proceso de competencia: el generador mejora para engañar al discriminador, y el discriminador mejora para no ser engañado. Con el tiempo, esta dinámica hace que el generador produzca datos cada vez más realistas.

Si el entrenamiento del generador es muy bueno puede generar imágenes que al discriminador le sean difíciles de clasificar, el objetivo es que la precisión no disminuya a medida que mejora la generación de imágenes (Google Developers, s.f.).

Generador

Nuestro generador U-Net sigue una arquitectura encoder-decoder con skip-connections. Podemos destacar varias cosas de la arquitectura. En primer lugar, el uso de LeakyReLU en la parte del encoder con una slope de 0.2 y de ReLU normal en la parte del decoder (Nazeri, Ng & Ebrahimi, 2018). Por otro lado, el uso de función de activación tanh en la última capa (Isola et al., 2017; Ashwini, Pasham & Sameer, 2022) y, por último, la utilización de Batch Normalization, que nos ayuda a evitar que el entrenamiento de la GAN converja a una serie de parámetros que siempre devuelven el mismo output.  Este escenario se llama mode-collapse o the Helvetica scenario (Nazeri, Ng & Ebrahimi, 2018).

Patch Discriminator

Un Patch Discriminator es un tipo de discriminador usado en GANs (por ejemplo en la arquitectura Pix2Pix) que evalúa si pequeñas subregiones de la imagen parecen realistas.

En lugar de dar un único valor de salida , el Patch Discriminator produce una matriz de valores, donde cada valor indica si una pequeña subregión de la imagen se considera real o falso.

Es una buena idea usarlo para colorización por varios motivos: hace que el discriminador se enfoque en los detalles locales, lo que provoca que el generador acierte en cada parte de la imagen. Además, tiene menos parámetros que un discriminador global y permite un entrenamiento más estable al evaluar diferentes regiones de la imagen.

Función de pérdida

En nuestra función de pérdida, debemos destacar la aplicación de One Sided Label Smoothing; ya que, de esta manera, evitamos que la red emita outputs extremadamente seguros y aprenda mejor a la hora de realizar el entrenamiento (Nazeri, Ng & Ebrahimi, 2018). El smoothing lo realizamos en la etiqueta positiva (pasamos de 1 a 0.9) y la etiqueta negativa la dejamos intacta.

La función de pérdida establece dos modos. El modo vanilla introduce el uso de Binary Cross Entropy con Logistic Loss, mientras que el modo lsgan introduce Mean Squared Error Loss.

Esta función de pérdida guía tanto al discriminador como al generador en su entrenamiento: el discriminador aprende a diferenciar imágenes reales de generadas, mientras que el generador intenta engañarlo generando imágenes que el discriminador clasifique como reales.

Nuestro modelo

El primer paso es obtener las imágenes en el formato adecuado; en este caso, utilizamos imágenes de tamaño 256×256 en el espacio de color LAB en lugar de RGB. 

En el método __init__, se definen e inicializan el generador y el discriminador. Se configuran las pérdidas GAN y L1 (coherencia de color), y se establecen los optimizadores Adam para ambos modelos.

El método optimize gestiona el entrenamiento. Primero, se pasa la imagen en escala de grises al generador, obteniendo fake_color.

Luego, se entrena el discriminador (backward_D):

  • Se le alimentan las imágenes generadas (fake_color), marcadas como “falsas” .
  • Se le alimentan imágenes reales del conjunto de entrenamiento, marcadas como “reales”.
  • Se calcula la pérdida promedio de “falsas” y “reales”, y se retropropaga para actualizar los pesos del discriminador.

A continuación, se entrena el generador (backward_G):

  • Se alimentan al discriminador las imágenes generadas, pero con etiquetas “reales” para calcular la pérdida adversarial (engañar al discriminador).
  • Se calcula la pérdida L1, que mide la diferencia a nivel de píxel con las imágenes reales (multiplicada por 100).
  • La suma de la pérdida adversarial y la L1 es la pérdida total del generador, que se utiliza para actualizar sus pesos mediante retropropagación.

Finalmente, optimize coordina todo el ciclo de entrenamiento: primero genera la predicción, luego optimiza el discriminador y, a continuación, optimiza el generador, repitiendo este proceso para cada lote del conjunto de entrenamiento. Este proceso debe evitar el mode-collapse, un escenario en el que el generador converge a una serie de parámetros que siempre devuelven el mismo output, lo que limita la diversidad de las imágenes generadas (Shariatnia, 2020).
El código entrena el modelo durante 50 épocas utilizando el conjunto de datos de entrenamiento. En cada época, se procesan los lotes (batches) del conjunto de entrenamiento, y en cada iteración, el modelo optimiza tanto el generador como el discriminador. El generador crea imágenes coloreadas a partir de las imágenes en escala de grises, mientras que el discriminador evalúa si las imágenes generadas son reales o no. Las pérdidas de ambos modelos se registran, y cada 350 iteraciones se visualiza el rendimiento del modelo y se imprime el estado de las pérdidas. El objetivo es que, con el tiempo, el generador aprenda a crear imágenes más realistas y el discriminador mejore su capacidad para distinguir entre imágenes reales y generadas.

Materiales

En este apartado, vamos a compartir el código más importante y necesario para realizar esta tarea de colorización de imágenes en blanco y negro.

El generador U-Net de nuestra GAN (se necesita tener instalada la librería PyTorch de Python, entre otras):

class UnetBlock(nn.Module):
    def __init__(self, nf, ni, submodule=None, input_c=None, dropout=False,
                 innermost=False, outermost=False):
        super().__init__()
        self.outermost = outermost
        if input_c is None: input_c = nf
        downconv = nn.Conv2d(input_c, ni, kernel_size=4,
                             stride=2, padding=1, bias=False)
        downrelu = nn.LeakyReLU(0.2, True)
        downnorm = nn.BatchNorm2d(ni)
        uprelu = nn.ReLU(True)
        upnorm = nn.BatchNorm2d(nf)

        if outermost:
            upconv = nn.ConvTranspose2d(ni * 2, nf, kernel_size=4,
                                        stride=2, padding=1)
            down = [downconv]
            up = [uprelu, upconv, nn.Tanh()] #Utilizamos función de activación tanh en la última capa como se indica en Image Colorization using Generative Adversarial Networks
            model = down + [submodule] + up
        elif innermost:
            upconv = nn.ConvTranspose2d(ni, nf, kernel_size=4,
                                        stride=2, padding=1, bias=False)
            down = [downrelu, downconv]
            up = [uprelu, upconv, upnorm]
            model = down + up
        else:
            upconv = nn.ConvTranspose2d(ni * 2, nf, kernel_size=4,
                                        stride=2, padding=1, bias=False)
            down = [downrelu, downconv, downnorm]
            up = [uprelu, upconv, upnorm]
            if dropout: up += [nn.Dropout(0.5)]
            model = down + [submodule] + up
        self.model = nn.Sequential(*model)

    def forward(self, x):
        if self.outermost:
            return self.model(x)
        else:
            return torch.cat([x, self.model(x)], 1)

class Unet(nn.Module):
    def __init__(self, input_c=1, output_c=2, n_down=8, num_filters=64):
        super().__init__()
        unet_block = UnetBlock(num_filters * 8, num_filters * 8, innermost=True)
        for _ in range(n_down - 5):
            unet_block = UnetBlock(num_filters * 8, num_filters * 8, submodule=unet_block, dropout=True)
        out_filters = num_filters * 8
        for _ in range(3):
            unet_block = UnetBlock(out_filters // 2, out_filters, submodule=unet_block)
            out_filters //= 2
        self.model = UnetBlock(output_c, out_filters, input_c=input_c, submodule=unet_block, outermost=True)

    def forward(self, x):
        return self.model(x)

Nuestro Patch Discriminator:

class PatchDiscriminator(nn.Module):
    def __init__(self, input_c, num_filters=64, n_down=3):
        super().__init__()
        model = [self.get_layers(input_c, num_filters, norm=False)]
        model += [self.get_layers(num_filters * 2 ** i, num_filters * 2 ** (i + 1), s=1 if i == (n_down-1) else 2)
                          for i in range(n_down)] 
        model += [self.get_layers(num_filters * 2 ** n_down, 1, s=1, norm=False, act=False)] 
        self.model = nn.Sequential(*model)

    def get_layers(self, ni, nf, k=4, s=2, p=1, norm=True, act=True): 
        layers = [nn.Conv2d(ni, nf, k, s, p, bias=not norm)]          
        if norm: layers += [nn.BatchNorm2d(nf)]
        if act: layers += [nn.LeakyReLU(0.2, True)]
        return nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

La función de pérdida:

class GANLoss(nn.Module):
    def __init__(self, gan_mode='vanilla', real_label=0.9, fake_label=0.0): #Aplicamos un One-Sided Label Smoothing como se indica en Image Colorizarion using Generative Adversarial Networks
        super().__init__()
        self.register_buffer('real_label', torch.tensor(real_label))
        self.register_buffer('fake_label', torch.tensor(fake_label))
        if gan_mode == 'vanilla':
            self.loss = nn.BCEWithLogitsLoss()
        elif gan_mode == 'lsgan':
            self.loss = nn.MSELoss()

    def get_labels(self, preds, target_is_real):
        if target_is_real:
            labels = self.real_label
        else:
            labels = self.fake_label
        return labels.expand_as(preds)

    def __call__(self, preds, target_is_real):
        labels = self.get_labels(preds, target_is_real)
        loss = self.loss(preds, labels)
        return loss

La inicialización del modelo y la definición del mismo:

def init_weights(net, init='norm', gain=0.02):

    def init_func(m):
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and 'Conv' in classname:
            if init == 'norm':
                nn.init.normal_(m.weight.data, mean=0.0, std=gain)
            elif init == 'xavier':
                nn.init.xavier_normal_(m.weight.data, gain=gain)
            elif init == 'kaiming':
                nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')

            if hasattr(m, 'bias') and m.bias is not None:
                nn.init.constant_(m.bias.data, 0.0)
        elif 'BatchNorm2d' in classname:
            nn.init.normal_(m.weight.data, 1., gain)
            nn.init.constant_(m.bias.data, 0.)

    net.apply(init_func)
    print(f"model initialized with {init} initialization")
    return net

def init_model(model, device):
    model = model.to(device)
    model = init_weights(model)
    return model

class MainModel(nn.Module):
    def __init__(self, net_G=None, lr_G=2e-4, lr_D=2e-4, #Utilizamos initial learning rate de 2x10-4 como se indica en Image Colorization using Generative Adversarial Networks
                 beta1=0.5, beta2=0.999, lambda_L1=100.): #Utilizamos lambda = 100 y beta1 = 0.5 (para reducir el momento) como indica la misma fuente
        super().__init__()

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(self.device)
        self.lambda_L1 = lambda_L1

        if net_G is None:
            self.net_G = init_model(Unet(input_c=1, output_c=2, n_down=8, num_filters=64), self.device)
        else:
            self.net_G = net_G.to(self.device)
        self.net_D = init_model(PatchDiscriminator(input_c=3, n_down=3, num_filters=64), self.device)
        self.GANcriterion = GANLoss(gan_mode='vanilla').to(self.device)
        self.L1criterion = nn.L1Loss()
        self.opt_G = optim.Adam(self.net_G.parameters(), lr=lr_G, betas=(beta1, beta2))
        self.opt_D = optim.Adam(self.net_D.parameters(), lr=lr_D, betas=(beta1, beta2))

    def set_requires_grad(self, model, requires_grad=True):
        for p in model.parameters():
            p.requires_grad = requires_grad

    def setup_input(self, data):
        self.L = data['L'].to(self.device)
        self.ab = data['ab'].to(self.device)

    def forward(self):
        self.fake_color = self.net_G(self.L)

    def backward_D(self):
        fake_image = torch.cat([self.L, self.fake_color], dim=1)
        fake_preds = self.net_D(fake_image.detach())
        self.loss_D_fake = self.GANcriterion(fake_preds, False)
        real_image = torch.cat([self.L, self.ab], dim=1)
        real_preds = self.net_D(real_image)
        self.loss_D_real = self.GANcriterion(real_preds, True)
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
        self.loss_D.backward()

    def backward_G(self):
        fake_image = torch.cat([self.L, self.fake_color], dim=1)
        fake_preds = self.net_D(fake_image)
        self.loss_G_GAN = self.GANcriterion(fake_preds, True)
        self.loss_G_L1 = self.L1criterion(self.fake_color, self.ab) * self.lambda_L1
        self.loss_G = self.loss_G_GAN + self.loss_G_L1
        self.loss_G.backward()

    def optimize(self):
        self.forward()
        self.net_D.train()
        self.set_requires_grad(self.net_D, True)
        self.opt_D.zero_grad()
        self.backward_D()
        self.opt_D.step()

        self.net_G.train()
        self.set_requires_grad(self.net_D, False)
        self.opt_G.zero_grad()
        self.backward_G()
        self.opt_G.step()

El bucle de entrenamiento (con 50 épocas):

def train_model(model, train_dl, epochs, display_every=350):
    data = next(iter(val_dl)) 
    for e in range(epochs):
        loss_meter_dict = create_loss_meters() 
        i = 0                                 
        for data in tqdm(train_dl):
            model.setup_input(data)
            model.optimize()
            update_losses(model, loss_meter_dict, count=data['L'].size(0)) 
            i += 1
            if i % display_every == 0:
                print(f"\nEpoch {e+1}/{epochs}")ç
                print(f"Iteration {i}/{len(train_dl)}")
                log_results(loss_meter_dict) 
                visualize(model, data, save=False) 

model = MainModel()
train_model(model, train_dl, 50)

Resultados

La primera fila son las imágenes que se introducen al generador en blanco y negro. La segunda fila son el resultado de la colorización por parte del generador y la tercera fila son las imágenes originales.

Discusión

Podemos observar que la colorización se ha realizado correctamente en todas las imágenes, aunque existen ligeros defectos en algunas de ellas. En la segunda imagen desde la izquierda, el suelo no se ha coloreado correctamente en ciertas zonas. Por otro lado, en la segunda imagen desde la derecha, el perro presenta una mancha verde en la cabeza. Esto se puede deber a que la imagen presenta una gran cantidad de este color y pueden ocurrir errores por este motivo. Por último, en la imagen de la derecha, parece que se ha aplicado un filtro sepia y el resultado queda entre medias de la imagen real y de la imagen en blanco y negro.

Referencias

Este entrada ha sido realizada por Mario Gutiérrez, Diego Rivero y Alejandro Tapia.

About MARIO GUTIERREZ DELGADO