content: # Copyright 2021 Google LLC
  #
  # 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
  #
  #      http://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.
  """GAN for generating CJK characters.
  
  The vast majority of this code is adapted from the pix2pix GAN described in
  https://www.tensorflow.org/tutorials/generative/pix2pix. Changes include the
  specific tensor dimensions, some tuning of magic numbers, and some changes to
  loss functions.
  
  TODO(ambuc): This file has type annotations because they're useful for a human
  reader, but the build system doesn't yet enforce them with a strictly-typed
  python build rule.
  """
  
  import time
  from typing import List, Text, Tuple
  
  from IPython import display
  import matplotlib.pyplot as plt
  import tensorflow as tf
  
  _LAMBDA = 100
  
  
  def _load_image(filename: Text) -> List[List[tf.Tensor]]:
      """Given the filename of a PNG, returns a list of three tensors: a, b, a+b.
  
    Args:
      filename: Path to a file. The file must be a PNG and greyscale and 256x256.
  
    Returns:
      A list of tensors: a, b, and a+b.
  
    """
      image = tf.io.read_file(filename)
      image = tf.image.decode_png(image, channels=1)  # greyscale
  
      # Our images have a width which is divisible by three.
      w = tf.shape(image)[1] // 3
  
      return [
          tf.cast(image[:, n * w:(n + 1) * w, :], tf.float32) for n in range(3)
      ]
  
  
  def make_datasets(files_glob: Text) -> Tuple[tf.data.Dataset, tf.data.Dataset]:
      """Makes the train/test datasets.
  
    Args:
      files_glob: A glob (like "/tmp/folder/*.png") of all the input images.
  
    Returns:
      A pair of train, test datasets of type tf.data.Dataset.
  
    """
      ds = tf.data.Dataset.list_files(files_glob).map(
          _load_image, num_parallel_calls=tf.data.AUTOTUNE).shuffle(400).batch(1)
  
      train_dataset_a = ds.shard(num_shards=3, index=0)
      train_dataset_b = ds.shard(num_shards=3, index=1)
      train_ds = train_dataset_a.concatenate(train_dataset_b)
  
      test_ds = ds.shard(num_shards=3, index=2)
  
      return train_ds, test_ds
  
  
  def _downsample(filters: int,
                  size: int,
                  apply_batchnorm: bool = True) -> tf.keras.Sequential:
      """Downsampler from https://www.tensorflow.org/tutorials/generative/pix2pix#build_the_generator.
  
    Args:
      filters: The number of filters.
      size: The size of the input tensor width at this step.
      apply_batchnorm: Whether or not to apply batch normalization. Probably
        should be false on the input layer, and true elsewhere.
  
    Returns:
      A sequential model.
    """
      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
  
  
  def _upsample(filters: int,
                size: int,
                apply_dropout: bool = False) -> tf.keras.Sequential:
      """Upsampler from https://www.tensorflow.org/tutorials/generative/pix2pix#build_the_generator.
  
    Args:
      filters: The number of filters.
      size: The size of the input tensor width at this step.
      apply_dropout: Whether or not to apply dropout. Probably should be true for
        the first few layers and false elsewhere.
  
    Returns:
      A sequential model.
  
    """
      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
  
  
  def make_generator() -> tf.keras.Model:
      """Creates a generator.
  
    99% of this is copied directly from
    https://www.tensorflow.org/tutorials/generative/pix2pix#build_the_generator,
    except for the input shape (now two channels, two greyscale images instead of
    one RGB image) and output shape (one channel, one greyscale image instead of
    one RGB image).
  
    Returns:
      a tf.keras.Model which returns a 256x256x1 tensor.
    """
      inputs = tf.keras.layers.Input(shape=[256, 256, 2])
  
      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)
      ]
  
      x = inputs
  
      skips = []
      for down in [
              _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)
      ]:
          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])
  
      last = tf.keras.layers.Conv2DTranspose(
          1,  # one output channel, i.e. greyscale
          4,
          strides=2,
          padding='same',
          kernel_initializer=tf.random_normal_initializer(0., 0.02),
          activation='tanh')  # (bs, 256, 256, 3)
  
      x = last(x)
  
      return tf.keras.Model(inputs=inputs, outputs=x)
  
  
  def generator_loss(loss_object: tf.keras.losses.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
  
  
  def make_discriminator() -> tf.keras.Model:
      """Returns a discriminator.
  
    This is 99% the same as
    https://www.tensorflow.org/tutorials/generative/pix2pix#build_the_discriminator,
    except that the shape of the input and output tensors are different.
  
    Returns:
      A tf.keras.model which accepts a 256x256x2 tensor and compares it to a
      target 256x256x1 tensor.
    """
      initializer = tf.random_normal_initializer(0., 0.02)
  
      input_img = tf.keras.layers.Input(shape=[256, 256, 2], name='input_image')
      target_img = tf.keras.layers.Input(shape=[256, 256, 1],
                                         name='target_image')
  
      x = tf.keras.layers.concatenate([input_img,
                                       target_img])  # (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=[input_img, target_img], outputs=last)
  
  
  def discriminator_loss(loss_object: tf.keras.losses.Loss, disc_real_output,
                         disc_generated_output) -> float:
      """Returns discriminator loss.
  
    100% the same as
    https://www.tensorflow.org/tutorials/generative/pix2pix#build_the_discriminator.
  
    Args:
      loss_object: A reusable loss_object of type
        tf.keras.losses.BinaryCrossentropy.
      disc_real_output: A set of real images.
      disc_generated_output: A set of generator images.
  
    Returns:
      The sum of some loss functions.
    """
      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)
  
      return real_loss + generated_loss
  
  
  def generate_images(model: tf.keras.Model, input_a: tf.Tensor,
                      input_b: tf.Tensor, target: tf.Tensor) -> None:
      """In Colab, prints [a | b | real(a,b) | predicted(a,b)] to the display.
  
    Args:
      model: The generator to use.
      input_a: the LHS image.
      input_b: the RHS image.
      target: The real(a,b) composition.
    """
      x = tf.concat([input_a, input_b], 3)
      x = tf.reshape(x, [256, 256, 2])
      prediction = model(x[tf.newaxis, ...], training=True)
  
      images = [input_a[0], input_b[0], target[0], prediction[0]]
      fig, axes = plt.subplots(1, 4)
      titles = [
          'Input Image A', 'Input Image B', 'Ground Truth', 'Predicted Image'
      ]
  
      for image, axis, title in zip(images, axes, titles):
          axis.set_title(title)
          axis.imshow(image[:, :, 0])
          axis.axis('off')
      fig.show()
  
  
  @tf.function
  def train_step(generator: tf.keras.Model,
                 generator_optimizer: tf.keras.optimizers.Optimizer,
                 discriminator: tf.keras.Model,
                 discriminator_optimizer: tf.keras.optimizers.Optimizer,
                 loss_object: tf.keras.losses.Loss, inp_a: tf.Tensor,
                 inp_b: tf.Tensor, target: tf.Tensor, epoch: int,
                 summary_writer: tf.summary.SummaryWriter) -> None:
      """Trains the models for one (1) epoch.
  
    See https://www.tensorflow.org/tutorials/generative/pix2pix#training.
  
    Args:
      generator: A generator model,
      generator_optimizer: and an optimizer for the generator.
      discriminator: A discriminator model,
      discriminator_optimizer: and an optimizer for the generator.
      loss_object: A reusable BinaryCrossentropy object.
      inp_a: A full-width image of the left-most component.
      inp_b: A full-width image of the right-most component.
      target: The human-authored image of the a+b character.
      epoch: The index of the epoch we're in.
      summary_writer: A SummaryWriter object for writing.... summaries.
    """
      with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
          inp_x = tf.concat([inp_a, inp_b], 3)
          gen_output = generator(inp_x, training=True)
  
          disc_real_output = discriminator([inp_x, target], training=True)
          disc_generated_output = discriminator([inp_x, gen_output],
                                                training=True)
  
          gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(
              loss_object, disc_generated_output, gen_output, target)
          disc_loss = discriminator_loss(loss_object, disc_real_output,
                                         disc_generated_output)
  
      # TODO(ambuc): Should this simply be gen_l1_loss?
      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)
  
  
  def fit(generator: tf.keras.Model,
          generator_optimizer: tf.keras.optimizers.Optimizer,
          discriminator: tf.keras.Model,
          discriminator_optimizer: tf.keras.optimizers.Optimizer,
          loss_object: tf.keras.losses.Loss, train_ds: tf.data.Dataset,
          epochs: int, test_ds: tf.data.Dataset, checkpoint: tf.train.Checkpoint,
          checkpoint_prefix: Text,
          summary_writer: tf.summary.SummaryWriter) -> None:
      """Runs for |epochs| and trains the models.
  
    Args:
      generator: A generator model,
      generator_optimizer: and an optimizer for the generator.
      discriminator: A discriminator model,
      discriminator_optimizer: and an optimizer for the generator.
      loss_object: A reusable BinaryCrossentropy object.
      train_ds:
      epochs: The number of epochs to train for.
      test_ds:
      checkpoint:
      checkpoint_prefix:
      summary_writer: A SummaryWriter object for writing.... summaries.
    """
      for epoch in range(epochs):
          start = time.time()
  
          display.clear_output(wait=True)
  
          for a, b, ab in test_ds.take(1):
              generate_images(generator, a, b, ab)
          print('Epoch: ', epoch)
  
          for n, (inp_a, inp_b, target) in train_ds.enumerate():
              print('.', end='')
              if (n + 1) % 100 == 0:
                  print()
              train_step(generator, generator_optimizer, discriminator,
                         discriminator_optimizer, loss_object, inp_a, inp_b,
                         target, epoch, summary_writer)
          print()
  
          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)
  