#@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.

"""
Explanation: Copyright 2019 The TensorFlow Authors.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
"""


import tensorflow as tf

import os
import time

from matplotlib import pyplot as plt
from IPython import display

!pip install -U tensorboard

"""
Explanation: Pix2Pix
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://tensorflow.google.cn/tutorials/generative/pix2pix">
    <img src="https://tensorflow.google.cn/images/tf_logo_32px.png" />
    在 tensorFlow.google.cn 上查看</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/generative/pix2pix.ipynb">
    <img src="https://tensorflow.google.cn/images/colab_logo_32px.png" />
    在 Google Colab 中运行</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/generative/pix2pix.ipynb">
    <img src="https://tensorflow.google.cn/images/GitHub-Mark-32px.png" />
    在 GitHub 上查看源代码</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/tutorials/generative/pix2pix.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png" />下载 notebook</a>
  </td>
</table>

此笔记本将演示如何使用条件 GAN 进行图像到图像转换，如 Image-to-Image Translation with Conditional Adversarial Networks 中所述。利用这种技术，我们可以将黑白照片转换成彩色照片，将 Google 地图转换为 Google 地球等。在这里，我们将建筑物的外立面转换为实际建筑物。
在示例中，我们将使用 CMP Facade 数据库，该数据库由布拉格捷克理工大学的机器感知中心提供。为了使我们的示例简短，我们将使用此数据集（由上述论文作者创建）的预处理副本。
在单个 V100 GPU 上，每个周期大约需要 15 秒。
下面是将模型训练 200 个周期后生成的输出。
 
导入 TensorFlow 和其他库
End of explanation
"""


_URL = 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz'

path_to_zip = tf.keras.utils.get_file('facades.tar.gz',
                                      origin=_URL,
                                      extract=True)

PATH = os.path.join(os.path.dirname(path_to_zip), 'facades/')

BUFFER_SIZE = 400
BATCH_SIZE = 1
IMG_WIDTH = 256
IMG_HEIGHT = 256

def load(image_file):
  image = tf.io.read_file(image_file)
  image = tf.image.decode_jpeg(image)

  w = tf.shape(image)[1]

  w = w // 2
  real_image = image[:, :w, :]
  input_image = image[:, w:, :]

  input_image = tf.cast(input_image, tf.float32)
  real_image = tf.cast(real_image, tf.float32)

  return input_image, real_image

inp, re = load(PATH+'train/100.jpg')
# casting to int for matplotlib to show the image
plt.figure()
plt.imshow(inp/255.0)
plt.figure()
plt.imshow(re/255.0)

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 286 x 286 x 3
  input_image, real_image = resize(input_image, real_image, 286, 286)

  # randomly cropping to 256 x 256 x 3
  input_image, real_image = random_crop(input_image, real_image)

  if tf.random.uniform(()) &gt; 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

"""
Explanation: 加载数据集
您可以从此处下载该数据集和类似数据集。如论文中所述，我们将随机抖动和镜像应用于训练数据集。

在随机抖动中，图像的大小被调整为 286 x 286，然后随机裁切为 256 x 256
在随机镜像中，图像被水平（即从左到右）随机翻转。
End of explanation
"""


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

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

"""
Explanation: 正如您在下面的图像中看到的那样，它们正在经历论文中介绍的随机抖动，目的为：

将图像调整为更大的高度和宽度
随机裁切为目标大小
随机水平翻转图像
End of explanation
"""


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

test_dataset = tf.data.Dataset.list_files(PATH+'test/*.jpg')
test_dataset = test_dataset.map(load_image_test)
test_dataset = test_dataset.batch(BATCH_SIZE)

"""
Explanation: 输入流水线
End of explanation
"""


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)

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)

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

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

  up_stack = [
    upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)
    upsample(512, 4), # (bs, 16, 16, 1024)
    upsample(256, 4), # (bs, 32, 32, 512)
    upsample(128, 4), # (bs, 64, 64, 256)
    upsample(64, 4), # (bs, 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') # (bs, 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)

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

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

"""
Explanation: 构建生成器

生成器的架构是经过修改的 U-Net。
编码器中的每个块为 (Conv -> Batchnorm -> Leaky ReLU)
解码器中的每个块为（Transposed Conv -> Batchnorm -> Dropout（应用于前三个块）-> ReLU）
编码器和解码器之间存在跳跃连接（如在 U-Net 中）。
End of explanation
"""


LAMBDA = 100

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

"""
Explanation: 生成器损失
它是所生成图像和一个图像数组的 sigmoid 交叉熵损失。
论文还包括 L1 损失，它是所生成图像与目标图像之间的 MAE（平均绝对误差）。
这样可使生成的图像在结构上与目标图像相似。
计算总生成器损失的公式 = gan_loss + LAMBDA * l1_loss，其中 LAMBDA =100。该值由论文作者决定。



生成器的训练过程如下所示：
End of explanation
"""


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]) # (bs, 256, 256, channels*2)

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

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

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

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

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

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

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

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

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

"""
Explanation: 构建判别器

判别器是一个 PatchGAN。
判别器中的每个块为 (Conv -> BatchNorm -> Leaky ReLU)
最后一层之后的输出形状为 (batch_size, 30, 30, 1)
输出的每个 30x30 图像块将输入图像的 70x70 区域分类（这种架构称为 PatchGAN）。
判别器接收 2 个输入。
输入图像和目标图像，应分类为实际图像。
输入图像和所生成图像（生成器的输出），应分类为伪图像。
我们在代码 (tf.concat([inp, tar], axis=-1)) 中将这 2 个输入串联
End of explanation
"""


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

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

"""
Explanation: 判别器损失

判别器损失函数使用 2 个输入：实际图像和生成图像
real_loss 是实际图像和一个图像数组（因为这些是实际图像）的 sigmoid 交叉熵损失
generated_loss 是生成图像和一个零数组（因为这些是伪图像）的 sigmoid 交叉熵损失
接着，total_loss 是 real_loss 和 generate_loss 的和
End of explanation
"""


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)

"""
Explanation: 判别器的训练过程如下所示。
要详细了解架构和超参数，请参阅论文。

定义优化器和检查点 saver
End of explanation
"""


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 between [0, 1] to plot it.
    plt.imshow(display_list[i] * 0.5 + 0.5)
    plt.axis('off')
  plt.show()

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

"""
Explanation: 生成图像
编写函数以在训练期间绘制一些图像。

我们将图像从测试数据集传递到生成器。
然后，生成器会将输入图像转换为输出。
最后一步是绘制预测，瞧！

注：在这里，training=True 是有意的，因为在基于测试数据集运行模型时，我们需要批次统计信息。如果我们使用 training = False，将获得从训练数据集中学习的累积统计信息（我们不需要）。
End of explanation
"""


EPOCHS = 150

import datetime
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, epoch):
  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=epoch)
    tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch)
    tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch)
    tf.summary.scalar('disc_loss', disc_loss, step=epoch)

"""
Explanation: 训练

为每个示例输入生成一个输出。
判别器接收 input_image 和生成的图像作为第一个输入。第二个输入为 input_image 和 target_image。
接下来，我们计算生成器和判别器损失。
随后，我们计算损失相对于生成器和判别器变量（输入）的梯度，并将其应用于优化器。
随后将损失记录到 TensorBoard。
End of explanation
"""


def fit(train_ds, epochs, test_ds):
  for epoch in range(epochs):
    start = time.time()

    display.clear_output(wait=True)

    for example_input, example_target in test_ds.take(1):
      generate_images(generator, example_input, example_target)
    print("Epoch: ", epoch)

    # Train
    for n, (input_image, target) in train_ds.enumerate():
      print('.', end='')
      if (n+1) % 100 == 0:
        print()
      train_step(input_image, target, epoch)
    print()

    # saving (checkpoint) the model every 20 epochs
    if (epoch + 1) % 20 == 0:
      checkpoint.save(file_prefix = checkpoint_prefix)

    print ('Time taken for epoch {} is {} sec\n'.format(epoch + 1,
                                                        time.time()-start))
  checkpoint.save(file_prefix = checkpoint_prefix)

"""
Explanation: 实际训练循环：

迭代周期数。
在每个周期，它都会清除显示，并运行 generate_images 来显示进度。
在每个周期，它都会迭代训练数据集，同时为每个样本打印 '.'。
它每 20 个周期保存一个检查点。
End of explanation
"""


#docs_infra: no_execute
%load_ext tensorboard
%tensorboard --logdir {log_dir}

"""
Explanation: 此训练循环会保存日志，您可以在 TensorBoard 中轻松查看日志来监视训练进度。在本地工作时，您将启动一个单独的 Tensorboard 进程。在笔记本中，如果要使用 TensorBoard 进行监视，最简单的方法是在开始训练之前启动查看器。
要启动查看器，请将以下代码粘贴到代码单元中：
End of explanation
"""


fit(train_dataset, EPOCHS, test_dataset)

"""
Explanation: 现在运行训练循环：
End of explanation
"""


display.IFrame(
    src="https://tensorboard.dev/experiment/lZ0C6FONROaUMfjYkVyJqw",
    width="100%",
    height="1000px")

"""
Explanation: 如果要公开共享 TensorBoard 结果，可以通过将以下代码复制到代码单元中的方式将日志上传到 TensorBoard.dev。
注：此操作需要一个 Google 帐号。
!tensorboard dev upload --logdir  {log_dir}
小心：此命令不会终止。它可以连续上传长时间运行实验的结果。数据上传后，您需要使用笔记本工具中的“Interrupt Execution”选项将其停止。
您可以在 TensorBoard.dev 上查看此笔记本先前运行的结果。
TensorBoard.dev 是一种托管式体验，用于托管、跟踪 ML 实验并与所有人共享。
也可以使用 &lt;iframe&gt; 将其包含在行内：
End of explanation
"""


!ls {checkpoint_dir}

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

"""
Explanation: 与简单的分类或回归模型相比，从 GAN 解释日志更为微妙。要检查的内容包括：

检查两个模型是否均未“获胜”。如果 gen_gan_loss 或 disc_loss 变得很低，则表明此模型正在支配另一个模型，并且您未能成功训练组合模型。
值 log(2) = 0.69 是这些损失的一个适当参考点，因为它表示困惑度为 2：判别器对这两个选项的平均不确定性是相等的。
对于 disc_loss，如果值小于 0.69，则表示判别器对实际 + 生成图像组合的效果要优于随机数。
对于 gen_gan_loss，如果值小于 0.69，则表示生成器在欺骗判别器方面的表现要优于随机数。
随着训练的进行，gen_l1_loss 应当下降。

恢复最新的检查点并进行测试
End of explanation
"""


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

"""
Explanation: 使用测试数据集生成
End of explanation
"""


