{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n#/kaggle/input/gan-getting-started/monet_jpg\n#/kaggle/input/gan-getting-started/photo_tfrec\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"CycleGAN\nCycleGAN is a model that aims to solve the image-to-image translation problem. The goal of the image-to-image translation problem is to learn the mapping between an input image and an output image using a training set of aligned image pairs. However, obtaining paired examples isn't always feasible. CycleGAN tries to learn this mapping without requiring paired input-output images, using cycle-consistent adversarial networks."},{"metadata":{},"cell_type":"markdown","source":"## Setup"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\n\nimport tensorflow_addons as tfa\nimport tensorflow_datasets as tfds\n\ntfds.disable_progress_bar()\nautotune = tf.data.experimental.AUTOTUNE\nfrom sklearn.model_selection import train_test_split\n\nimport os, random, json, PIL, shutil, re","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Prepare the dataset**"},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_PATH ='/kaggle/input/gan-getting-started'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"MONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\n\nPHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\n\ndef count_data_items(filenames):\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)\n\nn_monet_samples = count_data_items(MONET_FILENAMES)\nn_photo_samples = count_data_items(PHOTO_FILENAMES)\n\n\nprint(f'Monet TFRecord files: {len(MONET_FILENAMES)}')\nprint(f'Monet image files: {n_monet_samples}')\nprint(f'Photo TFRecord files: {len(PHOTO_FILENAMES)}')\nprint(f'Photo image files: {n_photo_samples}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Size of the random crops to be used during training.\ninput_img_size = (256, 256, 3)\n# Weights initializer for the layers.\nkernel_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n# Gamma initializer for instance normalization.\ngamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\nbuffer_size = 256\nbatch_size = 1\n\nIMAGE_SIZE = [256, 256]\n\ndef decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        'image_name': tf.io.FixedLenFeature([], tf.string),\n        'image':      tf.io.FixedLenFeature([], tf.string),\n        'target':     tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create `Dataset` objects"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Apply the preprocessing operations to the training data ### after this STEP IT SHOULD BE 4 dimensional\nAUTO = tf.data.experimental.AUTOTUNE\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO).cache().shuffle(buffer_size).batch(batch_size) \n    return dataset\n\nmonet_ds = load_dataset(MONET_FILENAMES, labeled=True)\nphoto_ds = load_dataset(PHOTO_FILENAMES, labeled=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualize some samples"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_photo = photo_ds\ntrain_monet = monet_ds\n_, ax = plt.subplots(4, 2, figsize=(10, 15))\nfor i, samples in enumerate(zip(train_photo.take(4), train_monet.take(4))):\n    photo = (((samples[0][0] * 127.5) + 127.5).numpy()).astype(np.uint8)\n    monet = (((samples[1][0] * 127.5) + 127.5).numpy()).astype(np.uint8)\n    ax[i, 0].imshow(photo)\n    ax[i, 1].imshow(monet)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Building blocks used in the CycleGAN generators and discriminators"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ReflectionPadding2D(layers.Layer):\n    \"\"\"Implements Reflection Padding as a layer.\n\n    Args:\n        padding(tuple): Amount of padding for the\n        spatial dimensions.\n\n    Returns:\n        A padded tensor with the same type as the input tensor.\n    \"\"\"\n\n    def __init__(self, padding=(1, 1), **kwargs):\n        self.padding = tuple(padding)\n        super(ReflectionPadding2D, self).__init__(**kwargs)\n\n    def call(self, input_tensor, mask=None):\n        padding_width, padding_height = self.padding\n        padding_tensor = [\n            [0, 0],\n            [padding_height, padding_height],\n            [padding_width, padding_width],\n            [0, 0],\n        ]\n        return tf.pad(input_tensor, padding_tensor, mode=\"REFLECT\")\n\n\ndef residual_block(\n    x,\n    activation,\n    kernel_initializer=kernel_init,\n    kernel_size=(3, 3),\n    strides=(1, 1),\n    padding=\"valid\",\n    gamma_initializer=gamma_init,\n    use_bias=False,\n):\n    dim = x.shape[-1]\n    input_tensor = x\n\n    x = ReflectionPadding2D()(input_tensor)\n    x = layers.Conv2D(\n        dim,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = activation(x)\n\n    x = ReflectionPadding2D()(x)\n    x = layers.Conv2D(\n        dim,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = layers.add([input_tensor, x])\n    return x\n\n\ndef downsample(\n    x,\n    filters,\n    activation,\n    kernel_initializer=kernel_init,\n    kernel_size=(3, 3),\n    strides=(2, 2),\n    padding=\"same\",\n    gamma_initializer=gamma_init,\n    use_bias=False,\n):\n    x = layers.Conv2D(\n        filters,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    if activation:\n        x = activation(x)\n    return x\n\n\ndef upsample(\n    x,\n    filters,\n    activation,\n    kernel_size=(3, 3),\n    strides=(2, 2),\n    padding=\"same\",\n    kernel_initializer=kernel_init,\n    gamma_initializer=gamma_init,\n    use_bias=False,\n):\n    x = layers.Conv2DTranspose(\n        filters,\n        kernel_size,\n        strides=strides,\n        padding=padding,\n        kernel_initializer=kernel_initializer,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    if activation:\n        x = activation(x)\n    return x\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build the generators\n\nThe generator consists of downsampling blocks: nine residual blocks\nand upsampling blocks. The structure of the generator is the following:\n\n```\nc7s1-64 ==> Conv block with `relu` activation, filter size of 7\nd128 ====|\n         |-> 2 downsampling blocks\nd256 ====|\nR256 ====|\nR256     |\nR256     |\nR256     |\nR256     |-> 9 residual blocks\nR256     |\nR256     |\nR256     |\nR256 ====|\nu128 ====|\n         |-> 2 upsampling blocks\nu64  ====|\nc7s1-3 => Last conv block with `tanh` activation, filter size of 7.\n```"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_resnet_generator(\n    filters=64,\n    num_downsampling_blocks=2,\n    num_residual_blocks=9,\n    num_upsample_blocks=2,\n    gamma_initializer=gamma_init,\n    name=None,\n):\n    img_input = layers.Input(shape=input_img_size, name=name + \"_img_input\")\n    x = ReflectionPadding2D(padding=(3, 3))(img_input)\n    x = layers.Conv2D(filters, (7, 7), kernel_initializer=kernel_init, use_bias=False)(\n        x\n    )\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = layers.Activation(\"relu\")(x)\n\n    # Downsampling\n    for _ in range(num_downsampling_blocks):\n        filters *= 2\n        x = downsample(x, filters=filters, activation=layers.Activation(\"relu\"))\n\n    # Residual blocks\n    for _ in range(num_residual_blocks):\n        x = residual_block(x, activation=layers.Activation(\"relu\"))\n\n    # Upsampling\n    for _ in range(num_upsample_blocks):\n        filters //= 2\n        x = upsample(x, filters, activation=layers.Activation(\"relu\"))\n\n    # Final block\n    x = ReflectionPadding2D(padding=(3, 3))(x)\n    x = layers.Conv2D(3, (7, 7), padding=\"valid\")(x)\n    x = layers.Activation(\"tanh\")(x)\n\n    model = keras.models.Model(img_input, x, name=name)\n    return model\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build the discriminators\n\nThe discriminators implement the following architecture:\n`C64->C128->C256->C512`"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_discriminator(\n    filters=64, kernel_initializer=kernel_init, num_downsampling=3, name=None\n):\n    img_input = layers.Input(shape=input_img_size, name=name + \"_img_input\")\n    x = layers.Conv2D(\n        filters,\n        (4, 4),\n        strides=(2, 2),\n        padding=\"same\",\n        kernel_initializer=kernel_initializer,\n    )(img_input)\n    x = layers.LeakyReLU(0.2)(x)\n\n    num_filters = filters\n    for num_downsample_block in range(3):\n        num_filters *= 2\n        if num_downsample_block < 2:\n            x = downsample(\n                x,\n                filters=num_filters,\n                activation=layers.LeakyReLU(0.2),\n                kernel_size=(4, 4),\n                strides=(2, 2),\n            )\n        else:\n            x = downsample(\n                x,\n                filters=num_filters,\n                activation=layers.LeakyReLU(0.2),\n                kernel_size=(4, 4),\n                strides=(1, 1),\n            )\n\n    x = layers.Conv2D(\n        1, (4, 4), strides=(1, 1), padding=\"same\", kernel_initializer=kernel_initializer\n    )(x)\n\n    model = keras.models.Model(inputs=img_input, outputs=x, name=name)\n    return model\n\n\n# Get the generators\ngen_G = get_resnet_generator(name=\"generator_G\")\ngen_F = get_resnet_generator(name=\"generator_F\")\n\n# Get the discriminators\ndisc_X = get_discriminator(name=\"discriminator_X\")\ndisc_Y = get_discriminator(name=\"discriminator_Y\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build the CycleGAN model\n\nWe will override the `train_step()` method of the `Model` class\nfor training via `fit()`."},{"metadata":{"trusted":true},"cell_type":"code","source":"class CycleGan(keras.Model):\n    def __init__(\n        self,\n        generator_G,\n        generator_F,\n        discriminator_X,\n        discriminator_Y,\n        lambda_cycle=10.0,\n        lambda_identity=0.5,\n    ):\n        super(CycleGan, self).__init__()\n        self.gen_G = generator_G\n        self.gen_F = generator_F\n        self.disc_X = discriminator_X\n        self.disc_Y = discriminator_Y\n        self.lambda_cycle = lambda_cycle\n        self.lambda_identity = lambda_identity\n\n    def compile(\n        self,\n        gen_G_optimizer,\n        gen_F_optimizer,\n        disc_X_optimizer,\n        disc_Y_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n    ):\n        super(CycleGan, self).compile()\n        self.gen_G_optimizer = gen_G_optimizer\n        self.gen_F_optimizer = gen_F_optimizer\n        self.disc_X_optimizer = disc_X_optimizer\n        self.disc_Y_optimizer = disc_Y_optimizer\n        self.generator_loss_fn = gen_loss_fn\n        self.discriminator_loss_fn = disc_loss_fn\n        self.cycle_loss_fn = keras.losses.MeanAbsoluteError()\n        self.identity_loss_fn = keras.losses.MeanAbsoluteError()\n\n    def train_step(self, batch_data):\n        # x is photo and y is monet\n        real_x, real_y = batch_data\n\n        # For CycleGAN, we need to calculate different\n        # kinds of losses for the generators and discriminators.\n        # We will perform the following steps here:\n        #\n        # 1. Pass real images through the generators and get the generated images\n        # 2. Pass the generated images back to the generators to check if we\n        #    we can predict the original image from the generated image.\n        # 3. Do an identity mapping of the real images using the generators.\n        # 4. Pass the generated images in 1) to the corresponding discriminators.\n        # 5. Calculate the generators total loss (adverserial + cycle + identity)\n        # 6. Calculate the discriminators loss\n        # 7. Update the weights of the generators\n        # 8. Update the weights of the discriminators\n        # 9. Return the losses in a dictionary\n\n        with tf.GradientTape(persistent=True) as tape:\n            # Photo to fake Monet\n            fake_y = self.gen_G(real_x, training=True)\n            # Monet to fake photo -> y2x\n            fake_x = self.gen_F(real_y, training=True)\n\n            # Cycle (Photo to fake monet to fake photo): x -> y -> x\n            cycled_x = self.gen_F(fake_y, training=True)\n            # Cycle (Monet to fake photo to fake monet) y -> x -> y\n            cycled_y = self.gen_G(fake_x, training=True)\n\n            # Identity mapping\n            same_x = self.gen_F(real_x, training=True)\n            same_y = self.gen_G(real_y, training=True)\n\n            # Discriminator output\n            disc_real_x = self.disc_X(real_x, training=True)\n            disc_fake_x = self.disc_X(fake_x, training=True)\n\n            disc_real_y = self.disc_Y(real_y, training=True)\n            disc_fake_y = self.disc_Y(fake_y, training=True)\n\n            # Generator adverserial loss\n            gen_G_loss = self.generator_loss_fn(disc_fake_y)\n            gen_F_loss = self.generator_loss_fn(disc_fake_x)\n\n            # Generator cycle loss\n            cycle_loss_G = self.cycle_loss_fn(real_y, cycled_y) * self.lambda_cycle\n            cycle_loss_F = self.cycle_loss_fn(real_x, cycled_x) * self.lambda_cycle\n\n            # Generator identity loss\n            id_loss_G = (\n                self.identity_loss_fn(real_y, same_y)\n                * self.lambda_cycle\n                * self.lambda_identity\n            )\n            id_loss_F = (\n                self.identity_loss_fn(real_x, same_x)\n                * self.lambda_cycle\n                * self.lambda_identity\n            )\n\n            # Total generator loss\n            total_loss_G = gen_G_loss + cycle_loss_G + id_loss_G\n            total_loss_F = gen_F_loss + cycle_loss_F + id_loss_F\n\n            # Discriminator loss\n            disc_X_loss = self.discriminator_loss_fn(disc_real_x, disc_fake_x)\n            disc_Y_loss = self.discriminator_loss_fn(disc_real_y, disc_fake_y)\n\n        # Get the gradients for the generators\n        grads_G = tape.gradient(total_loss_G, self.gen_G.trainable_variables)\n        grads_F = tape.gradient(total_loss_F, self.gen_F.trainable_variables)\n\n        # Get the gradients for the discriminators\n        disc_X_grads = tape.gradient(disc_X_loss, self.disc_X.trainable_variables)\n        disc_Y_grads = tape.gradient(disc_Y_loss, self.disc_Y.trainable_variables)\n\n        # Update the weights of the generators\n        self.gen_G_optimizer.apply_gradients(\n            zip(grads_G, self.gen_G.trainable_variables)\n        )\n        self.gen_F_optimizer.apply_gradients(\n            zip(grads_F, self.gen_F.trainable_variables)\n        )\n\n        # Update the weights of the discriminators\n        self.disc_X_optimizer.apply_gradients(\n            zip(disc_X_grads, self.disc_X.trainable_variables)\n        )\n        self.disc_Y_optimizer.apply_gradients(\n            zip(disc_Y_grads, self.disc_Y.trainable_variables)\n        )\n\n        return {\n            \"G_loss\": total_loss_G,\n            \"F_loss\": total_loss_F,\n            \"D_X_loss\": disc_X_loss,\n            \"D_Y_loss\": disc_Y_loss,\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create a callback that periodically saves generated images"},{"metadata":{"trusted":true},"cell_type":"code","source":"class GANMonitor(keras.callbacks.Callback):\n    \"\"\"A callback to generate and save images after each epoch\"\"\"\n\n    def __init__(self, num_img=4):\n        self.num_img = num_img\n\n    def on_epoch_end(self, epoch, logs=None):\n        _, ax = plt.subplots(4, 2, figsize=(12, 12))\n        for i, img in enumerate(train_photo.take(self.num_img)):\n            prediction = self.model.gen_G(img)[0].numpy()\n            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n            img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n            ax[i, 0].imshow(img)\n            ax[i, 1].imshow(prediction)\n            ax[i, 0].set_title(\"Input image\")\n            ax[i, 1].set_title(\"Translated image\")\n            ax[i, 0].axis(\"off\")\n            ax[i, 1].axis(\"off\")\n\n            prediction = keras.preprocessing.image.array_to_img(prediction)\n            prediction.save(\n                \"generated_img_{i}_{epoch}.png\".format(i=i, epoch=epoch + 1)\n            )\n        plt.show()\n        plt.close()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train the end-to-end model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Loss function for evaluating adversarial loss\nadv_loss_fn = keras.losses.MeanSquaredError()\n\n# Define the loss function for the generators\ndef generator_loss_fn(fake):\n    fake_loss = adv_loss_fn(tf.ones_like(fake), fake)\n    return fake_loss\n\n\n# Define the loss function for the discriminators\ndef discriminator_loss_fn(real, fake):\n    real_loss = adv_loss_fn(tf.ones_like(real), real)\n    fake_loss = adv_loss_fn(tf.zeros_like(fake), fake)\n    return (real_loss + fake_loss) * 0.5\n\n\n# Create cycle gan model\ncycle_gan_model = CycleGan(\n    generator_G=gen_G, generator_F=gen_F, discriminator_X=disc_X, discriminator_Y=disc_Y\n)\n\n# Compile the model\ncycle_gan_model.compile(\n    gen_G_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    gen_F_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    disc_X_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    disc_Y_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    gen_loss_fn=generator_loss_fn,\n    disc_loss_fn=discriminator_loss_fn,\n)\n# Callbacks\nplotter = GANMonitor()\ncheckpoint_filepath = \"./model_checkpoints/cyclegan_checkpoints.{epoch:03d}\"\nmodel_checkpoint_callback = keras.callbacks.ModelCheckpoint(\n    filepath=checkpoint_filepath\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Here we will train the model for just one epoch as each epoch takes around\n# 7 minutes on a single P100 backed machine.\ncycle_gan_model.fit(\n    tf.data.Dataset.zip((train_photo, train_monet)),\n    epochs=100,\n    callbacks=[plotter, model_checkpoint_callback],\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create submission file"},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\n! rm -rf ../images\n! mkdir ../images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"photo_ds_unbatch = load_dataset(PHOTO_FILENAMES).take(7038)\n\ni = 1\nfor img in photo_ds_unbatch:\n    prediction = cycle_gan_model.gen_G(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    im = PIL.Image.fromarray(prediction)\n    im.save(\"../images/\" + str(i) + \".jpg\")\n    print(f'{i} images saved')\n    i += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}