{"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. Images are optionally converted to grayscale.\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":"# Version\n\nThis version allows better selection of parameters, datasets and preprocessing. History saving was fixed.\n\n<a href=\"#training_results\"><b>Go to training results.</b></a>","metadata":{}},{"cell_type":"code","source":"# convert training images to grayscale\nTO_GRAYSCALE = False\n# debug flag for testing the pipeline\nTEST_PIPELINE_ONLY = False\n# early stopping patience, set to -1 to disable\nEARLY_STOPPING_PATIENCE = -1\n# dataset selection\nSOURCE_DATASET = 'landscapes' # landscapes, animal-faces\nTARGET_DATASET = 'sketches' # monet, sketches\n# Adam parameters\nLR = 2e-4\nBETA_1 = 0.5\n# training\nepochs = 300\nsteps_per_epoch = 400\n\n# show image at the end of each epoch\nSHOW_SAMPLES_DURING_TRAINING = True\n# save generator each #n epochs, set to -1 to disable\nPERIODIC_SAVES = -1 \n\n# Time limits\n# - safe limit on Kaggle is about HOUR_TO_SECONDS * 6\n# - it takes about 30 minutes to run 17 epochs at 400 steps per epoch\nHOUR_TO_SECONDS = 3600\nMINUTE_TO_SECONDS = 60\nTIME_LIMIT = HOUR_TO_SECONDS * 6","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:54:55.917315Z","iopub.execute_input":"2021-12-28T17:54:55.917673Z","iopub.status.idle":"2021-12-28T17:54:56.018841Z","shell.execute_reply.started":"2021-12-28T17:54:55.917599Z","shell.execute_reply":"2021-12-28T17:54:56.018159Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-12-28T17:54:56.020423Z","iopub.execute_input":"2021-12-28T17:54:56.020736Z","iopub.status.idle":"2021-12-28T17:55:00.273391Z","shell.execute_reply.started":"2021-12-28T17:54:56.0207Z","shell.execute_reply":"2021-12-28T17:55:00.272512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:00.274772Z","iopub.execute_input":"2021-12-28T17:55:00.275054Z","iopub.status.idle":"2021-12-28T17:55:01.110167Z","shell.execute_reply.started":"2021-12-28T17:55:00.275021Z","shell.execute_reply":"2021-12-28T17:55:01.109318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"AUTOTUNE = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:01.11315Z","iopub.execute_input":"2021-12-28T17:55:01.113671Z","iopub.status.idle":"2021-12-28T17:55:01.117948Z","shell.execute_reply.started":"2021-12-28T17:55:01.113629Z","shell.execute_reply":"2021-12-28T17:55:01.117229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.__version__","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:01.119437Z","iopub.execute_input":"2021-12-28T17:55:01.11997Z","iopub.status.idle":"2021-12-28T17:55:01.133575Z","shell.execute_reply.started":"2021-12-28T17:55:01.119919Z","shell.execute_reply":"2021-12-28T17:55:01.132684Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"markdown","source":"## Target Dataset","metadata":{}},{"cell_type":"markdown","source":"Sketch 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.\n\nMonet dataset is from https://www.kaggle.com/c/gan-getting-started.","metadata":{}},{"cell_type":"code","source":"if TARGET_DATASET == 'sketches':\n    target_dir = '/kaggle/input/imagenetsketch/sketch/*/*.JPEG'\nelif TARGET_DATASET == 'monet':\n    target_dir = '/kaggle/input/gan-getting-started/monet_jpg/*.jpg'\nelse:\n    raise ValueError('Invalid dataset name!')\nprint(target_dir)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:01.1353Z","iopub.execute_input":"2021-12-28T17:55:01.136124Z","iopub.status.idle":"2021-12-28T17:55:01.143185Z","shell.execute_reply.started":"2021-12-28T17:55:01.13609Z","shell.execute_reply":"2021-12-28T17:55:01.142357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_height, img_width = 256, 256\nchannels = 3","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:01.144599Z","iopub.execute_input":"2021-12-28T17:55:01.144976Z","iopub.status.idle":"2021-12-28T17:55:01.151193Z","shell.execute_reply.started":"2021-12-28T17:55:01.144922Z","shell.execute_reply":"2021-12-28T17:55:01.150508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_img(image, to_gray=False):\n    # convert to 3-channel\n    image = tf.image.decode_jpeg(image, channels=3)\n    if to_gray:\n        image = tf.image.rgb_to_grayscale(image)\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_grayscale(file_path):\n    img = tf.io.read_file(file_path)\n    img = decode_img(img, to_gray=True)\n    return img\n\ndef process_path(file_path):\n    img = tf.io.read_file(file_path)\n    img = decode_img(img)\n    return img\n\nif TO_GRAYSCALE:\n    process_fun = process_path_grayscale\nelse:\n    process_fun = process_path","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:01.15333Z","iopub.execute_input":"2021-12-28T17:55:01.154004Z","iopub.status.idle":"2021-12-28T17:55:01.162577Z","shell.execute_reply.started":"2021-12-28T17:55:01.153949Z","shell.execute_reply":"2021-12-28T17:55:01.161802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_ds = tf.data.Dataset.list_files(target_dir, shuffle=True)\n# batch images so they have the right shape for neural net\nstyled_ds = list_ds.map(process_fun, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:01.164126Z","iopub.execute_input":"2021-12-28T17:55:01.164754Z","iopub.status.idle":"2021-12-28T17:55:03.171344Z","shell.execute_reply.started":"2021-12-28T17:55:01.164716Z","shell.execute_reply":"2021-12-28T17:55:03.170634Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"default_cmap = 'gray' if TO_GRAYSCALE else None\ndef plot(image, ax=None, cmap=default_cmap):\n    if ax is None:\n        ax = plt.gca()\n    ax.imshow(image * 0.5 + 0.5, cmap=cmap)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:03.174551Z","iopub.execute_input":"2021-12-28T17:55:03.174837Z","iopub.status.idle":"2021-12-28T17:55:03.180988Z","shell.execute_reply.started":"2021-12-28T17:55:03.1748Z","shell.execute_reply":"2021-12-28T17:55:03.180156Z"},"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-12-28T17:55:03.182294Z","iopub.execute_input":"2021-12-28T17:55:03.182561Z","iopub.status.idle":"2021-12-28T17:55:03.919213Z","shell.execute_reply.started":"2021-12-28T17:55:03.182524Z","shell.execute_reply":"2021-12-28T17:55:03.916262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Source Domain Dataset\n\nOne of:\n- https://www.kaggle.com/c/gan-getting-started\n- https://www.kaggle.com/andrewmvd/animal-faces","metadata":{}},{"cell_type":"code","source":"if SOURCE_DATASET == 'animal-faces':\n    source_dir = '/kaggle/input/animal-faces/afhq/train/*/*.jpg'\nelif SOURCE_DATASET == 'landscapes':\n    source_dir = '/kaggle/input/gan-getting-started/photo_jpg/*.jpg'\nelse:\n    raise ValueError('Invalid dataset name!')\nprint(source_dir)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:03.920156Z","iopub.execute_input":"2021-12-28T17:55:03.92038Z","iopub.status.idle":"2021-12-28T17:55:03.92631Z","shell.execute_reply.started":"2021-12-28T17:55:03.920352Z","shell.execute_reply":"2021-12-28T17:55:03.925602Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_ds = tf.data.Dataset.list_files(source_dir, shuffle=True)\nphoto_ds = list_ds.map(process_fun, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:03.927878Z","iopub.execute_input":"2021-12-28T17:55:03.928814Z","iopub.status.idle":"2021-12-28T17:55:07.804101Z","shell.execute_reply.started":"2021-12-28T17:55:03.928573Z","shell.execute_reply":"2021-12-28T17:55:07.803373Z"},"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-12-28T17:55:07.80606Z","iopub.execute_input":"2021-12-28T17:55:07.806543Z","iopub.status.idle":"2021-12-28T17:55:08.472343Z","shell.execute_reply.started":"2021-12-28T17:55:07.806508Z","shell.execute_reply":"2021-12-28T17:55:08.47077Z"},"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 = 1 if TO_GRAYSCALE else 3\nINPUT_CHANNELS = 1 if TO_GRAYSCALE else 3","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:08.47372Z","iopub.execute_input":"2021-12-28T17:55:08.474171Z","iopub.status.idle":"2021-12-28T17:55:08.478597Z","shell.execute_reply.started":"2021-12-28T17:55:08.474137Z","shell.execute_reply":"2021-12-28T17:55:08.477634Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def downsample(filters, size, apply_instancenorm=True, name=None):\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(name=name)\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-12-28T17:55:08.480293Z","iopub.execute_input":"2021-12-28T17:55:08.480713Z","iopub.status.idle":"2021-12-28T17:55:08.490039Z","shell.execute_reply.started":"2021-12-28T17:55:08.480676Z","shell.execute_reply":"2021-12-28T17:55:08.489132Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False, apply_instancenorm=True, name=None):\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(name=name)\n    result.add(layers.Conv2DTranspose(filters, size, strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n    \n    if apply_instancenorm:\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-12-28T17:55:08.491947Z","iopub.execute_input":"2021-12-28T17:55:08.492242Z","iopub.status.idle":"2021-12-28T17:55:08.50246Z","shell.execute_reply.started":"2021-12-28T17:55:08.492206Z","shell.execute_reply":"2021-12-28T17:55:08.501884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Resized Convolution\ndef upsample_rc(input_layer, filters, size=3, strides=2, apply_instancenorm=True, name=None, 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    block = tf.image.resize(images=input_layer, method='bilinear', \n                            size=(input_layer.shape[1]*2, input_layer.shape[2]*2))\n    \n    block = layers.Conv2D(filters, size, \n                     strides=strides, \n                     padding='same', \n                     use_bias=False, \n                     kernel_initializer=initializer, \n                     name=f'decoder_{name}')(block)\n\n    if apply_instancenorm:\n        block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(block)\n        \n    if apply_dropout:\n        block = layers.Dropout(0.5)(block)\n\n    block = layers.ReLU()(block)\n    \n    return block","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T17:55:08.503576Z","iopub.execute_input":"2021-12-28T17:55:08.504091Z","iopub.status.idle":"2021-12-28T17:55:08.515392Z","shell.execute_reply.started":"2021-12-28T17:55:08.504055Z","shell.execute_reply":"2021-12-28T17:55:08.514641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def transformer_block(input_layer, size=3, strides=1, name='block_x'):\n    filters = input_layer.shape[-1]\n    initializer = tf.random_normal_initializer(0., 0.02)\n    \n    block = layers.Conv2D(filters, size, strides=strides, padding='same', use_bias=False, \n                     kernel_initializer=initializer, name=f'transformer_{name}_1')(input_layer)\n#     block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n    block = layers.ReLU()(block)\n    \n    block = layers.Conv2D(filters, size, strides=strides, padding='same', use_bias=False, \n                     kernel_initializer=initializer, name=f'transformer_{name}_2')(block)\n#     block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n    \n    block = layers.Add()([block, input_layer])\n\n    return block\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T17:55:08.51666Z","iopub.execute_input":"2021-12-28T17:55:08.517049Z","iopub.status.idle":"2021-12-28T17:55:08.527976Z","shell.execute_reply.started":"2021-12-28T17:55:08.517015Z","shell.execute_reply":"2021-12-28T17:55:08.527086Z"},"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,INPUT_CHANNELS])\n\n    # bs = batch size, assuming (h, w) = (256, 256)\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False, name='down_1'), # (bs, 128, 128, 64)\n        downsample(128, 4, name='down_2'), # (bs, 64, 64, 128)\n        downsample(256, 4, name='down_3'), # (bs, 32, 32, 256)\n        downsample(512, 4, name='down_4'), # (bs, 16, 16, 512)\n        downsample(512, 4, name='down_5'), # (bs, 8, 8, 512)\n        downsample(512, 4, name='down_6'), # (bs, 4, 4, 512)\n        downsample(512, 4, name='down_7'), # (bs, 2, 2, 512)\n        downsample(512, 4, name='down_8'), # (bs, 1, 1, 512)\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True, name='up_1'), # (bs, 2, 2, 1024)\n        upsample(512, 4, apply_dropout=True, name='up_2'), # (bs, 4, 4, 1024)\n        upsample(512, 4, apply_dropout=True, name='up_3'), # (bs, 8, 8, 1024)\n        upsample(512, 4, name='up_4'), # (bs, 16, 16, 1024)\n        upsample(256, 4, name='up_5'), # (bs, 32, 32, 512)\n        upsample(128, 4, name='up_6'), # (bs, 64, 64, 256)\n        upsample(64, 4, name='up_7'), # (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',\n                                  name='up_8_final'\n                                 ) # (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, name='Generator')","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:08.529322Z","iopub.execute_input":"2021-12-28T17:55:08.529558Z","iopub.status.idle":"2021-12-28T17:55:08.709099Z","shell.execute_reply.started":"2021-12-28T17:55:08.529535Z","shell.execute_reply":"2021-12-28T17:55:08.708271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def GeneratorV2(transformer_blocks: int = 6):\n    inputs = layers.Input(shape=[img_height,img_width,3], name='input_image')\n\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False, name='down_1'), # (bs, 128, 128, 64)\n        downsample(128, 4, name='down_2'), # (bs, 64, 64, 128)\n        downsample(256, 4, name='down_3'), # (bs, 32, 32, 256)\n#         downsample(512, 4, name='down_4'), # (bs, 16, 16, 512)\n#         downsample(512, 4, name='down_5'), # (bs, 8, 8, 512)\n#         downsample(512, 4, name='down_6'), # (bs, 4, 4, 512)\n#         downsample(512, 4, name='down_7'), # (bs, 2, 2, 512)\n#         downsample(512, 4, name='down_8'), # (bs, 1, 1, 512)\n    ]\n    \n    up_stack = [\n#         upsample(512, 4, apply_dropout=True, name='up_8'), # (bs, 2, 2, 1024)\n#         upsample(512, 4, apply_dropout=True, name='up_7'), # (bs, 4, 4, 1024)\n#         upsample(512, 4, apply_dropout=True, name='up_6'), # (bs, 8, 8, 1024)\n#         upsample(512, 4, name='up_5'), # (bs, 16, 16, 1024)\n#         upsample(256, 4, name='up_4'), # (bs, 32, 32, 512)\n        upsample(128, 4, name='up_3'), # (bs, 64, 64, 256)\n        upsample(64, 4, name='up_2'), # (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',\n                                  name='up_1_final'\n                                 ) # (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    # Transformer\n    for n in range(transformer_blocks):\n        x = transformer_block(x, 3, 1, name=f'T_{n+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, name='Generator')","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2021-12-28T17:55:08.711022Z","iopub.execute_input":"2021-12-28T17:55:08.711333Z","iopub.status.idle":"2021-12-28T17:55:08.72231Z","shell.execute_reply.started":"2021-12-28T17:55:08.711296Z","shell.execute_reply":"2021-12-28T17:55:08.721411Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def GeneratorV3():\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, name='down_1'), # (bs, 128, 128, 64)\n        downsample(128, 4, name='down_2'), # (bs, 64, 64, 128)\n        downsample(256, 4, name='down_3'), # (bs, 32, 32, 256)\n        downsample(512, 4, name='down_4'), # (bs, 16, 16, 512)\n        downsample(512, 4, name='down_5'), # (bs, 8, 8, 512)\n        downsample(512, 4, name='down_6'), # (bs, 4, 4, 512)\n        downsample(512, 4, name='down_7'), # (bs, 2, 2, 512)\n        downsample(512, 4, name='down_8'), # (bs, 1, 1, 512)\n    ]\n\n    up_stack = [\n        {'filters': 512, 'size': 4, 'apply_dropout': True, 'name': 'up_-1'}, # (bs, 2, 2, 1024)\n        {'filters': 512, 'size': 4, 'apply_dropout': True, 'name': 'up_-2'}, # (bs, 4, 4, 1024)\n        {'filters': 512, 'size': 4, 'apply_dropout': True, 'name': 'up_-3'}, # (bs, 8, 8, 1024)\n        {'filters': 512, 'size': 4, 'apply_dropout': False, 'name': 'up_-4'}, # (bs, 16, 16, 1024)\n        {'filters': 256, 'size': 4, 'apply_dropout': False, 'name': 'up_-5'}, # (bs, 32, 32, 512)\n        {'filters': 128, 'size': 4, 'apply_dropout': False, 'name': 'up_-6'}, # (bs, 64, 64, 256)\n        {'filters': 64, 'size': 4, 'apply_dropout': False, 'name': 'up_-7'}, # (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',\n                                  name='up_-8_final'\n                                 ) # (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_params, skip in zip(up_stack, skips):\n        x = upsample_rc(x, **up_params)\n        x = layers.Concatenate()([x, skip])\n\n    x = last(x)\n\n    return keras.Model(inputs=inputs, outputs=x, name='Generator')","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T17:55:08.723969Z","iopub.execute_input":"2021-12-28T17:55:08.724232Z","iopub.status.idle":"2021-12-28T17:55:08.739428Z","shell.execute_reply.started":"2021-12-28T17:55:08.724199Z","shell.execute_reply":"2021-12-28T17:55:08.738766Z"},"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, INPUT_CHANNELS], 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-12-28T17:55:08.742682Z","iopub.execute_input":"2021-12-28T17:55:08.743022Z","iopub.status.idle":"2021-12-28T17:55:08.753194Z","shell.execute_reply.started":"2021-12-28T17:55:08.742989Z","shell.execute_reply":"2021-12-28T17:55:08.752242Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"styled_generator = Generator() # transforms photos to drawings\nphoto_generator = Generator() # transforms drawings to photos\n\nstyled_discriminator = Discriminator() # differentiates real and generated drawings\nphoto_discriminator = Discriminator() # differentiates real and generated photos","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:08.754477Z","iopub.execute_input":"2021-12-28T17:55:08.755013Z","iopub.status.idle":"2021-12-28T17:55:11.172422Z","shell.execute_reply.started":"2021-12-28T17:55:08.754977Z","shell.execute_reply":"2021-12-28T17:55:11.171686Z"},"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-12-28T17:55:11.173782Z","iopub.execute_input":"2021-12-28T17:55:11.174046Z","iopub.status.idle":"2021-12-28T17:55:11.224932Z","shell.execute_reply.started":"2021-12-28T17:55:11.174013Z","shell.execute_reply":"2021-12-28T17:55:11.224303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example.shape","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:11.226324Z","iopub.execute_input":"2021-12-28T17:55:11.226587Z","iopub.status.idle":"2021-12-28T17:55:11.233634Z","shell.execute_reply.started":"2021-12-28T17:55:11.226555Z","shell.execute_reply":"2021-12-28T17:55:11.232814Z"},"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-12-28T17:55:11.234898Z","iopub.execute_input":"2021-12-28T17:55:11.235769Z","iopub.status.idle":"2021-12-28T17:55:17.633579Z","shell.execute_reply.started":"2021-12-28T17:55:11.235732Z","shell.execute_reply":"2021-12-28T17:55:17.632852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CycleGAN","metadata":{}},{"cell_type":"code","source":"# discriminator loss is for parts of the image, compute mean to present one number\n_readable_loss = True\n\nclass 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        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        if _readable_loss:\n            return {\n                \"styled_gen_loss\": tf.math.reduce_mean(total_styled_gen_loss),\n                \"photo_gen_loss\": tf.math.reduce_mean(total_photo_gen_loss),\n                \"styled_disc_loss\": tf.math.reduce_mean(styled_disc_loss),\n                \"photo_disc_loss\": tf.math.reduce_mean(photo_disc_loss)\n            }\n        else:\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-12-28T17:55:17.63845Z","iopub.execute_input":"2021-12-28T17:55:17.638661Z","iopub.status.idle":"2021-12-28T17:55:17.659466Z","shell.execute_reply.started":"2021-12-28T17:55:17.638636Z","shell.execute_reply":"2021-12-28T17:55:17.658764Z"},"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":"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-12-28T17:55:17.661445Z","iopub.execute_input":"2021-12-28T17:55:17.661647Z","iopub.status.idle":"2021-12-28T17:55:17.670873Z","shell.execute_reply.started":"2021-12-28T17:55:17.661624Z","shell.execute_reply":"2021-12-28T17:55:17.670191Z"},"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":"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-12-28T17:55:17.673884Z","iopub.execute_input":"2021-12-28T17:55:17.674642Z","iopub.status.idle":"2021-12-28T17:55:17.679829Z","shell.execute_reply.started":"2021-12-28T17:55:17.674587Z","shell.execute_reply":"2021-12-28T17:55:17.679054Z"},"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":"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-12-28T17:55:17.68101Z","iopub.execute_input":"2021-12-28T17:55:17.681721Z","iopub.status.idle":"2021-12-28T17:55:17.688684Z","shell.execute_reply.started":"2021-12-28T17:55:17.681678Z","shell.execute_reply":"2021-12-28T17:55:17.688031Z"},"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":"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-12-28T17:55:17.690047Z","iopub.execute_input":"2021-12-28T17:55:17.690816Z","iopub.status.idle":"2021-12-28T17:55:17.701218Z","shell.execute_reply.started":"2021-12-28T17:55:17.690778Z","shell.execute_reply":"2021-12-28T17:55:17.700475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"styled_generator_optimizer = tf.keras.optimizers.Adam(LR, beta_1=BETA_1)\nphoto_generator_optimizer = tf.keras.optimizers.Adam(LR, beta_1=BETA_1)\n\nstyled_discriminator_optimizer = tf.keras.optimizers.Adam(LR, beta_1=BETA_1)\nphoto_discriminator_optimizer = tf.keras.optimizers.Adam(LR, beta_1=BETA_1)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:17.702302Z","iopub.execute_input":"2021-12-28T17:55:17.702749Z","iopub.status.idle":"2021-12-28T17:55:17.710496Z","shell.execute_reply.started":"2021-12-28T17:55:17.702713Z","shell.execute_reply":"2021-12-28T17:55:17.709697Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cycle_gan_model = CycleGan(\n    styled_generator, photo_generator, styled_discriminator, photo_discriminator\n)\n\ncycle_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-12-28T17:55:17.711794Z","iopub.execute_input":"2021-12-28T17:55:17.712516Z","iopub.status.idle":"2021-12-28T17:55:17.736883Z","shell.execute_reply.started":"2021-12-28T17:55:17.712478Z","shell.execute_reply":"2021-12-28T17:55:17.736279Z"},"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":{"_kg_hide-input":true}},{"cell_type":"code","source":"len(photo_ds), len(styled_ds)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T17:55:17.737851Z","iopub.execute_input":"2021-12-28T17:55:17.738445Z","iopub.status.idle":"2021-12-28T17:55:17.745029Z","shell.execute_reply.started":"2021-12-28T17:55:17.738403Z","shell.execute_reply":"2021-12-28T17:55:17.744224Z"},"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"},"_kg_hide-input":true}},{"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":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T17:55:17.746352Z","iopub.execute_input":"2021-12-28T17:55:17.74715Z","iopub.status.idle":"2021-12-28T17:55:17.77159Z","shell.execute_reply.started":"2021-12-28T17:55:17.747028Z","shell.execute_reply":"2021-12-28T17:55:17.770938Z"},"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, cmap='gray')\n        axs[1].imshow(prediction, cmap='gray')\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-12-28T17:55:17.773814Z","iopub.execute_input":"2021-12-28T17:55:17.774237Z","iopub.status.idle":"2021-12-28T17:55:17.786238Z","shell.execute_reply.started":"2021-12-28T17:55:17.774177Z","shell.execute_reply":"2021-12-28T17:55:17.785506Z"},"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-12-28T17:55:17.787381Z","iopub.execute_input":"2021-12-28T17:55:17.787862Z","iopub.status.idle":"2021-12-28T17:55:17.798253Z","shell.execute_reply.started":"2021-12-28T17:55:17.787828Z","shell.execute_reply":"2021-12-28T17:55:17.79751Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# For some reason the standard History does not work correctly with the custom training loop.\n# This custom callback saves the history at the end of each batch and I reconstruct it later.\nclass LogCallback(keras.callbacks.Callback):\n    def __init__(self):\n        super().__init__()\n    \n    def on_train_begin(self, logs=None):\n        self.logs = []\n    \n    def on_train_batch_end(self, batch, logs=None):\n        self.logs.append(logs)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T17:55:17.799333Z","iopub.execute_input":"2021-12-28T17:55:17.79983Z","iopub.status.idle":"2021-12-28T17:55:17.80716Z","shell.execute_reply.started":"2021-12-28T17:55:17.799792Z","shell.execute_reply":"2021-12-28T17:55:17.806431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"L = LogCallback()\nif TEST_PIPELINE_ONLY:\n    epochs = 7\n    steps_per_epoch = 9\n\n_training_steps = epochs * steps_per_epoch\n    \n# limit the size of the buffer if there are too many steps to fit into memory\n_buffer_size = min(_training_steps, 3000)\n\ncallbacks=[\n        # limit training to fit within Kaggle time limit\n        tfa.callbacks.TimeStopping(seconds=TIME_LIMIT, verbose=1),\n        L,\n              ]\n\nif SHOW_SAMPLES_DURING_TRAINING:\n    callbacks.append(ShowSamplesCallback('/kaggle/working/imgs'))\n    \nif PERIODIC_SAVES > 0:\n    callbacks.append(SaveGeneratorCallback(save_freq, '/kaggle/working'))\n\nif EARLY_STOPPING_PATIENCE >= 0:\n    if not _readable_loss:\n        raise ValueError('This callback requires _readable_loss.')\n    callbacks.append(keras.callbacks.EarlyStopping(monitor='photo_gen_loss', patience=EARLY_STOPPING_PATIENCE, verbose=1, mode='min', restore_best_weights=True))\n    \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=steps_per_epoch, # batches per epoch\n    epochs=epochs,\n    # --- callbacks ---\n    callbacks=callbacks,\n)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:50:50.935234Z","iopub.execute_input":"2021-12-28T18:50:50.935519Z","iopub.status.idle":"2021-12-28T19:04:58.005471Z","shell.execute_reply.started":"2021-12-28T18:50:50.935488Z","shell.execute_reply":"2021-12-28T19:04:58.004381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<div id=\"training_results\"><b>Scroll up</b> for training log or <b>down</b> for graph and visualisations.</div>","metadata":{}},{"cell_type":"code","source":"import seaborn as sns","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:15:58.081604Z","iopub.execute_input":"2021-12-28T18:15:58.081946Z","iopub.status.idle":"2021-12-28T18:15:58.643066Z","shell.execute_reply.started":"2021-12-28T18:15:58.081906Z","shell.execute_reply":"2021-12-28T18:15:58.642301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if _readable_loss:\n    b_loss = [x['styled_gen_loss'] for x in L.logs]\nelse:\n    b_loss = [x['styled_gen_loss'].mean() for x in L.logs]\ne_loss = [np.mean(b_loss[i:i+steps_per_epoch]) for i in range(0, len(b_loss), steps_per_epoch)]\ndf = pd.DataFrame(data={'styled_gen_loss': e_loss})","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:15:58.644336Z","iopub.execute_input":"2021-12-28T18:15:58.644574Z","iopub.status.idle":"2021-12-28T18:15:58.655593Z","shell.execute_reply.started":"2021-12-28T18:15:58.644542Z","shell.execute_reply":"2021-12-28T18:15:58.654754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.to_csv('/kaggle/working/loss.csv',index=False, header=True)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:15:58.657052Z","iopub.execute_input":"2021-12-28T18:15:58.65757Z","iopub.status.idle":"2021-12-28T18:15:58.669317Z","shell.execute_reply.started":"2021-12-28T18:15:58.657535Z","shell.execute_reply":"2021-12-28T18:15:58.668533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!head /kaggle/working/loss.csv","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:15:58.673591Z","iopub.execute_input":"2021-12-28T18:15:58.673779Z","iopub.status.idle":"2021-12-28T18:15:59.387139Z","shell.execute_reply.started":"2021-12-28T18:15:58.673757Z","shell.execute_reply":"2021-12-28T18:15:59.386304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, ax = plt.subplots(1, figsize=(12,6))\nsns.lineplot(data=df, ax=ax)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:15:59.389152Z","iopub.execute_input":"2021-12-28T18:15:59.389718Z","iopub.status.idle":"2021-12-28T18:15:59.654632Z","shell.execute_reply.started":"2021-12-28T18:15:59.389676Z","shell.execute_reply":"2021-12-28T18:15:59.653815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# I'm not sure why but the below code does not display correct history - it is probably saved incorrectly inside the custom train_step\n\n\n# h_ = history.history\n\n\n# i = range(len(h_['styled_gen_loss']))\n\n# fig, ax = plt.subplots(1, figsize=(12,6))\n\n# sns.lineplot(x=i, y=[x.mean() for x in h_['styled_gen_loss']], ax=ax)\n# # sns.lineplot(x=i, y=h_['photo_gen_loss'], ax=ax)\n# # sns.lineplot(x=i, y=h_['styled_disc_loss'], ax=ax)\n# # sns.lineplot(x=i, y=h_['photo_disc_loss'], ax=ax)\n# # fig.legend(['styled_gen_loss', 'photo_gen_loss', 'styled_disc_loss', 'photo_disc_loss'])\n# fig.legend(['styled_gen_loss'])\n# plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T18:15:59.655941Z","iopub.execute_input":"2021-12-28T18:15:59.656391Z","iopub.status.idle":"2021-12-28T18:15:59.661759Z","shell.execute_reply.started":"2021-12-28T18:15:59.656352Z","shell.execute_reply":"2021-12-28T18:15:59.660869Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# pd.DataFrame.from_dict(h_).to_csv('/kaggle/working/history.csv',index=False)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T18:15:59.6629Z","iopub.execute_input":"2021-12-28T18:15:59.663628Z","iopub.status.idle":"2021-12-28T18:15:59.669804Z","shell.execute_reply.started":"2021-12-28T18:15:59.663566Z","shell.execute_reply":"2021-12-28T18:15:59.669114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !head /kaggle/working/history.csv","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T18:15:59.671992Z","iopub.execute_input":"2021-12-28T18:15:59.672197Z","iopub.status.idle":"2021-12-28T18:15:59.678916Z","shell.execute_reply.started":"2021-12-28T18:15:59.672175Z","shell.execute_reply":"2021-12-28T18:15:59.678189Z"},"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#     plot(img, ax[0, i])\n#     plot(prediction, ax[1, i])\n    ax[0, i].imshow(img, cmap=default_cmap)\n    ax[1, i].imshow(prediction, cmap=default_cmap)\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-12-28T18:24:53.26434Z","iopub.execute_input":"2021-12-28T18:24:53.264608Z","iopub.status.idle":"2021-12-28T18:24:54.706491Z","shell.execute_reply.started":"2021-12-28T18:24:53.264579Z","shell.execute_reply":"2021-12-28T18:24:54.705852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Saving/Loading","metadata":{}},{"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-12-28T18:16:00.603768Z","iopub.execute_input":"2021-12-28T18:16:00.606987Z","iopub.status.idle":"2021-12-28T18:16:11.745368Z","shell.execute_reply.started":"2021-12-28T18:16:00.606939Z","shell.execute_reply":"2021-12-28T18:16:11.744537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo_generator.save('/kaggle/working/saved_model/photo_generator')","metadata":{"execution":{"iopub.status.busy":"2021-12-28T18:16:11.774331Z","iopub.execute_input":"2021-12-28T18:16:11.774581Z","iopub.status.idle":"2021-12-28T18:16:21.065796Z","shell.execute_reply.started":"2021-12-28T18:16:11.774553Z","shell.execute_reply":"2021-12-28T18:16:21.065038Z"},"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-12-28T18:16:21.096883Z","iopub.execute_input":"2021-12-28T18:16:21.097162Z","iopub.status.idle":"2021-12-28T18:16:25.789992Z","shell.execute_reply.started":"2021-12-28T18:16:21.097133Z","shell.execute_reply":"2021-12-28T18:16:25.789268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_, ax = plt.subplots(2, 5, figsize=(16, 12))\nfor i, img in enumerate(photo_ds.take(5)):\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#     plot(img, ax[0, i])\n#     plot(prediction, ax[1, i])\n    ax[0, i].imshow(img, cmap=default_cmap)\n    ax[1, i].imshow(prediction, cmap=default_cmap)\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-12-28T18:24:41.097668Z","iopub.execute_input":"2021-12-28T18:24:41.09797Z","iopub.status.idle":"2021-12-28T18:24:42.309712Z","shell.execute_reply.started":"2021-12-28T18:24:41.097925Z","shell.execute_reply":"2021-12-28T18:24:42.308924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Work-in-progress experiments, visualisations, ...","metadata":{"_kg_hide-input":true}},{"cell_type":"markdown","source":"## Cyclic Loss Visualisation","metadata":{}},{"cell_type":"code","source":"for _ in range(3):\n    _, ax = plt.subplots(1, 3, figsize=(12, 8), squeeze=True)\n    for i, img in enumerate(photo_ds.take(1)):\n        prediction = styled_generator(img, training=False)\n        cycled = photo_generator(prediction, training=False)\n\n        prediction_img = prediction[0].numpy()\n        prediction_img = (prediction_img * 127.5 + 127.5).astype(np.uint8)\n\n        img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n        cycled_img = cycled[0].numpy()\n        cycled_img = (cycled_img * 127.5 + 127.5).astype(np.uint8)\n\n        \n        ax[0].imshow(img, cmap=default_cmap)\n        ax[1].imshow(prediction_img, cmap=default_cmap)\n        ax[2].imshow(cycled_img, cmap=default_cmap)\n#         plot(img, ax[0])\n#         plot(prediction_img, ax[1])\n#         plot(cycled_img, ax[2])\n\n        ax[0].set_title(\"Input Photo\")\n        ax[1].set_title(\"Generated\")\n        ax[2].set_title(\"Cycled\")\n        ax[0].axis(\"off\")\n        ax[1].axis(\"off\")\n        ax[2].axis(\"off\")\n    plt.tight_layout()\n    plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-12-28T18:25:32.025991Z","iopub.execute_input":"2021-12-28T18:25:32.026247Z","iopub.status.idle":"2021-12-28T18:25:34.23842Z","shell.execute_reply.started":"2021-12-28T18:25:32.026218Z","shell.execute_reply":"2021-12-28T18:25:34.237777Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Results, Comparison","metadata":{}},{"cell_type":"code","source":"import seaborn as sns\nimport pandas as pd\nfrom pathlib import Path\nimport matplotlib.pyplot as plt\n\nresults_dir = Path('/kaggle/input/mvi-results')","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:31:31.660709Z","iopub.execute_input":"2021-12-28T19:31:31.660998Z","iopub.status.idle":"2021-12-28T19:31:31.667595Z","shell.execute_reply.started":"2021-12-28T19:31:31.660948Z","shell.execute_reply":"2021-12-28T19:31:31.666892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_meta = pd.read_csv(results_dir / 'version_map.txt')","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:35:38.752682Z","iopub.execute_input":"2021-12-28T19:35:38.752951Z","iopub.status.idle":"2021-12-28T19:35:38.761767Z","shell.execute_reply.started":"2021-12-28T19:35:38.752922Z","shell.execute_reply":"2021-12-28T19:35:38.760922Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_meta","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:48:26.893825Z","iopub.execute_input":"2021-12-28T19:48:26.894428Z","iopub.status.idle":"2021-12-28T19:48:26.906327Z","shell.execute_reply.started":"2021-12-28T19:48:26.894391Z","shell.execute_reply":"2021-12-28T19:48:26.904076Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(df_meta.to_latex())","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:48:36.892754Z","iopub.execute_input":"2021-12-28T19:48:36.893434Z","iopub.status.idle":"2021-12-28T19:48:36.905893Z","shell.execute_reply.started":"2021-12-28T19:48:36.893393Z","shell.execute_reply":"2021-12-28T19:48:36.904816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"losses = []\nfor _, row in df_meta.iterrows():\n    v = row.version\n    p = (results_dir / f'{v:02d}').with_suffix('.csv')\n    df_v = pd.read_csv(p)\n    for epoch, loss in enumerate(df_v['styled_gen_loss'], start=1):\n        row_cp = row.copy()\n        row_cp['epoch'] = epoch\n        row_cp['loss'] = loss\n        losses.append(row_cp)\n\ndf = pd.DataFrame(losses)\ndf['version'] = df['version'].astype('category')\ndf = df.reset_index()","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:39:01.823538Z","iopub.execute_input":"2021-12-28T19:39:01.823810Z","iopub.status.idle":"2021-12-28T19:39:02.231408Z","shell.execute_reply.started":"2021-12-28T19:39:01.823772Z","shell.execute_reply":"2021-12-28T19:39:02.230690Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, ax = plt.subplots(1, figsize=(12,8))\nsns.lineplot(x=df['epoch'], y=df['loss'], hue=df['version'], ax=ax)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:39:05.464072Z","iopub.execute_input":"2021-12-28T19:39:05.464345Z","iopub.status.idle":"2021-12-28T19:39:05.823454Z","shell.execute_reply.started":"2021-12-28T19:39:05.464315Z","shell.execute_reply":"2021-12-28T19:39:05.822839Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df['dataset'] = df['source_ds'] + ' -> ' + df['target_ds']","metadata":{"execution":{"iopub.status.busy":"2021-12-28T19:53:43.747300Z","iopub.execute_input":"2021-12-28T19:53:43.747570Z","iopub.status.idle":"2021-12-28T19:53:43.756052Z","shell.execute_reply.started":"2021-12-28T19:53:43.747539Z","shell.execute_reply":"2021-12-28T19:53:43.755341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, ax = plt.subplots(1, figsize=(12,8))\nax.set_ylabel('mean loss')\nsns.lineplot(x=df['epoch'], y=df['loss'], hue=df['dataset'], ax=ax, estimator='mean', ci=0)","metadata":{"execution":{"iopub.status.busy":"2021-12-28T20:01:27.379995Z","iopub.execute_input":"2021-12-28T20:01:27.380265Z","iopub.status.idle":"2021-12-28T20:01:30.101207Z","shell.execute_reply.started":"2021-12-28T20:01:27.380236Z","shell.execute_reply":"2021-12-28T20:01:30.100512Z"},"trusted":true},"execution_count":null,"outputs":[]}]}