{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Disclaimer\n\nThis notebook was created as a frame to transpose (from PyTorch), step by step, the Contrastive Unpaired Translation, aka CUT [https://github.com/taesungp/contrastive-unpaired-translation]\n\nI was eager to have a shot at this topic.\n\nThis is actually my first attempt on CycleGAN and I'll use a mixture some of the available tutorials to grasp some intuition around this topic.\n\nFeel free to comment so that I can improve :)\n\nKudos\n\nMJO"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -q git+https://github.com/tensorflow/docs","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"#IMPORTS\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\nfrom tensorflow import keras\nimport tensorflow_addons as tfa\nfrom tensorflow.keras import layers\nfrom tensorflow.keras.callbacks import History\n\n\nimport os, time\nfrom kaggle_datasets import KaggleDatasets\nfrom IPython.display import clear_output\n\nimport tensorflow_docs.vis.embed as embed\nimport PIL\nfrom IPython import display\nimport imageio\n\nimport shutil","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's create a Configuration class to gather every variables."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Taken from https://www.kaggle.com/forwet/unpaired-data-cyclicgan-awesome-monets\n# Adapted to CUT\n# Configuration\nclass Configuration:\n    \"\"\"Class containing most of the parameters or hyperparameters used\n    throughout the notebook.\"\"\"\n    \n    epochs = 30\n    MONET_TFREC = \"/monet_tfrec/*.tfrec\"\n    MONET_JPG = \"/monet_jpg/*.jpg\"\n    PHOTO_TFREC = \"/photo_tfrec/*.tfrec\"\n    PHOTO_JPG = \"/photo_jpg/*.jpg\"\n    BATCH_SIZE = 8\n    IMAGE_SIZE = [256, 256, 3]\n    BUFFER = 10000\n    steps_per_epoch = 0\n    \n    #In CUT\n    lambda_cycle = 10\n    lambda_id = 0.5\n    \n    #In CUT\n    # Weights initializer for the layers.\n    kernel_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    # Gamma initializer for instance normalization.\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    \n    #Original\n    #loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n    \n    #In CUT\n    loss = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)\n\ncfg = Configuration()\nmonet_jpg = tf.io.gfile.glob(\"../input/gan-getting-started/monet_jpg/*.jpg\")\ncfg.steps_per_epoch = len(monet_jpg)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## TPU Initialization"},{"metadata":{"trusted":true},"cell_type":"code","source":"# TPU Setup\n# Taken from Kaggle's Tutorial\n# https://www.kaggle.com/amyjang/monet-cyclegan-tutorial\n\ntry:\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)\nexcept:\n    strategy = tf.distribute.get_strategy()\nprint('Number of replicas:', strategy.num_replicas_in_sync)\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Datasets creation through a class.\n\nDataAugmentation is realized here.\n\n[Maybe separate DataAugmentation for future testing]"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Taken from https://www.kaggle.com/forwet/unpaired-data-cyclicgan-awesome-monets\n# To do - CUT processing from publication\nclass MonetDataset:\n        def __init__(self,config):\n            \"\"\"Creates a data of TFRecord files.\"\"\"\n            self.cfg = config\n            # Specifing 'gan-getting-started' is mendatory as we\n            # load also the output datas from input folder\n            gcs_path = KaggleDatasets().get_gcs_path('gan-getting-started')\n            self.monet_files = tf.io.gfile.glob(gcs_path+self.cfg.MONET_TFREC)\n            self.photo_files = tf.io.gfile.glob(gcs_path + self.cfg.PHOTO_TFREC)\n            \n        def decode_image(self, image):\n            \"\"\"Function to preprocess the image prior to training.\"\"\"\n            img = tf.image.decode_jpeg(image, channels=3)\n            img = tf.cast(img, tf.float32)\n            img = img/127.5 - 1\n            img = tf.reshape(img, [*self.cfg.IMAGE_SIZE])\n            return img\n        \n        def read_tfrecord(self, instance):\n            \"\"\"Function to extract data from TFRecordDataset Instance.\"\"\"\n            tfrecordformat = {\n                    \"image_name\": tf.io.FixedLenFeature([], tf.string),\n                    \"image\": tf.io.FixedLenFeature([], tf.string),\n                    \"target\": tf.io.FixedLenFeature([], tf.string)\n                   }\n            example = tf.io.parse_single_example(instance, tfrecordformat)\n            return self.decode_image(example[\"image\"])\n        \n        def prepare_dataset(self, monet=True):\n            \"\"\"Main function to prepare the input pipeline.\n            Args: \n            monet- bool value\n            Determines if we wanna generate monet dataset or the photo dataset\"\"\"\n            dataset = tf.data.TFRecordDataset(self.monet_files if monet else self.photo_files, num_parallel_reads=AUTOTUNE)\n            dataset = dataset.map(self.read_tfrecord, num_parallel_calls=AUTOTUNE)\n            dataset = dataset.map(self.random_jitter, num_parallel_calls=AUTOTUNE)\n            dataset = dataset.repeat()\n            dataset = dataset.shuffle(self.cfg.BUFFER)\n            dataset = dataset.batch(self.cfg.BATCH_SIZE)\n            dataset = dataset.prefetch(AUTOTUNE)\n            return dataset\n\n        def random_crop(self, image):\n            \"\"\"Function to perform random cropping.\"\"\"\n            image = tf.image.random_crop(image, [*self.cfg.IMAGE_SIZE])\n            return image\n        \n        def random_jitter(self, image):\n            \"\"\"Function to perform random jittering.\"\"\"\n            image = tf.image.resize(image, [286, 286])\n            image = self.random_crop(image)\n            \n            if tf.random.uniform([], 0, 1) > 0.5:\n                image = tf.image.random_flip_left_right(image)\n            return image\n        \n        \n        \n        def visualize_data(self, data):\n            \"\"\"Utility function to visualize the samples in the dataset instance being provided.\"\"\"\n            fig, ax = plt.subplots(2, self.cfg.BATCH_SIZE//2, figsize=(8, 4)) # Figsize->W x H\n            ax = ax.flatten()\n            for i, im in zip(range(self.cfg.BATCH_SIZE), data):\n                im = im*0.5 + 0.5\n                ax[i].imshow(im)\n                ax[i].axis(\"off\")\n            plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Creation and quick looks"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Creating instance of dataset\ndataset = MonetDataset(Configuration())\n\n# Creating seperate monet and photo dataset\nmonet_dataset = dataset.prepare_dataset(monet=True)\nphoto_dataset = dataset.prepare_dataset(monet=False)\n\n# Checking some Monet examples\ndataset.visualize_data(next(iter(monet_dataset)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Checking some Photo examples\ndataset.visualize_data(next(iter(photo_dataset)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"In CUT, paddings are realized by using Reflection & Replica padding.\n\nLet's create these paddings as layers to put in our different networks.\n\n[Issue between TPU and tf.pad - Top priority]"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"#Issue between TPU and tf.pad - looking into it - MirrorPadGrad\n#Error :\n#  ...on XLA_TPU_JIT: MirrorPadGrad (No registered 'MirrorPadGrad' OpKernel for XLA_TPU_JIT...\n\n# Definition of two Padding Layers\n# [Taken and adapted from Keras tutorial on CycleGAN]\n# https://keras.io/examples/generative/cyclegan/\nclass ReflectionPadding2D(layers.Layer):\n    \"\"\"Implements Reflection Padding as a layer.\n\n    Args:\n        padding(tuple): Amount of padding for the\n        spatial dimensions.\n\n    Returns:\n        A padded tensor with the same type as the input tensor.\n    \"\"\"\n\n    def __init__(self, padding=(1, 1), **kwargs):\n        self.padding = tuple(padding)\n        super(ReflectionPadding2D, self).__init__(**kwargs)\n        \n    def compute_output_shape(self, input_shape):\n        return(input_shape[0], input_shape[1] + 2 * self.padding[0], input_shape[2] + 2 * self.padding[1], input_shape[3])\n\n    def call(self, input_tensor, mask=None):\n        padding_width, padding_height = self.padding\n        padding_tensor = [\n            [0, 0],\n            [padding_height, padding_height],\n            [padding_width, padding_width],\n            [0, 0],\n        ]\n        return tf.pad(input_tensor, padding_tensor, mode=\"REFLECT\")\n\nclass ReplicaPadding2D(layers.Layer):\n    \"\"\"Implements Reflection Padding as a layer.\n\n    Args:\n        padding(tuple): Amount of padding for the\n        spatial dimensions.\n\n    Returns:\n        A padded tensor with the same type as the input tensor.\n    \"\"\"\n\n    def __init__(self, padding=(1, 1), **kwargs):\n        self.padding = tuple(padding)\n        super(ReplicaPadding2D, self).__init__(**kwargs)\n        \n    def compute_output_shape(self, input_shape):\n        return(input_shape[0], input_shape[1] + 2 * self.padding[0], input_shape[2] + 2 * self.padding[1], input_shape[3])\n\n    def call(self, input_tensor, mask=None):\n        padding_width, padding_height = self.padding\n        padding_tensor = [\n            [0, 0],\n            [padding_height, padding_height],\n            [padding_width, padding_width],\n            [0, 0],\n        ]\n        return tf.pad(input_tensor, padding_tensor, mode=\"SYMMETRIC\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The generator networks will be composed of 9 Residual Blocks.\n\nLet's automatize the block generation."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Definition of Residual Block\n# For Resnet generator as in CUT\n# Issue with the Padding impacts here\ndef residual_block(x,\n                   activation,\n                   kernel_initializer=cfg.kernel_init,\n                   kernel_size=(3, 3),\n                   strides=(1, 1),\n                   padding=\"valid\",\n                   gamma_initializer=cfg.gamma_init,\n                   use_bias=False):\n    dim = x.shape[-1]\n    input_tensor = x\n\n    # x = ReflectionPadding2D()(input_tensor) #Issue with TPU and tf.pad... looking into it\n    x = layers.ZeroPadding2D()(input_tensor)\n    x = layers.Conv2D(\n        dim,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = activation(x)\n\n    x = layers.ZeroPadding2D()(x) #Should be ReflectionPadding2D\n    x = layers.Conv2D(\n        dim,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = layers.add([input_tensor, x])\n    return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Creation of the Downsampler and the Upsampler."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Downsampling step\n# ReflectionPadding should also be included here\ndef downsample(\n    x,\n    filters,\n    activation,\n    kernel_initializer=cfg.kernel_init,\n    kernel_size=(3, 3),\n    strides=(2, 2),\n    padding=\"same\",\n    gamma_initializer=cfg.gamma_init,\n    use_bias=False,\n):\n    x = layers.Conv2D(\n        filters,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    if activation:\n        x = activation(x)\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Downsampling step\n# ReplicationPadding should be included here\ndef upsample(\n    x,\n    filters,\n    activation,\n    kernel_size=(3, 3),\n    strides=(2, 2),\n    padding=\"same\",\n    kernel_initializer=cfg.kernel_init,\n    gamma_initializer=cfg.gamma_init,\n    use_bias=False,\n):\n    x = layers.Conv2DTranspose(\n        filters,\n        kernel_size,\n        strides=strides,\n        padding=padding,\n        kernel_initializer=kernel_initializer,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    if activation:\n        x = activation(x)\n    return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Creation of the total generator and discriminator."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Creating the Res Net generator\n# Combining Resnet Blocks, Downsampler, Upsampler\ndef get_resnet_generator(\n    filters=64,\n    num_downsampling_blocks=2,\n    num_residual_blocks=9,\n    num_upsample_blocks=2,\n    gamma_initializer=cfg.gamma_init,\n    name=None,\n):\n    img_input = layers.Input(shape=cfg.IMAGE_SIZE, name=name + \"_img_input\")\n    #x = ReflectionPadding2D(padding=(3, 3))(img_input)\n    x = layers.ZeroPadding2D(padding=(3, 3))(img_input)\n    x = layers.Conv2D(filters, (7, 7), kernel_initializer=cfg.kernel_init, use_bias=False)(\n        x\n    )\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = layers.Activation(\"relu\")(x)\n\n    # Downsampling\n    for _ in range(num_downsampling_blocks):\n        filters *= 2\n        x = downsample(x, filters=filters, activation=layers.Activation(\"relu\"))\n\n    # Residual blocks\n    for _ in range(num_residual_blocks):\n        x = residual_block(x, activation=layers.Activation(\"relu\"))\n\n    # Upsampling\n    for _ in range(num_upsample_blocks):\n        filters //= 2\n        x = upsample(x, filters, activation=layers.Activation(\"relu\"))\n\n    # Final block\n    x = layers.ZeroPadding2D(padding=(3, 3))(x) #Should be ReflectionPadding2D\n    x = layers.Conv2D(3, (7, 7), padding=\"valid\")(x)\n    x = layers.Activation(\"tanh\")(x)\n\n    model = keras.models.Model(img_input, x, name=name)\n    return model\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Creating the Discriminator\ndef get_discriminator(\n    filters=64, kernel_initializer=cfg.kernel_init, num_downsampling=3, name=None\n):\n    img_input = layers.Input(shape=cfg.IMAGE_SIZE, name=name + \"_img_input\")\n    x = layers.Conv2D(\n        filters,\n        (4, 4),\n        strides=(2, 2),\n        padding=\"same\",\n        kernel_initializer=kernel_initializer,\n    )(img_input)\n    x = layers.LeakyReLU(0.2)(x)\n\n    num_filters = filters\n    for num_downsample_block in range(3):\n        num_filters *= 2\n        if num_downsample_block < 2:\n            x = downsample(\n                x,\n                filters=num_filters,\n                activation=layers.LeakyReLU(0.2),\n                kernel_size=(4, 4),\n                strides=(2, 2),\n            )\n        else:\n            x = downsample(\n                x,\n                filters=num_filters,\n                activation=layers.LeakyReLU(0.2),\n                kernel_size=(4, 4),\n                strides=(1, 1),\n            )\n\n    x = layers.Conv2D(\n        1, (4, 4), strides=(1, 1), padding=\"same\", kernel_initializer=kernel_initializer\n    )(x)\n\n    model = keras.models.Model(inputs=img_input, outputs=x, name=name)\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Calling our fonction to declare our networks."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Call this to use prediction directly\ngen_monet = tf.keras.models.load_model('../input/monetcycleganoutputs/gen_monet.h5')\ngen_photo = tf.keras.models.load_model('../input/monetcycleganoutputs/gen_photo.h5')\ndisc_monet = tf.keras.models.load_model('../input/monetcycleganoutputs/disc_monet.h5')\ndisc_photo = tf.keras.models.load_model('../input/monetcycleganoutputs/disc_photo.h5')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Call this before training\n# Generate the 4 networks\n#with strategy.scope():\n    # Get the generators\n#    gen_monet = get_resnet_generator(name=\"generator_Monet\")\n#    gen_photo = get_resnet_generator(name=\"generator_Photo\")\n\n    # Get the discriminators\n#    disc_monet = get_discriminator(name=\"discriminator_Monet\")\n#    disc_photo = get_discriminator(name=\"discriminator_Photo\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Quicklook on the graphs.\n\nGenerator:"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Having a look to the generator graph\ntf.keras.utils.plot_model(gen_monet, show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Discriminator:"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Having a look to the discriminator graph\ntf.keras.utils.plot_model(disc_monet, show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"CycleGAN Model with our defined models."},{"metadata":{"trusted":true},"cell_type":"code","source":"# CycleGAN Model with our defined models\nclass CycleGan(keras.Model):\n    def __init__(\n        self,\n        generator_Monet,\n        generator_Photo,\n        discriminator_Monet,\n        discriminator_Photo,\n        lambda_cycle=cfg.lambda_cycle,\n        lambda_identity=cfg.lambda_id,\n    ):\n        super(CycleGan, self).__init__()\n        self.gen_Monet = generator_Monet\n        self.gen_Photo = generator_Photo\n        self.disc_Monet = discriminator_Monet\n        self.disc_Photo = discriminator_Photo\n        self.lambda_cycle = lambda_cycle\n        self.lambda_identity = lambda_identity\n\n    def compile(\n        self,\n        gen_Monet_optimizer,\n        gen_Photo_optimizer,\n        disc_Monet_optimizer,\n        disc_Photo_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n        cycl_loss_fn,\n        id_loss_fn\n    ):\n        super(CycleGan, self).compile()\n        self.gen_Monet_optimizer = gen_Monet_optimizer\n        self.gen_Photo_optimizer = gen_Photo_optimizer\n        self.disc_Monet_optimizer = disc_Monet_optimizer\n        self.disc_Photo_optimizer = disc_Photo_optimizer\n        self.generator_loss_fn = gen_loss_fn\n        self.discriminator_loss_fn = disc_loss_fn\n        self.cycl_loss_fn = cycl_loss_fn\n        self.id_loss_fn = id_loss_fn\n        \n\n    def train_step(self, batch_data):\n        real_monet, real_photo = batch_data\n\n        with tf.GradientTape(persistent=True) as tape:\n            \n            # Photo to fake Monet\n            fake_monet = self.gen_Monet(real_photo, training=True)\n            # Monet to fake Photo\n            fake_photo = self.gen_Photo(real_monet, training=True)\n\n            # Cycle (Monet to fake Photo to fake Monet): x -> y -> x\n            cycled_monet = self.gen_Monet(fake_photo, training=True)\n            # Cycle (Photo to fake Monet to fake Photo) y -> x -> y\n            cycled_photo = self.gen_Photo(fake_monet, training=True)\n\n            # Identity mapping\n            same_monet = self.gen_Monet(real_monet, training=True)\n            same_photo = self.gen_Photo(real_photo, training=True)\n\n            # Discriminator output\n            disc_real_monet = self.disc_Monet(real_monet, training=True)\n            disc_fake_monet = self.disc_Monet(fake_monet, training=True)\n\n            disc_real_photo = self.disc_Photo(real_photo, training=True)\n            disc_fake_photo = self.disc_Photo(fake_photo, training=True)\n\n            # Generator adverserial loss\n            gen_Monet_loss = self.generator_loss_fn(disc_fake_photo)\n            gen_Photo_loss = self.generator_loss_fn(disc_fake_monet)\n\n            # Generator cycle loss\n            cycle_loss_Monet = self.cycl_loss_fn(real_monet, cycled_monet) * self.lambda_cycle\n            cycle_loss_Photo = self.cycl_loss_fn(real_photo, cycled_photo) * self.lambda_cycle\n\n            # Generator identity loss\n            id_loss_Monet = (\n                self.id_loss_fn(real_monet, same_monet)\n                * self.lambda_cycle\n                * self.lambda_identity\n            )\n            id_loss_Photo = (\n                self.id_loss_fn(real_photo, same_photo)\n                * self.lambda_cycle\n                * self.lambda_identity\n            )\n\n            # Total generator loss\n            total_loss_Monet = gen_Monet_loss + cycle_loss_Monet + id_loss_Monet\n            total_loss_Photo = gen_Photo_loss + cycle_loss_Photo + id_loss_Photo\n\n            # Discriminator loss\n            disc_Monet_loss = self.discriminator_loss_fn(disc_real_monet, disc_fake_monet)\n            disc_Photo_loss = self.discriminator_loss_fn(disc_real_photo, disc_fake_photo)\n\n        # Get the gradients for the generators\n        grads_Monet = tape.gradient(total_loss_Monet, self.gen_Monet.trainable_variables)\n        grads_Photo = tape.gradient(total_loss_Photo, self.gen_Photo.trainable_variables)\n\n        # Get the gradients for the discriminators\n        disc_Monet_grads = tape.gradient(disc_Monet_loss, self.disc_Monet.trainable_variables)\n        disc_Photo_grads = tape.gradient(disc_Photo_loss, self.disc_Photo.trainable_variables)\n\n        # Update the weights of the generators\n        self.gen_Monet_optimizer.apply_gradients(\n            zip(grads_Monet, self.gen_Monet.trainable_variables)\n        )\n        self.gen_Photo_optimizer.apply_gradients(\n            zip(grads_Photo, self.gen_Photo.trainable_variables)\n        )\n\n        # Update the weights of the discriminators\n        self.disc_Monet_optimizer.apply_gradients(\n            zip(disc_Monet_grads, self.disc_Monet.trainable_variables)\n        )\n        self.disc_Photo_optimizer.apply_gradients(\n            zip(disc_Photo_grads, self.disc_Photo.trainable_variables)\n        )\n\n        return {\n            \"Monet_generator_loss\": total_loss_Monet,\n            \"Photo_generator_loss\": total_loss_Photo,\n            \"Monet_discriminator_loss\": disc_Monet_loss,\n            \"Photo_discriminator_loss\": disc_Photo_loss,\n        }\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CallBack function for GIF creation\n\nTaking one Photo to check how the CycleGAN is learning during each epoch.\n\nCreate a 'on_epoch_end' callback modification to generate predictions of our photo at each epoch. We will be able to generate a gif output through these."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Call these lines before training\n# You might want to check the different photos\n# in photo to select which one you want to monitor","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#global im_to_gif\n#im_to_gif = np.zeros((30,256,256,3))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#photo = next(iter(photo_dataset))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Taking one Photo from which we will check evolution\nnum_photo = 0\n#plt.imshow(photo[num_photo]*0.5 + 0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Generate a CallBack function to save\n# the prediction, for each epoch, of the Photo above \n#class GANMonitor(keras.callbacks.Callback):\n#    \"\"\"A callback to generate and save images after each epoch\"\"\"\n\n#    def on_epoch_end(self, epoch, logs=None):\n#        prediction = gen_monet(photo, training=False)[num_photo].numpy()\n#        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n#        im_to_gif[epoch] = prediction     ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Deploying everything on TPU."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Call this for before training\n# Defining losses, model and compiling\n#with strategy.scope():\n#    # Define the loss function for the generators\n#    def generator_loss_fn(fake):\n#        fake_loss = cfg.loss(tf.ones_like(fake), fake)\n#        return fake_loss\n\n\n    # Define the loss function for the discriminators\n#    def discriminator_loss_fn(real, fake):\n#        real_loss = cfg.loss(tf.ones_like(real), real)\n#        fake_loss = cfg.loss(tf.zeros_like(fake), fake)\n#        return (real_loss + fake_loss) * 0.5\n    \n#    def cyclic_loss_fn(real, cycled):\n#        return tf.reduce_mean(tf.abs(real - cycled))\n    \n#    def id_loss_fn(real, same):\n#        return tf.reduce_mean(tf.abs(real - same))\n\n    # Create cycle gan model\n#    cycle_gan_model = CycleGan(\n#        generator_Monet=gen_monet,\n#        generator_Photo=gen_photo,\n#        discriminator_Monet=disc_monet,\n#        discriminator_Photo=disc_photo\n#    )\n\n    # Compile the model\n#    cycle_gan_model.compile(\n#        gen_Monet_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n#        gen_Photo_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n#        disc_Monet_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n#        disc_Photo_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n#        gen_loss_fn=generator_loss_fn,\n#        disc_loss_fn=discriminator_loss_fn,\n#        cycl_loss_fn=cyclic_loss_fn,\n#        id_loss_fn=id_loss_fn\n#    )\n    # Callbacks\n#    plotter = GANMonitor()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Finally, training!\n\n[77 s / epoch => should be less than 40 min for 30 epochs]"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Training\n#with strategy.scope():\n#    history = cycle_gan_model.fit(tf.data.Dataset.zip((monet_dataset, photo_dataset)),\n#                        epochs=cfg.epochs,\n#                        steps_per_epoch=cfg.steps_per_epoch,\n#                        callbacks=[History(),\n#                                   plotter])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Saving outputs\n# Download them afterwards - refresh folder\n# You'll have to reupload them through 'Add data'\n# to use your own\n#gen_monet.save('gen_monet.h5')\n#gen_photo.save('gen_photo.h5')\n#disc_monet.save('disc_monet.h5')\n#disc_photo.save('disc_photo.h5')\n#import pickle\n#with open('history.pkl','wb') as f:\n#    pickle.dump(history.history, f)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's first have a look at how the different networks behaved during training.\n\n[Taking the mean might not be the good way, but we can still monitor something]"},{"metadata":{"trusted":true},"cell_type":"code","source":"def smooth_curve(points, factor=0.8):\n    smoothed_points = []\n    for point in points:\n        if smoothed_points:\n            previous = smoothed_points[-1]\n            smoothed_points.append(previous * factor + point * (1 - factor))\n        else:\n            smoothed_points.append(point)\n    return smoothed_points\n\n\ndef plot_smoothed_acc_and_loss(history, factor=0.8, load=False):\n    monet_g = []\n    photo_g = []\n    monet_d = []\n    photo_d = []\n    if load==True:\n        for i in range(np.array(history[\"Monet_generator_loss\"]).shape[0]):\n            monet_g.append(np.array(history[\"Monet_generator_loss\"][i]).squeeze().mean())\n            photo_g.append(np.array(history[\"Photo_generator_loss\"][i]).squeeze().mean())\n            monet_d.append(np.array(history[\"Monet_discriminator_loss\"][i]).squeeze().mean())\n            photo_d.append(np.array(history[\"Photo_discriminator_loss\"][i]).squeeze().mean())\n    else:\n        for i in range(np.array(history.history[\"Monet_generator_loss\"]).shape[0]):\n            monet_g.append(np.array(history.history[\"Monet_generator_loss\"][i]).squeeze().mean())\n            photo_g.append(np.array(history.history[\"Photo_generator_loss\"][i]).squeeze().mean())\n            monet_d.append(np.array(history.history[\"Monet_discriminator_loss\"][i]).squeeze().mean())\n            photo_d.append(np.array(history.history[\"Photo_discriminator_loss\"][i]).squeeze().mean())\n    \n    fig, axs = plt.subplots(1, 2, figsize=(15, 5))\n    axs[0].plot(np.arange(1,31,1),\n        smooth_curve(monet_g, factor=factor),\n        label=\"Monet\")\n    axs[0].plot(np.arange(1,31,1),\n        smooth_curve(photo_g, factor=factor),\n        label=\"Photo\")\n    axs[0].set_title(\"Smoothed generator loss\")\n    axs[0].legend()\n\n    axs[1].plot(np.arange(1,31,1),\n        smooth_curve(monet_d, factor=factor),\n        label=\"Monet\")\n    axs[1].plot(np.arange(1,31,1),\n        smooth_curve(photo_d, factor=factor),\n        label=\"Photo\")\n    axs[1].set_title(\"Smoothed discriminator loss\")\n    axs[1].legend()\n\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Call this if loading outputs\nimport pickle\nwith (open(\"../input/monetcycleganoutputs/history.pkl\", \"rb\")) as openfile:\n    history = pickle.load(openfile)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Switch 'load' to false if you have trained the model\nplot_smoothed_acc_and_loss(history, 0.8, load=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's make a simple visualization of how our model trains using a gif generator.\n\n[Taken and adapted from Tensorflow tutorial on DCGan]"},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_gif(num_photo=0, load=False):\n    if load == True:\n        anim_file = '../input/monetcycleganoutputs/CycleGAN.gif'\n    else:\n        # Creating a gif from each predictions\n        anim_file = 'CycleGAN.gif'\n        init_pic = np.array(photo[num_photo]*0.5 + 0.5)\n        with imageio.get_writer(anim_file, mode='I') as writer:\n            # Three first frames are the converted picture\n            writer.append_data(init_pic)\n            writer.append_data(init_pic)\n            writer.append_data(init_pic)\n            for i in range(im_to_gif.shape[0]):\n                writer.append_data(im_to_gif[i])\n                writer.append_data(im_to_gif[i])\n                writer.append_data(im_to_gif[i])\n            for i in range(int(im_to_gif.shape[0])):\n                writer.append_data(im_to_gif[-1])\n    return anim_file","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Switch 'load' to false if you have trained the model and put the correct photo number\nanim_file = create_gif(num_photo=num_photo,\n                       load=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def gen_input_img(num_photo=0, load=False):\n    fig, ax = plt.subplots(figsize=(5,5))\n    \n    if load == True:\n        img = np.array(PIL.Image.open('../input/monetcycleganoutputs/input_image.png'))\n        plt.imshow(img)\n        ax.axis(\"off\")\n        \n    else:\n        img = photo[3]*0.5 + 0.5\n        plt.imshow(img)\n        ax.axis(\"off\")\n        plt.title('Input photo')    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Switch 'load' to false if you have trained the model and put the correct photo number\ngen_input_img(num_photo=num_photo,\n              load=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Prediction evolution according to epoch\nembed.embed_file(anim_file)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's have another look on another Photo for the gist of it."},{"metadata":{"trusted":true},"cell_type":"code","source":"photo = next(iter(photo_dataset))\npredict_img = gen_monet.predict(tf.expand_dims(photo[0], axis=0))\nfig, ax = plt.subplots(1, 2, figsize=(8,8))\nax = ax.flatten()\nax[0].imshow(photo[0]*0.5 + 0.5)\nax[0].set_title('Input photo')\nax[0].axis(\"off\")\nout  = (predict_img[0]*127.5 + 127.5).astype(np.uint8)\nax[1].imshow(out)\nax[1].set_title('Output fake Monet')\nax[1].axis(\"off\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generating the output datas"},{"metadata":{"trusted":true},"cell_type":"code","source":"# !mkdir ../images\n\n# photo_jpg = tf.io.gfile.glob(\"../input/gan-getting-started/photo_jpg/*.jpg\")\n\n# for i, image in zip(range(1, len(photo_jpg)+1), photo_dataset):\n#     prediction = gen_monet(image, training=False)[0].numpy()\n#     prediction = (prediction*127.5 + 127.5).astype(np.uint8)\n#     im = PIL.Image.fromarray(prediction)\n#     im.save(f\"../images/{i}.jpg\")\n#     if(i%100==0):\n#         print(f\"Processed {i} images\")\n        \n# shutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Still working on the CUT transposition.\n\nIf anyone has ideas over the Padding issues, feel free to contact me.\n\nTODO :\n\n- Mirror padding\n- Adding the NCE layers\n- Recheck GAN and NCE loss criterion"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}