Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
tensorflow
GitHub Repository: tensorflow/docs-l10n
Path: blob/master/site/pt-br/tutorials/generative/pix2pix.ipynb
25118 views
Kernel: Python 3

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: tradução de imagem para imagem com um GAN condicional

Este tutorial demonstra como construir e treinar uma rede adversária generativa condicional (cGAN) chamada pix2pix que aprende um mapeamento de imagens de entrada para imagens de saída, conforme descrito em Tradução de imagem para imagem com redes adversárias condicionais{:.external} por Isola et al. (2017). O pix2pix não é específico para uma aplicação – ele pode ser aplicado a uma ampla gama de tarefas, incluindo sintetizar fotos a partir de mapas de etiquetas, gerar fotos coloridas a partir de imagens em preto e branco, transformar fotos do Google Maps em imagens aéreas e até mesmo transformar esboços em fotos.

Neste exemplo, sua rede gerará imagens de fachadas de edifícios usando o banco de dados de fachadas CMP fornecido pelo Center for Machine Perception{:.external} da Universidade Técnica Tcheca em Praga{:.external}. Para ser breve, você usará uma cópia pré-processada{:.external} deste dataset criado pelos autores do pix2pix.

Na cGAN pix2pix, você condiciona as imagens de entrada e gera as imagens de saída correspondentes. As cGANs foram propostas pela primeira vez em Redes Adversarias Gerativas Condicionais (Mirza e Osindero, 2014)

A arquitetura da sua rede conterá:

  • Um gerador com arquitetura baseada em U-Net{:.external}.

  • Um discriminador representado por um classificador convolucional PatchGAN (proposto no artigo pix2pix{:.external}).

Observe que cada época pode levar cerca de 15 segundos em uma única GPU V100.

Abaixo estão alguns exemplos da saída gerada pelo pix2pix cGAN após treinamento por 200 épocas no dataset de fachadas (80k etapas).

sample output_1 sample output_2

Importe o TensorFlow e outras bibliotecas

import tensorflow as tf import os import pathlib import time import datetime from matplotlib import pyplot as plt from IPython import display

Carregue o dataset

Baixe os dados do CMP Facade Database (30 MB). Datasets adicionais estão disponíveis no mesmo formato aqui {:.external}. No Colab você pode selecionar outros datasets no menu suspenso. Observe que alguns dos outros datasets são significativamente maiores (edges2handbags tem 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 imagem original tem tamanho 256 x 512 contendo duas imagens 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)

Você precisa separar as imagens reais da fachada do edifício das imagens do rótulo de arquitetura – todas terão o tamanho 256 x 256 .

Defina uma função que carregue arquivos de imagem e produza dois tensores de imagem:

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

Plote uma amostra das imagens de entrada (imagem do rótulo de arquitetura) e reais (foto da fachada do edifício):

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)

Conforme descrito no artigo pix2pix{:.external}, você precisa aplicar jitter e espelhamento aleatórios para pré-processar o conjunto de treinamento.

Defina várias funções que:

  1. Redimensionem cada imagem de 256 x 256 para uma altura e largura maiores de 286 x 286.

  2. Recortem a imagem para 256 x 256 de forma aleatória.

  3. Virem a imagem aleatoriamente na horizontal, ou seja, da esquerda para a direita (espelhamento aleatório).

  4. Normalizem as imagens para o intervalo [-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

Você pode inspecionar algumas das saídas pré-processadas:

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()

Depois de verificar se o carregamento e o pré-processamento funcionam, vamos definir algumas funções helper que carregam e pré-processam os datasets de treinamento e teste:

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

Crie um pipeline de entrada com 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)

Construa o gerador

O gerador do seu pix2pix cGAN é uma U-Net modificada{:.external}. Uma U-Net consiste em um encoder (downsampler) e um decoder (upsampler). (Você pode descobrir mais sobre ela no tutorial Segmentação de imagens e no site do projeto U-Net{:.external}.)

  • Cada bloco no encoder é: Convolution -> Batch normalization -> Leaky ReLU

  • Cada bloco no decoder é: Transposed convolution -> Batch normalization -> Dropout (aplicado aos primeiros 3 blocos) -> ReLU

  • Existem conexões de salto entre o encoder e o decoder (como no U-Net).

Defina o downsampler (encoder):

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 o upsampler (decoder):

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 o gerador com o downsampler e o 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)

Visualize a arquitetura do modelo do gerador:

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

Teste o gerador:

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

Defina a perda do gerador

As GANs aprendem uma perda que se adapta aos dados, enquanto que as cGANs aprendem uma perda estruturada que penaliza uma possível estrutura diferente da saída da rede e da imagem alvo, conforme descrito no artigo pix2pix{:.external}.

  • A perda do gerador é uma perda de entropia cruzada sigmóide das imagens geradas e de um array de uns.

  • O artigo pix2pix também menciona a perda L1, que é um MAE (erro médio absoluto) entre a imagem gerada e a imagem alvo.

  • Isto permite que a imagem gerada se torne estruturalmente semelhante à imagem alvo.

  • A fórmula para calcular a perda total do gerador é gan_loss + LAMBDA * l1_loss, onde LAMBDA = 100. Este valor foi decidido pelos autores do artigo.

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

O procedimento de treinamento do gerador é o seguinte:

Generator Update Image

Treine o discriminador.

O discriminador na cGAN pix2pix é um classificador convolucional PatchGAN - ele tenta classificar se cada patch de imagem é real ou não, conforme descrito no artigo pix2pix{:.external}.

  • Cada bloco no discriminador é: Convolution -> Batch normalization -> Leaky ReLU.

  • O formato da saída após a última camada é (batch_size, 30, 30, 1).

  • Cada patch de imagem 30 x 30 da saída classifica uma porção 70 x 70 da imagem de entrada.

  • O discriminador recebe duas entradas:

    • A imagem de entrada e a imagem de destino, que deve ser classificada como real.

    • A imagem de entrada e a imagem gerada (saída do gerador), que deve ser classificada como falsa.

    • Use tf.concat([inp, tar], axis=-1) para concatenar essas duas entradas.

Vamos definir o 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)

Visualize a arquitetura do modelo do discriminador:

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

Teste o 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 a perda do discriminador

  • A função discriminator_loss recebe duas entradas: imagens reais e imagens geradas.

  • real_loss é uma perda de entropia cruzada sigmóide das imagens reais e de um array de uns (já que essas são as imagens reais).

  • generated_loss é uma perda de entropia cruzada sigmóide das imagens geradas e um array de zeros (já que essas são as imagens falsas).

  • A total_loss é a soma de real_loss e 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

O procedimento de treinamento para o discriminador é mostrado abaixo.

Para saber mais sobre a arquitetura e os hiperparâmetros você pode consultar o artigo pix2pix{:.external}.

Discriminator Update Image

Defina os otimizadores e um gravador de checkpoints

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)

Gere as imagens

Escreva uma função para plotar algumas imagens durante o treinamento.

  • Passe imagens do dataset de testes para o gerador.

  • O gerador irá então traduzir a imagem de entrada na saída.

  • O último passo é plotar um gráfico das previsões e voilà!

Observação: O training=True é intencional aqui, pois o que você quer são as estatísticas do lote ao executar o modelo no dataset de teste. Se você usar training=False, obterá as estatísticas acumuladas aprendidas no dataset de treinamento (o que não é o que você quer).

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()

Teste a função:

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

Treinamento

  • Para cada exemplo a entrada gera uma saída.

  • O discriminador recebe a input_image e a imagem gerada como primeira entrada. A segunda entrada é input_image e a target_image.

  • Em seguida, calcule o gerador e a perda do discriminador.

  • Depois, calcule os gradientes de perda em relação às variáveis ​​​​do gerador e do discriminador (entradas) e aplique-os ao otimizador.

  • Por fim, registre as perdas no 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)

O ciclo de treinamento propriamente dito. Como este tutorial pode ser executado em mais de um dataset, e os datasets variam muito em tamanho, o loop de treinamento é configurado para funcionar em passos em vez de épocas.

  • Itera sobre o número de passos.

  • A cada 10 passos imprime um ponto (.).

  • A cada 1000 etapas: limpe a tela e execute generate_images para mostrar o progresso.

  • A cada 5000 passos: grave um checkpoint.

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 loop de treinamento salva logs que você pode visualizar no TensorBoard para monitorar o progresso do treinamento.

Se você trabalha numa máquina local, você iniciaria um processo TensorBoard separado. Ao trabalhar num notebook, inicie o visualizador antes de iniciar o treinamento para monitorar com o TensorBoard.

Inicie o TensorBoard Viewer (desculpe, isto não é mostrado no tensorflow.org):

%load_ext tensorboard %tensorboard --logdir {log_dir}

Você pode ver os resultados de uma execução anterior deste notebook em TensorBoard.dev.

Finalmente, execute o loop de treinamento:

fit(train_dataset, test_dataset, steps=40000)

A interpretação dos logs é mais sutil ao treinar uma GAN (ou uma cGAN como pix2pix) em comparação com um modelo simples de classificação ou regressão. Eis algumas coisas que você pode procurar:

  • Verifique que nem o gerador nem o modelo discriminador "ganharam". Se gen_gan_loss ou disc_loss ficar muito baixo, é um indicador de que este modelo está dominando o outro e você não está treinando com sucesso o modelo combinado.

  • O valor log(2) = 0.69 é um bom ponto de referência para estas perdas, pois indica uma perplexidade de 2. O discriminador é, em média, igualmente incerto sobre as duas opções.

  • Para a disc_loss, um valor abaixo de 0.69 significa que o discriminador está se saindo melhor do que o aleatório no conjunto combinado de imagens reais e geradas.

  • Para a gen_gan_loss, um valor abaixo de 0.69 significa que o gerador está se saindo melhor do que o aleatório em enganar o discriminador.

  • À medida que o treinamento avança, gen_l1_loss deve diminuir.

Restaure o checkpoint mais recente e teste a rede

!ls {checkpoint_dir}
# Restoring the latest checkpoint in checkpoint_dir checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

Gere algumas imagens usando o dataset de testes

# 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)