{"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":"code","source":"import os, random, json, PIL, shutil, re\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nimport tensorflow as tf\nimport tensorflow.keras.layers as L\nimport tensorflow_addons as tfa\nfrom tensorflow.keras import Model, losses, optimizers\nfrom tensorflow import keras\nfrom tensorflow.keras import layers","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-11-24T19:32:51.401424Z","iopub.execute_input":"2021-11-24T19:32:51.401753Z","iopub.status.idle":"2021-11-24T19:32:51.408904Z","shell.execute_reply.started":"2021-11-24T19:32:51.401717Z","shell.execute_reply":"2021-11-24T19:32:51.407692Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Setting up TPU","metadata":{}},{"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print(f'Running on TPU {tpu.master()}')\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\n\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')\nAUTO = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:51.415517Z","iopub.execute_input":"2021-11-24T19:32:51.415973Z","iopub.status.idle":"2021-11-24T19:32:57.553530Z","shell.execute_reply.started":"2021-11-24T19:32:51.415931Z","shell.execute_reply":"2021-11-24T19:32:57.552252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model Parameters","metadata":{}},{"cell_type":"code","source":"HEIGHT = 256\nWIDTH = 256\nCHANNELS = 3\nEPOCHS = 50\nBATCH_SIZE = 1\n\n# Generates tensors with a normal distribution\ninitializer = tf.random_normal_initializer(0., 0.02)\ngamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:57.555461Z","iopub.execute_input":"2021-11-24T19:32:57.555709Z","iopub.status.idle":"2021-11-24T19:32:57.561533Z","shell.execute_reply.started":"2021-11-24T19:32:57.555667Z","shell.execute_reply":"2021-11-24T19:32:57.560271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = KaggleDatasets().get_gcs_path()\n\nMONET_FILENAMES = tf.io.gfile.glob(str(PATH + '/monet_tfrec/*.tfrec'))\nPHOTO_FILENAMES = tf.io.gfile.glob(str(PATH + '/photo_tfrec/*.tfrec'))\n\ndef count_files(filenames):\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)\n\nn_monet_samples = count_files(MONET_FILENAMES)\nn_photo_samples = count_files(PHOTO_FILENAMES)\n\nprint(f'Monet TFRecord files: {len(MONET_FILENAMES)}')\nprint(f'Monet image files: {n_monet_samples}')\nprint(f'Photo TFRecord files: {len(PHOTO_FILENAMES)}')\nprint(f'Photo image files: {n_photo_samples}')","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:57.562790Z","iopub.execute_input":"2021-11-24T19:32:57.563176Z","iopub.status.idle":"2021-11-24T19:32:58.142440Z","shell.execute_reply.started":"2021-11-24T19:32:57.563147Z","shell.execute_reply":"2021-11-24T19:32:58.141288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=CHANNELS)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [HEIGHT, WIDTH, CHANNELS])\n    return image\n\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        'image_name': tf.io.FixedLenFeature([], tf.string),\n        'image':      tf.io.FixedLenFeature([], tf.string),\n        'target':     tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image\n\ndef load_dataset(filenames):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)\n    return dataset\n\ndef get_gan_dataset(monet_files, photo_files, augment=None, repeat=True, shuffle=True, batch_size=1):\n\n    monet_ds = load_dataset(monet_files)\n    photo_ds = load_dataset(photo_files)\n\n    if repeat:\n        monet_ds = monet_ds.repeat()\n        photo_ds = photo_ds.repeat()\n    if shuffle:\n        monet_ds = monet_ds.shuffle(2048)\n        photo_ds = photo_ds.shuffle(2048)\n        \n    monet_ds = monet_ds.batch(batch_size, drop_remainder=True)\n    photo_ds = photo_ds.batch(batch_size, drop_remainder=True)\n    monet_ds = monet_ds.cache()\n    photo_ds = photo_ds.cache()\n    monet_ds = monet_ds.prefetch(AUTO)\n    photo_ds = photo_ds.prefetch(AUTO)\n    \n    gan_ds = tf.data.Dataset.zip((monet_ds, photo_ds))\n    \n    return gan_ds\n\n","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.144090Z","iopub.execute_input":"2021-11-24T19:32:58.144416Z","iopub.status.idle":"2021-11-24T19:32:58.159303Z","shell.execute_reply.started":"2021-11-24T19:32:58.144377Z","shell.execute_reply":"2021-11-24T19:32:58.158153Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def display_samples(dataset, row, col):\n    ds_iter = iter(dataset)\n    plt.figure(figsize=(15, int(15*row/col)))\n    for j in range(row*col):\n        example_sample = next(ds_iter)\n        plt.subplot(row,col,j+1)\n        plt.axis('off')\n        plt.imshow(example_sample[0] * 0.5 + 0.5)\n    plt.show()\n        \ndef display_generated_samples(dataset, model, no_of_images):\n    # The iter() function returns an iterator object for the given object.\n    ds_iter = iter(dataset)\n    \n    for image in range(no_of_images):\n        plt.subplots(figsize=(15, 10))\n        \n        random_image = next(ds_iter)\n        generated_image = model.predict(random_image)\n        \n        \n        plt.subplot(121)\n        plt.title(\"Input image\")\n        plt.imshow(random_image[0] * 0.5 + 0.5)\n        plt.axis('off')\n        \n        plt.subplot(122)\n        plt.title(\"Generated image\")\n        plt.imshow(generated_image[0] * 0.5 + 0.5)\n        plt.axis('off')\n        plt.show()\n        \ndef predict_and_save(input_ds, generator_model, output_path):\n    i = 1\n    for img in input_ds:\n        prediction = generator_model(img, training=False)[0].numpy() # make predition\n        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)   # re-scale\n        im = PIL.Image.fromarray(prediction)\n        im.save(f'{output_path}{str(i)}.jpg')\n        i += 1","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.162653Z","iopub.execute_input":"2021-11-24T19:32:58.163172Z","iopub.status.idle":"2021-11-24T19:32:58.177619Z","shell.execute_reply.started":"2021-11-24T19:32:58.163129Z","shell.execute_reply":"2021-11-24T19:32:58.176950Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Model functions\ndef downsample(filters, size, apply_instancenorm = True, strides=2):\n    \n    result = keras.Sequential()\n    result.add(L.Conv2D(filters, size, strides=strides, padding='same', kernel_initializer=initializer, use_bias=False))\n    \n    if apply_instancenorm == True:\n        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n    result.add(L.LeakyReLU())\n\n    return result\n","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.179312Z","iopub.execute_input":"2021-11-24T19:32:58.179927Z","iopub.status.idle":"2021-11-24T19:32:58.196180Z","shell.execute_reply.started":"2021-11-24T19:32:58.179884Z","shell.execute_reply":"2021-11-24T19:32:58.195271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(filters, size, dropout=False, strides=2):\n    \n    result = keras.Sequential()\n    result.add(L.Conv2DTranspose(filters, size, strides=strides, padding='same', kernel_initializer=initializer, use_bias=False))\n\n    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout == True:\n        result.add(L.Dropout(0.5))\n    result.add(L.ReLU())\n\n    return result","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.198390Z","iopub.execute_input":"2021-11-24T19:32:58.199329Z","iopub.status.idle":"2021-11-24T19:32:58.209525Z","shell.execute_reply.started":"2021-11-24T19:32:58.199283Z","shell.execute_reply":"2021-11-24T19:32:58.208875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\n\ndef create_generator():\n    \n    # First layer - using keras.layers.inputLayer -- specifying explicitly.\n    inputs = L.Input(shape=[HEIGHT, WIDTH, CHANNELS])\n    \n    \n    model_1 = downsample(64,4,apply_instancenorm = False)\n    down_stack = []\n    down_stack.append(model_1)\n\n    down_stack.append(downsample(128, 4))\n    down_stack.append(downsample(256, 4))\n    down_stack.append(downsample(512, 4))\n    down_stack.append(downsample(512, 4))\n    down_stack.append(downsample(512, 4))\n    down_stack.append(downsample(512, 4))\n    down_stack.append(downsample(512, 4))\n\n    model_2 = upsample(512, 4, apply_dropout = True)\n    up_stack = []\n    up_stack.append(model_2)\n    \n    up_stack.append(upsample(512,4 ,dropout = True))\n    up_stack.append(upsample(512,4,dropout = True))\n    up_stack.append(upsample(512,4,dropout = True))\n    up_stack.append(upsample(512,4))\n    up_stack.append(upsample(256,4))\n    up_stack.append(upsample(128,4))\n    up_stack.append(upsample(64,4))\n\n\n    #Last conv block with `tanh` activation. ( as per documentation)\n    final_layer = L.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=2,padding='same',kernel_initializer=initializer, activation='tanh') # (bs, 256, 256, 3)\n\n    x = inputs\n\n    # Downsampling through the model\n    skips = []\n    for down in down_stack:\n        # First layer 'down' is added to the second layer 'x'\n        x = down(x)\n        skips.append(x)\n \n    \n    # everything except the last downsample layer\n    skips = reversed(skips[:-1])\n\n    # Upsampling and establishing the skip connections\n    # https://www.kaggle.com/jesperdramsch/understanding-and-improving-cyclegans-tutorial -- Highly recommended to understand skip connections\n    \n    # This is for concatenating the first layer of upsampling and last layer of downsampling and so on.\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = L.Concatenate()([x, skip])\n\n    x = final_layer(x)\n\n    return Model(inputs=inputs, outputs=x)","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.210895Z","iopub.execute_input":"2021-11-24T19:32:58.211348Z","iopub.status.idle":"2021-11-24T19:32:58.228948Z","shell.execute_reply.started":"2021-11-24T19:32:58.211298Z","shell.execute_reply":"2021-11-24T19:32:58.228181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_discriminator():\n\n    input_layer = L.Input(shape=[HEIGHT, WIDTH, CHANNELS], name='input_image')\n\n    model = keras.Sequential()\n    model.add(input_layer)\n    \n    down_1 = downsample(64,4, False)\n    down_2 = downsample(128,4)\n    down_3 = downsample(256, 4)\n    \n    model.add(down_1)\n    model.add(down_2)\n    model.add(down_3)\n    \n    zero_pad_1 = L.ZeroPadding2D()\n    conv = L.Conv2D(512, 4, strides=1,kernel_initializer=initializer, use_bias=False)\n    norm1 = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)\n    leaky_relu = L.LeakyReLU()\n    zero_pad_2 = L.ZeroPadding2D()\n    last = L.Conv2D(1, 4, strides=1,kernel_initializer=initializer)\n    \n    model.add(zero_pad_1)\n    model.add(conv)\n    model.add(norm1)\n    model.add(leaky_relu)\n    model.add(zero_pad_2)\n    model.add(last)\n\n#     return Model(inputs=input_layer, outputs=last)\n    return model","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.230046Z","iopub.execute_input":"2021-11-24T19:32:58.230418Z","iopub.status.idle":"2021-11-24T19:32:58.248216Z","shell.execute_reply.started":"2021-11-24T19:32:58.230387Z","shell.execute_reply":"2021-11-24T19:32:58.247302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    monet_G = create_generator() # transforms photos to Monet-esque paintings\n    photo_G = create_generator() # transforms Monet paintings to be more like photos\n\n    monet_D = create_discriminator() # differentiates real Monet paintings and generated Monet paintings\n    photo_D = create_discriminator() # differentiates real photos and generated photos\n\n\nclass CycleGan(Model):\n    def __init__(self, monet_G, photo_G, monet_D, photo_D, lambda_cycle = 10 ,lambda_identity=0.5):\n        super(CycleGan, self).__init__()\n        self.m_generator = monet_G\n        self.p_generator = photo_G\n        self.m_disc = monet_D\n        self.p_disc = photo_D\n        self.lambda_cycle = lambda_cycle\n        self.lambda_identity = lambda_identity\n        \n    def compile(self, m_gen_optimizer, p_gen_optimizer, m_disc_optimizer, p_disc_optimizer, generator_loss_function, discriminator_loss_function,\n                cycle_loss_function,\n                identity_loss_function\n               ):\n        super(CycleGan, self).compile()\n        self.m_gen_optimizer = m_gen_optimizer\n        self.p_gen_optimizer = p_gen_optimizer\n        self.m_disc_optimizer = m_disc_optimizer\n        self.p_disc_optimizer = p_disc_optimizer\n        self.generator_loss_function = generator_loss_function\n        self.discriminator_loss_function = discriminator_loss_function\n        self.cycle_loss_function = cycle_loss_function\n        self.identity_loss_function = identity_loss_function\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            # photo to monet back to photo\n            fake_monet = self.m_generator(real_photo, training=True)\n            cycled_photo = self.p_generator(fake_monet, training=True)\n\n            # monet to photo back to monet\n            fake_photo = self.p_generator(real_monet, training=True)\n            cycled_monet = self.m_generator(fake_photo, training=True)\n\n            # generating itself\n            same_monet = self.m_generator(real_monet, training=True)\n            same_photo = self.p_generator(real_photo, training=True)\n\n            # discriminator used to check, inputing real images\n            disc_real_monet = self.m_disc(real_monet, 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_monet = self.m_disc(fake_monet, training=True)\n            disc_fake_photo = self.p_disc(fake_photo, training=True)\n\n            # evaluates generator loss\n            monet_gen_loss = self.generator_loss_function(disc_fake_monet)\n            photo_gen_loss = self.generator_loss_function(disc_fake_photo)\n            \n            \n            # evaluates total cycle consistency loss\n#             total_cycle_loss = self.cycle_loss_fn(real_monet, cycled_monet) * self.lambda_cycle + self.cycle_loss_fn(real_photo, cycled_photo)* self.lambda_cycle\n            monet_cycle_loss = self.cycle_loss_function(real_monet, cycled_monet) * self.lambda_cycle\n            photo_cycle_loss = self.cycle_loss_function(real_photo, cycled_photo) * self.lambda_cycle\n            \n            # evaluates total generator loss\n            total_monet_generator_loss = monet_gen_loss + monet_cycle_loss + (self.identity_loss_function(real_monet, same_monet) * self.lambda_cycle * self.lambda_identity)\n            total_photo_generator_loss = photo_gen_loss + photo_cycle_loss + (self.identity_loss_function(real_photo, same_photo) * self.lambda_cycle * self.lambda_identity)\n\n            # evaluates total discriminator loss\n            monet_discriminator_loss = self.discriminator_loss_function(disc_real_monet, disc_fake_monet)\n            photo_discriminator_loss = self.discriminator_loss_function(disc_real_photo, disc_fake_photo)\n            \n            \n        # Compute gradients for gen and disc\n        monet_generator_gradients = tape.gradient(total_monet_generator_loss, self.m_generator.trainable_variables)\n        photo_generator_gradients = tape.gradient(total_photo_generator_loss, self.p_generator.trainable_variables)\n        monet_discriminator_gradients = tape.gradient(monet_discriminator_loss, self.m_disc.trainable_variables)\n        photo_discriminator_gradients = tape.gradient(photo_discriminator_loss, self.p_disc.trainable_variables)\n\n        # Apply the gradients to the optimizer\n        self.m_gen_optimizer.apply_gradients(zip(monet_generator_gradients, self.m_generator.trainable_variables))\n        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients, self.p_generator.trainable_variables))\n        self.m_disc_optimizer.apply_gradients(zip(monet_discriminator_gradients, self.m_disc.trainable_variables))\n        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients, self.p_disc.trainable_variables))\n        \n        return {\n            'monet_generator_loss': total_monet_generator_loss,\n            'photo_generator_loss': total_photo_generator_loss,\n            'monet_discriminator_loss': monet_discriminator_loss,\n            'photo_discriminator_loss': photo_discriminator_loss\n        }","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:32:58.249427Z","iopub.execute_input":"2021-11-24T19:32:58.249666Z","iopub.status.idle":"2021-11-24T19:33:03.279986Z","shell.execute_reply.started":"2021-11-24T19:32:58.249639Z","shell.execute_reply":"2021-11-24T19:33:03.278983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"adversarial_loss_function = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)   \nwith strategy.scope():\n    \n    # Discriminator loss\n    def calculate_discriminator_loss(real_image, generated_image):\n        real_loss = adversarial_loss_function( tf.ones_like(real_image), real_image )\n        generated_loss = adversarial_loss_function( tf.zeros_like(generated_image), generated_image )\n        total_discriminator_loss = real_loss + generated_loss\n        return total_discriminator_loss * 0.5\n    \n    # Generator loss\n    def calculate_generator_loss(generated_image):\n        total_generator_loss = losses.BinaryCrossentropy( from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(generated_image), generated_image )\n        return total_generator_loss\n    \n    # Cycle consistency loss (measures if original photo and the twice transformed photo to be similar to one another)\n    with strategy.scope():\n        def calculate_cycle_loss(real_image, cycled_image):\n            cycle_loss = tf.reduce_mean(tf.abs(real_image - cycled_image))\n            return cycle_loss\n\n    # Identity loss (compares the image with its generator (i.e. photo with photo generator))\n    with strategy.scope():\n        def calculate_identity_loss(real_image, same_image):\n            total_identity_loss = tf.reduce_mean(tf.abs(real_image - same_image))\n            return total_identity_loss","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:33:03.281352Z","iopub.execute_input":"2021-11-24T19:33:03.281584Z","iopub.status.idle":"2021-11-24T19:33:03.293237Z","shell.execute_reply.started":"2021-11-24T19:33:03.281558Z","shell.execute_reply":"2021-11-24T19:33:03.292289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Adam_optimizer = optimizers.Adam(2e-4, beta_1=0.5)\nwith strategy.scope():\n\n    # Create GAN\n    gan_model = CycleGan(monet_G, photo_G, monet_D, photo_D)\n\n    gan_model.compile(m_gen_optimizer=Adam_optimizer,\n                      p_gen_optimizer=Adam_optimizer,\n                      m_disc_optimizer=Adam_optimizer,\n                      p_disc_optimizer=Adam_optimizer,\n                      generator_loss_function=calculate_generator_loss,\n                      discriminator_loss_function=calculate_discriminator_loss,\n                      cycle_loss_function=calculate_cycle_loss,\n                      identity_loss_function=calculate_identity_loss\n                     )\n    \nhistory = gan_model.fit(get_gan_dataset(MONET_FILENAMES, PHOTO_FILENAMES, batch_size=BATCH_SIZE), \n                        steps_per_epoch=(n_monet_samples//BATCH_SIZE),\n                        epochs=1,\n                        verbose=1).history","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:33:03.294629Z","iopub.execute_input":"2021-11-24T19:33:03.294953Z","iopub.status.idle":"2021-11-24T19:35:27.215001Z","shell.execute_reply.started":"2021-11-24T19:33:03.294924Z","shell.execute_reply":"2021-11-24T19:35:27.213799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"display_generated_samples(load_dataset(PHOTO_FILENAMES).batch(1), monet_G, 10)\n","metadata":{"execution":{"iopub.status.busy":"2021-11-24T19:35:27.216982Z","iopub.execute_input":"2021-11-24T19:35:27.217226Z","iopub.status.idle":"2021-11-24T19:35:43.698055Z","shell.execute_reply.started":"2021-11-24T19:35:27.217196Z","shell.execute_reply":"2021-11-24T19:35:43.697128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}