{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Intro\n\nThe goal of this project is to drawings/sketches from photos. This type of task is often called style transfer and popular approach for this task is to use GANs ([review](https://arxiv.org/abs/1705.04058)).\n\nI've decided to start with [CycleGAN](https://arxiv.org/abs/1703.10593). CycleGAN allows for image-to-image translation without paired examples. This makes it fairly easy to use, as we do not need to get/create both the reference output for each image and instead our input are only images to be translated and images showing the desired style.\n\nOther interesting options are [ArtPDGAN](https://link.springer.com/chapter/10.1007/978-3-030-50436-6_21) and [Im2Pencil](https://arxiv.org/pdf/1903.08682.pdf), both of which are directly focused on drawings specifically but require paired images (original and reference).\n\n## Datasets & Implementation\n\nDatasets are from [\"I’m Something of a Painter Myself\"](https://www.kaggle.com/c/gan-getting-started) Kaggle competition and [ImageNet-Sketch](https://github.com/HaohanWang/ImageNet-Sketch).\n\nThe implementation of CycleGAN is based on the \"[Monet CycleGAN Tutorial](https://www.kaggle.com/amyjang/monet-cyclegan-tutorial)\".\n\n\n### Preprocessing\n\nImages are **resized** to 256x256 and **scaled** to <-1;1>. All images are converted to 3 channels for consistency.\n\n## Setup & Running\n\nThe notebook can be run directly on Kaggle without any setup. You can reference Kaggle's [Dockerfile](https://github.com/Kaggle/docker-python/blob/main/Dockerfile.tmpl) for local running. The main dependency is `Tensorflow 2.6.0`.","metadata":{}},{"cell_type":"markdown","source":"# Setup","metadata":{}},{"cell_type":"code","source":"from pathlib import Path\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport tensorflow_addons as tfa\n\n# from kaggle_datasets import KaggleDatasets\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:57:34.964962Z","iopub.execute_input":"2021-11-26T17:57:34.965231Z","iopub.status.idle":"2021-11-26T17:57:34.969696Z","shell.execute_reply.started":"2021-11-26T17:57:34.965201Z","shell.execute_reply":"2021-11-26T17:57:34.968946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:50:51.629367Z","iopub.execute_input":"2021-11-26T17:50:51.629651Z","iopub.status.idle":"2021-11-26T17:50:52.341258Z","shell.execute_reply.started":"2021-11-26T17:50:51.629616Z","shell.execute_reply":"2021-11-26T17:50:52.340490Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Enable TPU if needed.","metadata":{}},{"cell_type":"code","source":"# Not using TPU\n# try:\n#     raise ValueError('disable TPU')\n#     tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n#     print('Device:', tpu.master())\n#     tf.config.experimental_connect_to_cluster(tpu)\n#     tf.tpu.experimental.initialize_tpu_system(tpu)\n#     strategy = tf.distribute.experimental.TPUStrategy(tpu)\n# except:\n#     print('...fail...')\n#     strategy = tf.distribute.get_strategy()\n# print('Number of replicas:', strategy.num_replicas_in_sync)\n\nstrategy = tf.distribute.get_strategy()\nAUTOTUNE = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:50:52.344550Z","iopub.execute_input":"2021-11-26T17:50:52.345024Z","iopub.status.idle":"2021-11-26T17:50:52.354887Z","shell.execute_reply.started":"2021-11-26T17:50:52.344991Z","shell.execute_reply":"2021-11-26T17:50:52.354050Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.__version__","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:50:52.357375Z","iopub.execute_input":"2021-11-26T17:50:52.357901Z","iopub.status.idle":"2021-11-26T17:50:52.366406Z","shell.execute_reply.started":"2021-11-26T17:50:52.357864Z","shell.execute_reply":"2021-11-26T17:50:52.365628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"markdown","source":"## Sketch Dataset","metadata":{}},{"cell_type":"markdown","source":"Dataset is from https://github.com/HaohanWang/ImageNet-Sketch.\n\nMost images have 3 channels, some have 1. Sizes vary with median size of 570x600.","metadata":{}},{"cell_type":"code","source":"img_height, img_width = 256, 256\nchannels = 3","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:50:52.367897Z","iopub.execute_input":"2021-11-26T17:50:52.368405Z","iopub.status.idle":"2021-11-26T17:50:52.375374Z","shell.execute_reply.started":"2021-11-26T17:50:52.368369Z","shell.execute_reply":"2021-11-26T17:50:52.374661Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_img(image):\n    # convert to 3-channel\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = tf.image.resize(image, [img_height, img_width])\n    # scale to <-1,1>\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    return image\n\ndef process_path(file_path):\n    # does not work when using TPU\n    img = tf.io.read_file(file_path)\n    img = decode_img(img)\n    return img","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:50:52.377301Z","iopub.execute_input":"2021-11-26T17:50:52.378404Z","iopub.status.idle":"2021-11-26T17:50:52.385707Z","shell.execute_reply.started":"2021-11-26T17:50:52.378374Z","shell.execute_reply":"2021-11-26T17:50:52.384969Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_ds = tf.data.Dataset.list_files('../input/imagenetsketch/sketch/*/*.JPEG', shuffle=True)\n# batch images so they have the right shape for neural net\nstyled_ds = list_ds.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:50:52.387238Z","iopub.execute_input":"2021-11-26T17:50:52.387734Z","iopub.status.idle":"2021-11-26T17:51:03.895614Z","shell.execute_reply.started":"2021-11-26T17:50:52.387626Z","shell.execute_reply":"2021-11-26T17:51:03.894780Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot(image, ax=None):\n    if ax is None:\n        ax = plt.gca()\n    ax.imshow(image * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:03.896964Z","iopub.execute_input":"2021-11-26T17:51:03.897216Z","iopub.status.idle":"2021-11-26T17:51:03.902284Z","shell.execute_reply.started":"2021-11-26T17:51:03.897181Z","shell.execute_reply":"2021-11-26T17:51:03.901211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fix, axs = plt.subplots(1, 5, figsize=(17,17))\nfor batch, ax in zip(styled_ds.take(5), axs):\n    plot(batch[0], ax)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:03.903528Z","iopub.execute_input":"2021-11-26T17:51:03.904208Z","iopub.status.idle":"2021-11-26T17:51:04.669570Z","shell.execute_reply.started":"2021-11-26T17:51:03.904170Z","shell.execute_reply":"2021-11-26T17:51:04.668874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Photo Dataset\n\nFrom https://www.kaggle.com/c/gan-getting-started.","metadata":{}},{"cell_type":"code","source":"list_ds = tf.data.Dataset.list_files('../input/gan-getting-started/photo_jpg/*.jpg', shuffle=True)\nphoto_ds = list_ds.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:04.672264Z","iopub.execute_input":"2021-11-26T17:51:04.673046Z","iopub.status.idle":"2021-11-26T17:51:06.905712Z","shell.execute_reply.started":"2021-11-26T17:51:04.673010Z","shell.execute_reply":"2021-11-26T17:51:06.904829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fix, axs = plt.subplots(1, 5, figsize=(17,17))\nfor batch, ax in zip(photo_ds.take(5), axs):\n    plot(batch[0], ax)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:06.907247Z","iopub.execute_input":"2021-11-26T17:51:06.907850Z","iopub.status.idle":"2021-11-26T17:51:07.684528Z","shell.execute_reply.started":"2021-11-26T17:51:06.907814Z","shell.execute_reply":"2021-11-26T17:51:07.683575Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Architecture","metadata":{"execution":{"iopub.status.busy":"2021-11-25T15:45:34.917477Z","iopub.execute_input":"2021-11-25T15:45:34.918058Z","iopub.status.idle":"2021-11-25T15:45:34.942867Z","shell.execute_reply.started":"2021-11-25T15:45:34.918023Z","shell.execute_reply":"2021-11-25T15:45:34.942045Z"}}},{"cell_type":"markdown","source":"I'll be using CycleGAN with the UNET architecture. The networks are composed of `downsample` and `upsamble` blocks. Instance normalization is used which is included in Tensorflow Addons.","metadata":{}},{"cell_type":"markdown","source":"## Building Blocks","metadata":{}},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\n\ndef downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2D(filters, size, strides=2, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    result.add(layers.LeakyReLU())\n\n    return result","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:07.686086Z","iopub.execute_input":"2021-11-26T17:51:07.686475Z","iopub.status.idle":"2021-11-26T17:51:07.695509Z","shell.execute_reply.started":"2021-11-26T17:51:07.686443Z","shell.execute_reply":"2021-11-26T17:51:07.694584Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2DTranspose(filters, size, strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n\n    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout:\n        result.add(layers.Dropout(0.5))\n\n    result.add(layers.ReLU())\n\n    return result","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:07.697092Z","iopub.execute_input":"2021-11-26T17:51:07.697925Z","iopub.status.idle":"2021-11-26T17:51:07.708125Z","shell.execute_reply.started":"2021-11-26T17:51:07.697888Z","shell.execute_reply":"2021-11-26T17:51:07.707299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generator\n\nThe generator downsamples and then upsamples the image while using skip connections (UNET architecture).","metadata":{}},{"cell_type":"code","source":"def Generator():\n    inputs = layers.Input(shape=[img_height,img_width,3])\n\n    # bs = batch size, assuming (h, w) = (256, 256)\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False), # (bs, 128, 128, 64)\n        downsample(128, 4), # (bs, 64, 64, 128)\n        downsample(256, 4), # (bs, 32, 32, 256)\n        downsample(512, 4), # (bs, 16, 16, 512)\n        downsample(512, 4), # (bs, 8, 8, 512)\n        downsample(512, 4), # (bs, 4, 4, 512)\n        downsample(512, 4), # (bs, 2, 2, 512)\n        downsample(512, 4), # (bs, 1, 1, 512)\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)\n        upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)\n        upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)\n        upsample(512, 4), # (bs, 16, 16, 1024)\n        upsample(256, 4), # (bs, 32, 32, 512)\n        upsample(128, 4), # (bs, 64, 64, 256)\n        upsample(64, 4), # (bs, 128, 128, 128)\n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n                                  strides=2,\n                                  padding='same',\n                                  kernel_initializer=initializer,\n                                  activation='tanh') # (bs, 256, 256, 3)\n\n    x = inputs\n\n    # Downsampling through the model\n    skips = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n\n    skips = reversed(skips[:-1])\n\n    # Upsampling and establishing the skip connections\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = layers.Concatenate()([x, skip])\n\n    x = last(x)\n\n    return keras.Model(inputs=inputs, outputs=x)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:07.709753Z","iopub.execute_input":"2021-11-26T17:51:07.710156Z","iopub.status.idle":"2021-11-26T17:51:07.726101Z","shell.execute_reply.started":"2021-11-26T17:51:07.710121Z","shell.execute_reply":"2021-11-26T17:51:07.725235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Discriminator\n\nDiscriminator classifies the image as real or generated. The discriminator outputs a smaller 2D image with higher values indicating a real image.","metadata":{}},{"cell_type":"code","source":"def Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    inp = layers.Input(shape=[256, 256, 3], name='input_image')\n\n    x = inp\n\n    down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64)\n    down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)\n    down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)\n\n    zero_pad1 = layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)\n    conv = layers.Conv2D(512, 4, strides=1,\n                         kernel_initializer=initializer,\n                         use_bias=False)(zero_pad1) # (bs, 31, 31, 512)\n\n    norm1 = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(conv)\n\n    leaky_relu = layers.LeakyReLU()(norm1)\n\n    zero_pad2 = layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)\n\n    last = layers.Conv2D(1, 4, strides=1,\n                         kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1)\n\n    return tf.keras.Model(inputs=inp, outputs=last)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:07.736250Z","iopub.execute_input":"2021-11-26T17:51:07.736435Z","iopub.status.idle":"2021-11-26T17:51:07.746861Z","shell.execute_reply.started":"2021-11-26T17:51:07.736414Z","shell.execute_reply":"2021-11-26T17:51:07.746174Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    styled_generator = Generator() # transforms photos to drawings\n    photo_generator = Generator() # transforms drawings to photos\n\n    styled_discriminator = Discriminator() # differentiates real and generated drawings\n    photo_discriminator = Discriminator() # differentiates real and generated photos","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:07.748585Z","iopub.execute_input":"2021-11-26T17:51:07.748904Z","iopub.status.idle":"2021-11-26T17:51:10.409543Z","shell.execute_reply.started":"2021-11-26T17:51:07.748849Z","shell.execute_reply":"2021-11-26T17:51:10.408837Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Sample Run","metadata":{}},{"cell_type":"code","source":"example = next(iter(photo_ds))","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:10.410621Z","iopub.execute_input":"2021-11-26T17:51:10.412585Z","iopub.status.idle":"2021-11-26T17:51:10.465705Z","shell.execute_reply.started":"2021-11-26T17:51:10.412555Z","shell.execute_reply":"2021-11-26T17:51:10.465034Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example.shape","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:10.466801Z","iopub.execute_input":"2021-11-26T17:51:10.467059Z","iopub.status.idle":"2021-11-26T17:51:10.473149Z","shell.execute_reply.started":"2021-11-26T17:51:10.467025Z","shell.execute_reply":"2021-11-26T17:51:10.472303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"styled = styled_generator(example)\n\nplt.subplot(1, 2, 1)\nplt.title(\"Original\")\nplot(example[0])\nplt.show()\n\nplt.subplot(1, 2, 2)\nplt.title(\"Generated\")\nplot(styled[0])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:10.474641Z","iopub.execute_input":"2021-11-26T17:51:10.475379Z","iopub.status.idle":"2021-11-26T17:51:16.997369Z","shell.execute_reply.started":"2021-11-26T17:51:10.475342Z","shell.execute_reply":"2021-11-26T17:51:16.996619Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CycleGAN","metadata":{}},{"cell_type":"code","source":"class CycleGan(keras.Model):\n    def __init__(\n        self,\n        styled_generator,\n        photo_generator,\n        styled_discriminator,\n        photo_discriminator,\n        lambda_cycle=10,\n    ):\n        super(CycleGan, self).__init__()\n        self.s_gen = styled_generator\n        self.p_gen = photo_generator\n        self.s_disc = styled_discriminator\n        self.p_disc = photo_discriminator\n        self.lambda_cycle = lambda_cycle\n        \n    def compile(\n        self,\n        s_gen_optimizer,\n        p_gen_optimizer,\n        s_disc_optimizer,\n        p_disc_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n        cycle_loss_fn,\n        identity_loss_fn\n    ):\n        super(CycleGan, self).compile()\n        self.s_gen_optimizer = s_gen_optimizer\n        self.p_gen_optimizer = p_gen_optimizer\n        self.s_disc_optimizer = s_disc_optimizer\n        self.p_disc_optimizer = p_disc_optimizer\n        self.gen_loss_fn = gen_loss_fn\n        self.disc_loss_fn = disc_loss_fn\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n        \n    def train_step(self, batch_data):\n        # @TODO: rename\n        real_styled, real_photo = batch_data\n        \n        with tf.GradientTape(persistent=True) as tape:\n            # photo to styled back to photo\n            fake_styled = self.s_gen(real_photo, training=True)\n            cycled_photo = self.p_gen(fake_styled, training=True)\n\n            # styled to photo back to styled\n            fake_photo = self.p_gen(real_styled, training=True)\n            cycled_styled = self.s_gen(fake_photo, training=True)\n\n            # generating itself\n            same_styled = self.s_gen(real_styled, training=True)\n            same_photo = self.p_gen(real_photo, training=True)\n\n            # discriminator used to check, inputing real images\n            disc_real_styled = self.s_disc(real_styled, training=True)\n            disc_real_photo = self.p_disc(real_photo, training=True)\n\n            # discriminator used to check, inputing fake images\n            disc_fake_styled = self.s_disc(fake_styled, training=True)\n            disc_fake_photo = self.p_disc(fake_photo, training=True)\n\n            # evaluates generator loss\n            styled_gen_loss = self.gen_loss_fn(disc_fake_styled)\n            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)\n\n            # evaluates total cycle consistency loss\n            total_cycle_loss = self.cycle_loss_fn(real_styled, cycled_styled, self.lambda_cycle) + self.cycle_loss_fn(real_photo, cycled_photo, self.lambda_cycle)\n\n            # evaluates total generator loss\n            total_styled_gen_loss = styled_gen_loss + total_cycle_loss + self.identity_loss_fn(real_styled, same_styled, self.lambda_cycle)\n            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo, self.lambda_cycle)\n\n            # evaluates discriminator loss\n            styled_disc_loss = self.disc_loss_fn(disc_real_styled, disc_fake_styled)\n            photo_disc_loss = self.disc_loss_fn(disc_real_photo, disc_fake_photo)\n\n        # Calculate the gradients for generator and discriminator\n        styled_generator_gradients = tape.gradient(total_styled_gen_loss,\n                                                  self.s_gen.trainable_variables)\n        photo_generator_gradients = tape.gradient(total_photo_gen_loss,\n                                                  self.p_gen.trainable_variables)\n\n        styled_discriminator_gradients = tape.gradient(styled_disc_loss,\n                                                      self.s_disc.trainable_variables)\n        photo_discriminator_gradients = tape.gradient(photo_disc_loss,\n                                                      self.p_disc.trainable_variables)\n\n        # Apply the gradients to the optimizer\n        self.s_gen_optimizer.apply_gradients(zip(styled_generator_gradients,\n                                                 self.s_gen.trainable_variables))\n\n        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients,\n                                                 self.p_gen.trainable_variables))\n\n        self.s_disc_optimizer.apply_gradients(zip(styled_discriminator_gradients,\n                                                  self.s_disc.trainable_variables))\n\n        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients,\n                                                  self.p_disc.trainable_variables))\n        \n        return {\n            \"styled_gen_loss\": total_styled_gen_loss,\n            \"photo_gen_loss\": total_photo_gen_loss,\n            \"styled_disc_loss\": styled_disc_loss,\n            \"photo_disc_loss\": photo_disc_loss\n        }","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:16.999171Z","iopub.execute_input":"2021-11-26T17:51:16.999438Z","iopub.status.idle":"2021-11-26T17:51:17.017734Z","shell.execute_reply.started":"2021-11-26T17:51:16.999401Z","shell.execute_reply":"2021-11-26T17:51:17.016870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss functions","metadata":{"execution":{"iopub.status.busy":"2021-11-25T16:00:37.357996Z","iopub.execute_input":"2021-11-25T16:00:37.35828Z","iopub.status.idle":"2021-11-25T16:00:37.362523Z","shell.execute_reply.started":"2021-11-25T16:00:37.358249Z","shell.execute_reply":"2021-11-25T16:00:37.361787Z"}}},{"cell_type":"markdown","source":"### Discriminator\n\nDiscriminator should predict 1s for real images and 0s for fake images. The discriminator loss is the average of the real and generated loss.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def discriminator_loss(real, generated):\n        real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(real), real)\n\n        generated_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.zeros_like(generated), generated)\n\n        total_disc_loss = real_loss + generated_loss\n\n        return total_disc_loss * 0.5","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.019480Z","iopub.execute_input":"2021-11-26T17:51:17.020031Z","iopub.status.idle":"2021-11-26T17:51:17.031903Z","shell.execute_reply.started":"2021-11-26T17:51:17.019999Z","shell.execute_reply":"2021-11-26T17:51:17.031175Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Generator\n\nThe first loss tells us how well the generator tricked the discriminator. Ideally the discriminator should output only 1s for the generated images.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def generator_loss(generated):\n        return tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated), generated)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.034821Z","iopub.execute_input":"2021-11-26T17:51:17.035116Z","iopub.status.idle":"2021-11-26T17:51:17.041915Z","shell.execute_reply.started":"2021-11-26T17:51:17.035081Z","shell.execute_reply":"2021-11-26T17:51:17.041106Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The original and twice transformed photos should be similar. We calculate the the average of their (absolute) difference as cycle consistency loss. \n","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def calc_cycle_loss(real_image, cycled_image, LAMBDA):\n        loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n        return LAMBDA * loss1","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.043365Z","iopub.execute_input":"2021-11-26T17:51:17.043789Z","iopub.status.idle":"2021-11-26T17:51:17.051306Z","shell.execute_reply.started":"2021-11-26T17:51:17.043730Z","shell.execute_reply":"2021-11-26T17:51:17.050501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" The identity loss compares the input with the output of the generator.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def identity_loss(real_image, same_image, LAMBDA):\n        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n        return LAMBDA * 0.5 * loss","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.052704Z","iopub.execute_input":"2021-11-26T17:51:17.053182Z","iopub.status.idle":"2021-11-26T17:51:17.059831Z","shell.execute_reply.started":"2021-11-26T17:51:17.053149Z","shell.execute_reply":"2021-11-26T17:51:17.059058Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    styled_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\n    styled_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.061285Z","iopub.execute_input":"2021-11-26T17:51:17.061786Z","iopub.status.idle":"2021-11-26T17:51:17.069859Z","shell.execute_reply.started":"2021-11-26T17:51:17.061747Z","shell.execute_reply":"2021-11-26T17:51:17.069161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    cycle_gan_model = CycleGan(\n        styled_generator, photo_generator, styled_discriminator, photo_discriminator\n    )\n\n    cycle_gan_model.compile(\n        s_gen_optimizer = styled_generator_optimizer,\n        p_gen_optimizer = photo_generator_optimizer,\n        s_disc_optimizer = styled_discriminator_optimizer,\n        p_disc_optimizer = photo_discriminator_optimizer,\n        gen_loss_fn = generator_loss,\n        disc_loss_fn = discriminator_loss,\n        cycle_loss_fn = calc_cycle_loss,\n        identity_loss_fn = identity_loss\n    )","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.071351Z","iopub.execute_input":"2021-11-26T17:51:17.071823Z","iopub.status.idle":"2021-11-26T17:51:17.097820Z","shell.execute_reply.started":"2021-11-26T17:51:17.071787Z","shell.execute_reply":"2021-11-26T17:51:17.097213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Preparing zipped dataset\n\nWe need to combine the two datasets so we can provide images in pairs. Datasets also have different sizes and we need to deal with that.","metadata":{}},{"cell_type":"code","source":"len(photo_ds), len(styled_ds)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.099582Z","iopub.execute_input":"2021-11-26T17:51:17.099971Z","iopub.status.idle":"2021-11-26T17:51:17.106509Z","shell.execute_reply.started":"2021-11-26T17:51:17.099941Z","shell.execute_reply":"2021-11-26T17:51:17.105649Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Dataset combinaton example","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:08:02.504247Z","iopub.execute_input":"2021-11-26T09:08:02.504566Z","iopub.status.idle":"2021-11-26T09:08:02.508303Z","shell.execute_reply.started":"2021-11-26T09:08:02.504533Z","shell.execute_reply":"2021-11-26T09:08:02.507579Z"}}},{"cell_type":"code","source":"# in general yo uwant to batch after shuffling but it doesn't matter for batch_size=1\nds_A = tf.data.Dataset.from_tensor_slices(['a', 'b', 'c']).batch(1).shuffle(3, reshuffle_each_iteration=True).repeat()\nds_B = tf.data.Dataset.range(5).batch(1).shuffle(5, reshuffle_each_iteration=True).repeat()\nds_zipped_dummy = tf.data.Dataset.zip((ds_A, ds_B))\n\nlist(ds_zipped_dummy.take(10).as_numpy_iterator())","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.111199Z","iopub.execute_input":"2021-11-26T17:51:17.111598Z","iopub.status.idle":"2021-11-26T17:51:17.133201Z","shell.execute_reply.started":"2021-11-26T17:51:17.111573Z","shell.execute_reply":"2021-11-26T17:51:17.132483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Custom Callbacks","metadata":{}},{"cell_type":"code","source":"class ShowSamplesCallback(keras.callbacks.Callback):\n    def __init__(self, save_to_dir: str):\n        self.save_to_dir = save_to_dir\n    \n    def on_train_begin(self, logs=None):\n        self.original = []\n        self.styled = []\n    \n    def on_train_end(self, logs=None):\n        if self.save_to_dir:\n            dir_ = Path(self.save_to_dir)\n            dir_.mkdir(parents=True, exist_ok=True)\n            for i, (original, styled) in enumerate(zip(self.original, self.styled)):\n                tf.keras.preprocessing.image.save_img(dir_ / f'{i + 1}_original.jpg', original)\n                tf.keras.preprocessing.image.save_img(dir_ / f'{i + 1}_styled.jpg', styled)\n    \n    def on_epoch_end(self, epoch, logs=None):\n        styled_generator = self.model.s_gen\n        # sample predict\n        img_batch = next(iter(photo_ds.take(1)))\n        prediction = styled_generator(img_batch, training=False)[0].numpy()\n        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n        img = img_batch[0]\n        img = (img * 127.5 + 127.5).numpy().astype(np.uint8)\n        self.original.append(img)\n        self.styled.append(prediction)\n        \n        # plot\n        fig, axs = plt.subplots(1, 2, figsize=(6,6), squeeze=True)\n        axs[0].imshow(img)\n        axs[1].imshow(prediction)\n        axs[0].set_title(\"Input Photo [{:02d}]\".format(epoch + 1))\n        axs[1].set_title(\"Generated [{:02d}]\".format(epoch + 1))\n        axs[0].axis(\"off\")\n        axs[1].axis(\"off\")\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.134382Z","iopub.execute_input":"2021-11-26T17:51:17.134648Z","iopub.status.idle":"2021-11-26T17:51:17.146285Z","shell.execute_reply.started":"2021-11-26T17:51:17.134616Z","shell.execute_reply":"2021-11-26T17:51:17.145548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SaveGeneratorCallback(keras.callbacks.Callback):\n    def __init__(self, epochs: int, path: str):\n        self.epochs = epochs\n        self.dir_ = Path(path)\n    \n    def on_epoch_end(self, epoch, logs=None):\n        if epoch % self.epochs == 0:\n            styled_generator = self.model.s_gen\n            path = self.dir_ / f'weights.{epoch + 1}.hdf5'\n            print(f'[SaveGeneratorCallback] Saving weights to {path}.')\n            styled_generator.save_weights(path)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.147509Z","iopub.execute_input":"2021-11-26T17:51:17.148362Z","iopub.status.idle":"2021-11-26T17:51:17.158963Z","shell.execute_reply.started":"2021-11-26T17:51:17.148326Z","shell.execute_reply":"2021-11-26T17:51:17.158167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"HOUR_TO_SECONDS = 3600","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:17.160108Z","iopub.execute_input":"2021-11-26T17:51:17.160916Z","iopub.status.idle":"2021-11-26T17:51:17.167502Z","shell.execute_reply.started":"2021-11-26T17:51:17.160884Z","shell.execute_reply":"2021-11-26T17:51:17.166741Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_buffer_size = 3000\n\nhistory = cycle_gan_model.fit(\n    # zipped dataset - infinitely repeating\n    tf.data.Dataset.zip((\n        # buffer_size should be at least the size of the dataset for uniform shuffling\n        # - https://stackoverflow.com/a/47025850/3936732\n        # - however the dataset is too large for that so we limit the size\n        styled_ds.shuffle(_buffer_size, reshuffle_each_iteration=True).repeat(),\n        photo_ds.shuffle(_buffer_size, reshuffle_each_iteration=True).repeat(),\n    )),\n    \n    steps_per_epoch=1000, # batches per epoch\n    epochs=100,\n    # --- callbacks ---\n    callbacks=[\n        tfa.callbacks.TimeStopping(seconds=HOUR_TO_SECONDS * 6, verbose=1),\n        ShowSamplesCallback('/kaggle/working/imgs'),\n        SaveGeneratorCallback(50, '/kaggle/working'),\n              ],\n)","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:51:18.654570Z","iopub.execute_input":"2021-11-26T17:51:18.654863Z","iopub.status.idle":"2021-11-26T17:52:02.385053Z","shell.execute_reply.started":"2021-11-26T17:51:18.654831Z","shell.execute_reply":"2021-11-26T17:52:02.384157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualize","metadata":{}},{"cell_type":"code","source":"_, ax = plt.subplots(2, 5, figsize=(16, 12))\nfor i, img in enumerate(photo_ds.take(5)):\n    prediction = styled_generator(img, training=False)[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[0, i].imshow(img)\n    ax[1, i].imshow(prediction)\n    ax[0, i].set_title(\"Input Photo\")\n    ax[1, i].set_title(\"Generated\")\n    ax[0, i].axis(\"off\")\n    ax[1, i].axis(\"off\")\nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:53:32.528170Z","iopub.execute_input":"2021-11-26T17:53:32.528620Z","iopub.status.idle":"2021-11-26T17:53:33.740064Z","shell.execute_reply.started":"2021-11-26T17:53:32.528588Z","shell.execute_reply":"2021-11-26T17:53:33.738473Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Saving/Loading","metadata":{}},{"cell_type":"markdown","source":"## Loading weights only (saved with callback)","metadata":{}},{"cell_type":"code","source":"_TEST = False","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:53:37.744885Z","iopub.execute_input":"2021-11-26T17:53:37.745527Z","iopub.status.idle":"2021-11-26T17:53:37.750615Z","shell.execute_reply.started":"2021-11-26T17:53:37.745487Z","shell.execute_reply":"2021-11-26T17:53:37.749808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if _TEST:\n    generator = Generator()\n    generator.load_weights('/kaggle/working/weights.1.hdf5')\n    sample_photo = next(iter(photo_ds))\n    prediction = generator(sample_photo)\n    plot(sample_photo[0])\n    plt.show()\n    plot(prediction[0])\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:53:37.920704Z","iopub.execute_input":"2021-11-26T17:53:37.921705Z","iopub.status.idle":"2021-11-26T17:53:37.927042Z","shell.execute_reply.started":"2021-11-26T17:53:37.921633Z","shell.execute_reply":"2021-11-26T17:53:37.926166Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Final Model","metadata":{}},{"cell_type":"code","source":"!mkdir -p /kaggle/working/saved_model\nstyled_generator.save('/kaggle/working/saved_model/styled_generator')","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:53:38.522401Z","iopub.execute_input":"2021-11-26T17:53:38.523049Z","iopub.status.idle":"2021-11-26T17:53:49.445365Z","shell.execute_reply.started":"2021-11-26T17:53:38.523011Z","shell.execute_reply":"2021-11-26T17:53:49.444523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loaded_styled_generator = tf.keras.models.load_model('/kaggle/working/saved_model/styled_generator')","metadata":{"execution":{"iopub.status.busy":"2021-11-26T17:53:49.471975Z","iopub.execute_input":"2021-11-26T17:53:49.472233Z","iopub.status.idle":"2021-11-26T17:53:54.025852Z","shell.execute_reply.started":"2021-11-26T17:53:49.472203Z","shell.execute_reply":"2021-11-26T17:53:54.025099Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"n_rows = 3 * 2\nn_cols = 5\n_, ax = plt.subplots(n_rows, n_cols, figsize=(n_cols*4, n_rows*4))\nrow = 0\ni = 0\nfor img in photo_ds.take(n_rows//2*n_cols):\n    prediction = loaded_styled_generator(img, training=False)[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[row + 0, i].imshow(img)\n    ax[row + 1, i].imshow(prediction)\n    ax[row + 0, i].set_title(\"Input Photo\")\n    ax[row + 1, i].set_title(\"Generated\")\n    ax[row + 0, i].axis(\"off\")\n    ax[row + 1, i].axis(\"off\")\n    \n    i += 1\n    if i >= n_cols:\n        i = 0\n        row += 2\n    \nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-26T18:01:10.506862Z","iopub.execute_input":"2021-11-26T18:01:10.507424Z","iopub.status.idle":"2021-11-26T18:01:15.125487Z","shell.execute_reply.started":"2021-11-26T18:01:10.507385Z","shell.execute_reply":"2021-11-26T18:01:15.124695Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}