pix2pix.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

Licensed under the Apache License, Version 2.0 (the "License");

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

pix2pix: Traducción de imagen a imagen con un GAN condicional

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

En este tutorial se muestra cómo construir y entrenar una red generativa condicional adversarial (cGAN) llamada pix2pix, la cual aprende un mapeo desde las imágenes de entrada a las de salida, como se describe en Image-to-image translation with conditional adversarial networks{:.external} de Isola et al. (2017). pix2pix no es específico de una determinada aplicación, sino que puede aplicarse a una amplia gama de tareas, como sintetizar fotos a partir de mapas de etiquetas, generar fotos coloreadas a partir de imágenes en blanco y negro, convertir fotos de Google Maps en imágenes aéreas e incluso transformar bocetos en fotos.

En este ejemplo, su red generará imágenes de fachadas de edificios utilizando la base de datos CMP Facade Database proporcionada por el Center for Machine Perception{:.external} de la Universidad Técnica Checa de Praga{:.external}. Para abreviar, utilizará una copia preprocesada{:.external} de este conjunto de datos creado por los autores de pix2pix.

En el cGAN pix2pix, se condiciona a las imágenes de entrada y se generan las imágenes de salida correspondientes. Los cGAN se propusieron por primera vez en Conditional Generative Adversarial Nets (Mirza y Osindero, 2014)

La arquitectura de tu red incluirá:

• Un generador con una arquitectura basada en U-Net{:.external}.

• Un discriminador representado por un clasificador PatchGAN convolucional (propuesto en el documento pix2pix{:. external}).

Tenga en cuenta que cada época puede tardar unos 15 segundos en una sola GPU V100.

A continuación se muestran algunos ejemplos de la salida generada por el cGAN pix2pix después de entrenar durante 200 épocas en el conjunto de datos de fachadas (80,000 pasos).

Importar TensorFlow y otras bibliotecas

import tensorflow as tf

import os
import pathlib
import time
import datetime

from matplotlib import pyplot as plt
from IPython import display

Cargar el conjunto de datos

Descargue los datos de la base de datos de fachadas CMP (30MB). Los conjuntos de datos adicionales están disponibles en el mismo formato aquí{:.external}. En Colab puede seleccionar otros conjuntos de datos en el menú desplegable. Tenga en cuenta que algunos de los otros conjuntos de datos son significativamente mayores (edges2handbags tiene un tamaño de 8 GB).

dataset_name = "facades" #@param ["cityscapes", "edges2handbags", "edges2shoes", "facades", "maps", "night2day"]

_URL = f'http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/{dataset_name}.tar.gz'

path_to_zip = tf.keras.utils.get_file(
    fname=f"{dataset_name}.tar.gz",
    origin=_URL,
    extract=True)

path_to_zip  = pathlib.Path(path_to_zip)

PATH = path_to_zip.parent/dataset_name

list(PATH.parent.iterdir())

Cada imagen original es de tamaño 256 x 512 y contiene dos imágenes 256 x 256:

sample_image = tf.io.read_file(str(PATH / 'train/1.jpg'))
sample_image = tf.io.decode_jpeg(sample_image)
print(sample_image.shape)

plt.figure()
plt.imshow(sample_image)

Es necesario separar las imágenes reales de las fachadas de los edificios de las imágenes de las etiquetas de arquitectura, todas ellas de tamaño 256 x 256.

Defina una función que cargue archivos de imagen y dé como salida dos tensores de imágenes:

def load(image_file):
  # Read and decode an image file to a uint8 tensor
  image = tf.io.read_file(image_file)
  image = tf.io.decode_jpeg(image)

  # Split each image tensor into two tensors:
  # - one with a real building facade image
  # - one with an architecture label image 
  w = tf.shape(image)[1]
  w = w // 2
  input_image = image[:, w:, :]
  real_image = image[:, :w, :]

  # Convert both images to float32 tensors
  input_image = tf.cast(input_image, tf.float32)
  real_image = tf.cast(real_image, tf.float32)

  return input_image, real_image

Represente una muestra de las imágenes de entrada (imagen de la etiqueta de arquitectura) y reales (foto de la fachada del edificio):

inp, re = load(str(PATH / 'train/100.jpg'))
# Casting to int for matplotlib to display the images
plt.figure()
plt.imshow(inp / 255.0)
plt.figure()
plt.imshow(re / 255.0)

Tal y como se describe en el documento pix2pix{:.external}, es necesario aplicar el jittering aleatorio y el reflejo para preprocesar el conjunto de entrenamiento.

Defina varias funciones que:

1. Cambie el tamaño de cada imagen 256 x 256 a un alto y ancho mayores-286 x 286.

2. Recórtelo aleatoriamente a 256 x 256.

3. Dé la vuelta a la imagen horizontalmente de forma aleatoria, es decir, de izquierda a derecha (reflejo aleatorio).

4. Normalice las imágenes al rango [-1, 1].

# The facade training set consist of 400 images
BUFFER_SIZE = 400
# The batch size of 1 produced better results for the U-Net in the original pix2pix experiment
BATCH_SIZE = 1
# Each image is 256x256 in size
IMG_WIDTH = 256
IMG_HEIGHT = 256

def resize(input_image, real_image, height, width):
  input_image = tf.image.resize(input_image, [height, width],
                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  real_image = tf.image.resize(real_image, [height, width],
                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)

  return input_image, real_image

def random_crop(input_image, real_image):
  stacked_image = tf.stack([input_image, real_image], axis=0)
  cropped_image = tf.image.random_crop(
      stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])

  return cropped_image[0], cropped_image[1]

# Normalizing the images to [-1, 1]
def normalize(input_image, real_image):
  input_image = (input_image / 127.5) - 1
  real_image = (real_image / 127.5) - 1

  return input_image, real_image

@tf.function()
def random_jitter(input_image, real_image):
  # Resizing to 286x286
  input_image, real_image = resize(input_image, real_image, 286, 286)

  # Random cropping back to 256x256
  input_image, real_image = random_crop(input_image, real_image)

  if tf.random.uniform(()) > 0.5:
    # Random mirroring
    input_image = tf.image.flip_left_right(input_image)
    real_image = tf.image.flip_left_right(real_image)

  return input_image, real_image

Puede inspeccionar parte de la salida preprocesada:

plt.figure(figsize=(6, 6))
for i in range(4):
  rj_inp, rj_re = random_jitter(inp, re)
  plt.subplot(2, 2, i + 1)
  plt.imshow(rj_inp / 255.0)
  plt.axis('off')
plt.show()

Después de comprobar que la carga y el preprocesamiento funcionan, definamos un par de funciones de ayudante que carguen y preprocesen los conjuntos de entrenamiento y de prueba:

def load_image_train(image_file):
  input_image, real_image = load(image_file)
  input_image, real_image = random_jitter(input_image, real_image)
  input_image, real_image = normalize(input_image, real_image)

  return input_image, real_image

def load_image_test(image_file):
  input_image, real_image = load(image_file)
  input_image, real_image = resize(input_image, real_image,
                                   IMG_HEIGHT, IMG_WIDTH)
  input_image, real_image = normalize(input_image, real_image)

  return input_image, real_image

Cree una canalización de entrada con tf.data

train_dataset = tf.data.Dataset.list_files(str(PATH / 'train/*.jpg'))
train_dataset = train_dataset.map(load_image_train,
                                  num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.shuffle(BUFFER_SIZE)
train_dataset = train_dataset.batch(BATCH_SIZE)

try:
  test_dataset = tf.data.Dataset.list_files(str(PATH / 'test/*.jpg'))
except tf.errors.InvalidArgumentError:
  test_dataset = tf.data.Dataset.list_files(str(PATH / 'val/*.jpg'))
test_dataset = test_dataset.map(load_image_test)
test_dataset = test_dataset.batch(BATCH_SIZE)

Construya el generador

El generador de su cGAN pix2pix es un U-Netmodificado {:.external}. Una U-Net consta de un codificador (downsampler) y un decodificador (upsampler). (Puede encontrar más información al respecto en el tutorial Segmentación de imágenes y en la página web del proyecto U-Net{:.external}).

• Cada bloque del codificador es Convolution -> Batch normalization -> Leaky ReLU

• Cada bloque del descodificador es Transposed convolution -> Batch normalization -> Dropout (aplicado a los 3 primeros bloques) -> ReLU

• Hay conexiones de omisión entre el codificador y el decodificador (como en la U-Net).

Defina el reductor de muestreo o downsampler (codificador):

OUTPUT_CHANNELS = 3

def downsample(filters, size, apply_batchnorm=True):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
      tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
                             kernel_initializer=initializer, use_bias=False))

  if apply_batchnorm:
    result.add(tf.keras.layers.BatchNormalization())

  result.add(tf.keras.layers.LeakyReLU())

  return result

down_model = downsample(3, 4)
down_result = down_model(tf.expand_dims(inp, 0))
print (down_result.shape)

Defina el amplificador de muestreo o upsampler (decodificador):

def upsample(filters, size, apply_dropout=False):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
    tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                    padding='same',
                                    kernel_initializer=initializer,
                                    use_bias=False))

  result.add(tf.keras.layers.BatchNormalization())

  if apply_dropout:
      result.add(tf.keras.layers.Dropout(0.5))

  result.add(tf.keras.layers.ReLU())

  return result

up_model = upsample(3, 4)
up_result = up_model(down_result)
print (up_result.shape)

Defina el generador con el "downsampler" y el "upsampler":

def Generator():
  inputs = tf.keras.layers.Input(shape=[256, 256, 3])

  down_stack = [
    downsample(64, 4, apply_batchnorm=False),  # (batch_size, 128, 128, 64)
    downsample(128, 4),  # (batch_size, 64, 64, 128)
    downsample(256, 4),  # (batch_size, 32, 32, 256)
    downsample(512, 4),  # (batch_size, 16, 16, 512)
    downsample(512, 4),  # (batch_size, 8, 8, 512)
    downsample(512, 4),  # (batch_size, 4, 4, 512)
    downsample(512, 4),  # (batch_size, 2, 2, 512)
    downsample(512, 4),  # (batch_size, 1, 1, 512)
  ]

  up_stack = [
    upsample(512, 4, apply_dropout=True),  # (batch_size, 2, 2, 1024)
    upsample(512, 4, apply_dropout=True),  # (batch_size, 4, 4, 1024)
    upsample(512, 4, apply_dropout=True),  # (batch_size, 8, 8, 1024)
    upsample(512, 4),  # (batch_size, 16, 16, 1024)
    upsample(256, 4),  # (batch_size, 32, 32, 512)
    upsample(128, 4),  # (batch_size, 64, 64, 256)
    upsample(64, 4),  # (batch_size, 128, 128, 128)
  ]

  initializer = tf.random_normal_initializer(0., 0.02)
  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,
                                         strides=2,
                                         padding='same',
                                         kernel_initializer=initializer,
                                         activation='tanh')  # (batch_size, 256, 256, 3)

  x = inputs

  # Downsampling through the model
  skips = []
  for down in down_stack:
    x = down(x)
    skips.append(x)

  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    x = tf.keras.layers.Concatenate()([x, skip])

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

Visualice la arquitectura del modelo generador:

generator = Generator()
tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64)

Pruebe el generador:

gen_output = generator(inp[tf.newaxis, ...], training=False)
plt.imshow(gen_output[0, ...])

Defina la pérdida del generador

Las GAN aprenden que una pérdida se adapta a los datos, mientras que las cGAN aprenden una pérdida estructurada que penaliza una posible estructura que difiera de la salida de la red y de la imagen objetivo, como se describe en el documento pix2pix{:.external}.

• La pérdida del generador es una pérdida de entropía cruzada sigmoidea de las imágenes generadas y un arreglo de unos.

• El documento pix2pix también menciona la pérdida L1, que es un MAE (error promedio absoluto) entre la imagen generada y la imagen objetivo.

• Esto permite que la imagen generada se asemeje estructuralmente a la imagen objetivo.

• La fórmula para calcular la pérdida total del generador es gan_loss + LAMBDA * l1_loss, donde LAMBDA = 100. Este valor se decidió por los autores del documento.

LAMBDA = 100

loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def generator_loss(disc_generated_output, gen_output, target):
  gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output)

  # Mean absolute error
  l1_loss = tf.reduce_mean(tf.abs(target - gen_output))

  total_gen_loss = gan_loss + (LAMBDA * l1_loss)

  return total_gen_loss, gan_loss, l1_loss

El procedimiento de entrenamiento del generador es el siguiente:

Construya el discriminador

El discriminador en el cGAN pix2pix es un clasificador PatchGAN convolucional: intente clasificar si cada imagen patch es real o no es real, como se describe en el documento pix2pix{:.external}.

• Cada bloque del discriminador es Convolution -> Batch normalization -> Leaky ReLU.

• La forma de la salida después de la última capa es (batch_size, 30, 30, 1).

• Cada parche de imagen 30 x 30 de la salida clasifica una porción 70 x 70 de la imagen de entrada.

• El discriminador recibe 2 entradas:

  • La imagen de entrada y la imagen objetivo, que deben clasificarse como reales.

  • La imagen de entrada y la imagen generada (la salida del generador), que debe clasificarse como falsa.

  • Utilice tf.concat([inp, tar], axis=-1) para concatenar estas 2 entradas.

Vamos a definir el discriminador:

def Discriminator():
  initializer = tf.random_normal_initializer(0., 0.02)

  inp = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')
  tar = tf.keras.layers.Input(shape=[256, 256, 3], name='target_image')

  x = tf.keras.layers.concatenate([inp, tar])  # (batch_size, 256, 256, channels*2)

  down1 = downsample(64, 4, False)(x)  # (batch_size, 128, 128, 64)
  down2 = downsample(128, 4)(down1)  # (batch_size, 64, 64, 128)
  down3 = downsample(256, 4)(down2)  # (batch_size, 32, 32, 256)

  zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)  # (batch_size, 34, 34, 256)
  conv = tf.keras.layers.Conv2D(512, 4, strides=1,
                                kernel_initializer=initializer,
                                use_bias=False)(zero_pad1)  # (batch_size, 31, 31, 512)

  batchnorm1 = tf.keras.layers.BatchNormalization()(conv)

  leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)

  zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  # (batch_size, 33, 33, 512)

  last = tf.keras.layers.Conv2D(1, 4, strides=1,
                                kernel_initializer=initializer)(zero_pad2)  # (batch_size, 30, 30, 1)

  return tf.keras.Model(inputs=[inp, tar], outputs=last)

Visualice la arquitectura del modelo discriminador:

discriminator = Discriminator()
tf.keras.utils.plot_model(discriminator, show_shapes=True, dpi=64)

Pruebe el discriminador:

disc_out = discriminator([inp[tf.newaxis, ...], gen_output], training=False)
plt.imshow(disc_out[0, ..., -1], vmin=-20, vmax=20, cmap='RdBu_r')
plt.colorbar()

Defina la pérdida del discriminador

• La función discriminator_loss toma 2 entradas: imágenes reales e imágenes generadas.

• real_loss es una pérdida de entropía cruzada sigmoidea de las imágenes reales y un arreglo de unos (ya que éstas son las imágenes reales).

• generated_loss es una pérdida de entropía cruzada sigmoide de las imágenes generadas y un arreglo de ceros (ya que éstas son las imágenes falsas).

• La total_loss es la suma de real_loss y generated_loss.

def discriminator_loss(disc_real_output, disc_generated_output):
  real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output)

  generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output)

  total_disc_loss = real_loss + generated_loss

  return total_disc_loss

A continuación se muestra el procedimiento de entrenamiento del discriminador.

Para obtener más información sobre la arquitectura y los hiperparámetros, puede consultar el documento pix2pix{:.external}.

Defina los optimizadores y un salvador de puntos de verificación

generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

Generar imágenes

Escriba una función para trazar algunas imágenes durante el entrenamiento.

• Pasar imágenes del equipo de pruebas al generador.

• A continuación, el generador traducirá la imagen de entrada a la de salida.

• ¡El último paso consiste en representar gráficamente las predicciones y voila!

Nota: El training=True es intencional aquí ya que usted quiere las estadísticas por lotes, mientras ejecuta el modelo en el conjunto de datos de prueba. Si utiliza training=False, obtendrá las estadísticas acumuladas aprendidas del conjunto de datos de entrenamiento (que no desea).

def generate_images(model, test_input, tar):
  prediction = model(test_input, training=True)
  plt.figure(figsize=(15, 15))

  display_list = [test_input[0], tar[0], prediction[0]]
  title = ['Input Image', 'Ground Truth', 'Predicted Image']

  for i in range(3):
    plt.subplot(1, 3, i+1)
    plt.title(title[i])
    # Getting the pixel values in the [0, 1] range to plot.
    plt.imshow(display_list[i] * 0.5 + 0.5)
    plt.axis('off')
  plt.show()

Pruebe la función:

for example_input, example_target in test_dataset.take(1):
  generate_images(generator, example_input, example_target)

Entrenamiento

• Para cada entrada de ejemplo se genera una salida.

• El discriminador recibe como primera entrada la input_image y la imagen generada. La segunda entrada es la input_image y la target_image.

• A continuación, calcule la pérdida del generador y del discriminador.

• A continuación, calcule los gradientes de pérdida con respecto a las variables generadoras y discriminadoras (entradas) y aplíquelos al optimizador.

• Por último, registre las pérdidas en TensorBoard.

log_dir="logs/"

summary_writer = tf.summary.create_file_writer(
  log_dir + "fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

@tf.function
def train_step(input_image, target, step):
  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
    gen_output = generator(input_image, training=True)

    disc_real_output = discriminator([input_image, target], training=True)
    disc_generated_output = discriminator([input_image, gen_output], training=True)

    gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target)
    disc_loss = discriminator_loss(disc_real_output, disc_generated_output)

  generator_gradients = gen_tape.gradient(gen_total_loss,
                                          generator.trainable_variables)
  discriminator_gradients = disc_tape.gradient(disc_loss,
                                               discriminator.trainable_variables)

  generator_optimizer.apply_gradients(zip(generator_gradients,
                                          generator.trainable_variables))
  discriminator_optimizer.apply_gradients(zip(discriminator_gradients,
                                              discriminator.trainable_variables))

  with summary_writer.as_default():
    tf.summary.scalar('gen_total_loss', gen_total_loss, step=step//1000)
    tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=step//1000)
    tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=step//1000)
    tf.summary.scalar('disc_loss', disc_loss, step=step//1000)

El bucle de entrenamiento propiamente dicho. Como este tutorial puede ejecutarse con más de un conjunto de datos, y éstos varían mucho en tamaño, el bucle de entrenamiento está configurado para trabajar por pasos en vez de por épocas.

• Itera sobre el número de pasos.

• Cada 10 pasos imprime un punto (.).

• Cada 1,000 pasos: borre la pantalla y ejecute generate_images para mostrar el progreso.

• Cada 5,000 pasos: guarde un punto de verificación.

def fit(train_ds, test_ds, steps):
  example_input, example_target = next(iter(test_ds.take(1)))
  start = time.time()

  for step, (input_image, target) in train_ds.repeat().take(steps).enumerate():
    if (step) % 1000 == 0:
      display.clear_output(wait=True)

      if step != 0:
        print(f'Time taken for 1000 steps: {time.time()-start:.2f} sec\n')

      start = time.time()

      generate_images(generator, example_input, example_target)
      print(f"Step: {step//1000}k")

    train_step(input_image, target, step)

    # Training step
    if (step+1) % 10 == 0:
      print('.', end='', flush=True)


    # Save (checkpoint) the model every 5k steps
    if (step + 1) % 5000 == 0:
      checkpoint.save(file_prefix=checkpoint_prefix)

Este bucle de entrenamiento guarda registros que puede ver en TensorBoard para supervisar el progreso del entrenamiento.

Si trabaja en una máquina local, lanzaría un proceso TensorBoard separado. Cuando trabaje en un cuaderno, lance el visor antes de iniciar el entrenamiento para monitorizar con TensorBoard.

Inicie el visor TensorBoard (Lo sentimos, esto no se muestra en tensorflow.org):

%load_ext tensorboard
%tensorboard --logdir {log_dir}

Puede ver los resultados de una ejecución previa de este bloc de notas en TensorBoard.dev.

Por último, ejecute el bucle de entrenamiento:

fit(train_dataset, test_dataset, steps=40000)

La interpretación de los registros es más sutil cuando se entrena un GAN (o un cGAN como pix2pix) en comparación con un modelo sencillo de clasificación o regresión. Cosas en las que debe fijarse:

• Revise que ni el modelo generador ni el discriminador hayan "ganado". Si el gen_gan_loss o el disc_loss se vuelven muy bajos, es un indicador de que este modelo está dominando al otro, y usted no está entrenando correctamente el modelo combinado.

• El valor log(2) = 0.69 es un buen punto de referencia para estas pérdidas, ya que indica una perplejidad de 2: el discriminador tiene, por término medio, la misma incertidumbre sobre las dos opciones.

• Para el disc_loss, un valor inferior a 0,69 significa que el discriminador obtiene mejores resultados que el azar en el conjunto combinado de imágenes reales y generadas.

• Para el gen_gan_loss, un valor inferior a 0,69 significa que el generador lo hace mejor que el azar al momento de engañar al discriminador.

• A medida que progresa el entrenamiento, la gen_l1_loss debería disminuir.

Restaure el último punto de verificación y pruebe la red

!ls {checkpoint_dir}

# Restoring the latest checkpoint in checkpoint_dir
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

Genere algunas imágenes utilizando el conjunto de prueba

# Run the trained model on a few examples from the test set
for inp, tar in test_dataset.take(5):
  generate_images(generator, inp, tar)