{"cells":[{"metadata":{},"cell_type":"markdown","source":"Code inspired from [TensorFlow CycleGan](https://www.tensorflow.org/tutorials/generative/cyclegan) and [Keras CycleGan](https://keras.io/examples/generative/cyclegan/)\n\nAlso, [https://www.kaggle.com/amyjang/monet-cyclegan-tutorial](https://www.kaggle.com/amyjang/monet-cyclegan-tutorial)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -q git+https://github.com/tensorflow/examples.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom tqdm.auto import tqdm\n\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\nfrom tensorflow.keras import layers as L\nfrom tensorflow_examples.models.pix2pix import pix2pix\n\nfrom kaggle_datasets import KaggleDatasets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\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\n    \nprint(tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path()\n\nMONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\nPHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\n\nBUFFER_SIZE = 1000\nBATCH_SIZE = 1\nIMG_WIDTH = 128\nIMG_HEIGHT = 128\n\nIMAGE_SIZE = [IMG_WIDTH, IMG_HEIGHT]\n\ndef decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.image.resize(image, IMAGE_SIZE)\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        \"image_name\": tf.io.FixedLenFeature([], tf.string),\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    return dataset\n\nmonet_ds = load_dataset(MONET_FILENAMES, labeled=True).batch(BATCH_SIZE)\nphoto_ds = load_dataset(PHOTO_FILENAMES, labeled=True).batch(BATCH_SIZE)\n\nexample_monet = next(iter(monet_ds))\nexample_photo = next(iter(photo_ds))\n\nplt.subplot(121)\nplt.title('Photo')\nplt.imshow(example_photo[0] * 0.5 + 0.5)\n\nplt.subplot(122)\nplt.title('Monet')\nplt.imshow(example_monet[0] * 0.5 + 0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nOUTPUT_CHANNELS = 3\n\ngenerator_g = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm')\ngenerator_f = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm')\n\ndiscriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False)\ndiscriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass CycleGan(tf.keras.Model):\n    def __init__(\n        self,\n        generator_G,\n        generator_F,\n        discriminator_X,\n        discriminator_Y,\n        lambda_cycle=10.0,\n        lambda_identity=0.5,\n    ):\n        super(CycleGan, self).__init__()\n        self.gen_G = generator_G\n        self.gen_F = generator_F\n        self.disc_X = discriminator_X\n        self.disc_Y = discriminator_Y\n        self.lambda_cycle = lambda_cycle\n        self.lambda_identity = lambda_identity\n\n    def compile(\n        self,\n        gen_G_optimizer,\n        gen_F_optimizer,\n        disc_X_optimizer,\n        disc_Y_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n        cycle_loss_fn=tf.keras.losses.MeanAbsoluteError(),\n        identity_loss_fn=tf.keras.losses.MeanAbsoluteError()\n        \n    ):\n        super(CycleGan, self).compile()\n        self.gen_G_optimizer = gen_G_optimizer\n        self.gen_F_optimizer = gen_F_optimizer\n        self.disc_X_optimizer = disc_X_optimizer\n        self.disc_Y_optimizer = disc_Y_optimizer\n        self.generator_loss_fn = gen_loss_fn\n        self.discriminator_loss_fn = disc_loss_fn\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n\n    def train_step(self, batch_data):\n        # x is Horse and y is zebra\n        real_x, real_y = batch_data\n\n        # For CycleGAN, we need to calculate different\n        # kinds of losses for the generators and discriminators.\n        # We will perform the following steps here:\n        #\n        # 1. Pass real images through the generators and get the generated images\n        # 2. Pass the generated images back to the generators to check if we\n        #    we can predict the original image from the generated image.\n        # 3. Do an identity mapping of the real images using the generators.\n        # 4. Pass the generated images in 1) to the corresponding discriminators.\n        # 5. Calculate the generators total loss (adverserial + cycle + identity)\n        # 6. Calculate the discriminators loss\n        # 7. Update the weights of the generators\n        # 8. Update the weights of the discriminators\n        # 9. Return the losses in a dictionary\n\n        with tf.GradientTape(persistent=True) as tape:\n            # Horse to fake zebra\n            fake_y = self.gen_G(real_x, training=True)\n            # Zebra to fake horse -> y2x\n            fake_x = self.gen_F(real_y, training=True)\n\n            # Cycle (Horse to fake zebra to fake horse): x -> y -> x\n            cycled_x = self.gen_F(fake_y, training=True)\n            # Cycle (Zebra to fake horse to fake zebra) y -> x -> y\n            cycled_y = self.gen_G(fake_x, training=True)\n\n            # Identity mapping\n            same_x = self.gen_F(real_x, training=True)\n            same_y = self.gen_G(real_y, training=True)\n\n            # Discriminator output\n            disc_real_x = self.disc_X(real_x, training=True)\n            disc_fake_x = self.disc_X(fake_x, training=True)\n\n            disc_real_y = self.disc_Y(real_y, training=True)\n            disc_fake_y = self.disc_Y(fake_y, training=True)\n\n            # Generator adverserial loss\n            gen_G_loss = self.generator_loss_fn(disc_fake_y)\n            gen_F_loss = self.generator_loss_fn(disc_fake_x)\n\n            # Generator cycle loss\n            cycle_loss_G = self.cycle_loss_fn(real_y, cycled_y, self.lambda_cycle) * self.lambda_cycle\n            cycle_loss_F = self.cycle_loss_fn(real_x, cycled_x, self.lambda_cycle) * self.lambda_cycle\n\n            # Generator identity loss\n            id_loss_G = (\n                self.identity_loss_fn(real_y, same_y, self.lambda_cycle)\n                * self.lambda_cycle\n                * self.lambda_identity\n            )\n            id_loss_F = (\n                self.identity_loss_fn(real_x, same_x, self.lambda_cycle)\n                * self.lambda_cycle\n                * self.lambda_identity\n            )\n\n            # Total generator loss\n            total_loss_G = gen_G_loss + cycle_loss_G + id_loss_G\n            total_loss_F = gen_F_loss + cycle_loss_F + id_loss_F\n\n            # Discriminator loss\n            disc_X_loss = self.discriminator_loss_fn(disc_real_x, disc_fake_x)\n            disc_Y_loss = self.discriminator_loss_fn(disc_real_y, disc_fake_y)\n\n        # Get the gradients for the generators\n        grads_G = tape.gradient(total_loss_G, self.gen_G.trainable_variables)\n        grads_F = tape.gradient(total_loss_F, self.gen_F.trainable_variables)\n\n        # Get the gradients for the discriminators\n        disc_X_grads = tape.gradient(disc_X_loss, self.disc_X.trainable_variables)\n        disc_Y_grads = tape.gradient(disc_Y_loss, self.disc_Y.trainable_variables)\n\n        # Update the weights of the generators\n        self.gen_G_optimizer.apply_gradients(\n            zip(grads_G, self.gen_G.trainable_variables)\n        )\n        self.gen_F_optimizer.apply_gradients(\n            zip(grads_F, self.gen_F.trainable_variables)\n        )\n\n        # Update the weights of the discriminators\n        self.disc_X_optimizer.apply_gradients(\n            zip(disc_X_grads, self.disc_X.trainable_variables)\n        )\n        self.disc_Y_optimizer.apply_gradients(\n            zip(disc_Y_grads, self.disc_Y.trainable_variables)\n        )\n\n        return {\n            \"G_loss\": total_loss_G,\n            \"F_loss\": total_loss_F,\n            \"D_X_loss\": disc_X_loss,\n            \"D_Y_loss\": disc_Y_loss,\n        }\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass GANMonitor(tf.keras.callbacks.Callback):\n    \"\"\"A callback to generate and save images after each epoch\"\"\"\n\n    def __init__(self, num_img=4):\n        self.num_img = num_img\n\n    def on_epoch_end(self, epoch, logs=None):\n        _, ax = plt.subplots(4, 2, figsize=(12, 12))\n        for i, img in enumerate(photo_ds.take(4)):\n            prediction = self.model.gen_G(img)[0].numpy()\n            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n            img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n            ax[i, 0].imshow(img)\n            ax[i, 1].imshow(prediction)\n            ax[i, 0].set_title(\"Input image\")\n            ax[i, 1].set_title(\"Translated image\")\n            ax[i, 0].axis(\"off\")\n            ax[i, 1].axis(\"off\")\n\n            prediction = tf.keras.preprocessing.image.array_to_img(prediction)\n            prediction.save(\n                \"generated_img_{i}_{epoch}.png\".format(i=i, epoch=epoch + 1)\n            )\n        plt.show()\n        plt.close()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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\n    \ndef generator_loss(generated):\n        return tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated), generated)\n    \ndef calc_cycle_loss(real_image, cycled_image, LAMBDA):\n        loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n        return LAMBDA * loss1\n    \ndef identity_loss(real_image, same_image, LAMBDA):\n        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n        return LAMBDA * 0.5 * loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Loss function for evaluating adversarial loss\nadv_loss_fn = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.AUTO)\n\n# Define the loss function for the generators\ndef generator_loss_fn(fake):\n    fake_loss = adv_loss_fn(tf.ones_like(fake), fake)\n    return fake_loss\n\n\n# Define the loss function for the discriminators\ndef discriminator_loss_fn(real, fake):\n    real_loss = adv_loss_fn(tf.ones_like(real), real)\n    fake_loss = adv_loss_fn(tf.zeros_like(fake), fake)\n    return (real_loss + fake_loss) * 0.5\n\n\n# Create cycle gan model\ncycle_gan_model = CycleGan(\n    generator_G=generator_g, generator_F=generator_f, \n    discriminator_X=discriminator_x, discriminator_Y=discriminator_y\n)\n\n# Compile the model\ncycle_gan_model.compile(\n    gen_G_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    gen_F_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    disc_X_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n    disc_Y_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\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)\n# Callbacks\nplotter = GANMonitor()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cycle_gan_model.fit(\n    tf.data.Dataset.zip((photo_ds, monet_ds)),\n    epochs=15,\n    callbacks=[plotter],\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path()\n\nMONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\nPHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\n\nBUFFER_SIZE = 1000\nBATCH_SIZE = 1\nIMG_WIDTH = 256\nIMG_HEIGHT = 256\n\nIMAGE_SIZE = [IMG_WIDTH, IMG_HEIGHT]\n\ndef decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.image.resize(image, IMAGE_SIZE)\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        \"image_name\": tf.io.FixedLenFeature([], tf.string),\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    return dataset\n\nmonet_ds = load_dataset(MONET_FILENAMES, labeled=True).batch(BATCH_SIZE)\nphoto_ds = load_dataset(PHOTO_FILENAMES, labeled=True).batch(BATCH_SIZE)\n\nexample_monet = next(iter(monet_ds))\nexample_photo = next(iter(photo_ds))\n\nplt.subplot(121)\nplt.title('Photo')\nplt.imshow(example_photo[0] * 0.5 + 0.5)\n\nplt.subplot(122)\nplt.title('Monet')\nplt.imshow(example_monet[0] * 0.5 + 0.5)\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cycle_gan_model.fit(\n    tf.data.Dataset.zip((photo_ds, monet_ds)),\n    epochs=15,\n    callbacks=[plotter],\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\n! mkdir ../images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"i = 1\nfor img in tqdm(photo_ds, total=len(os.listdir(\"/kaggle/input/gan-getting-started/photo_jpg/\"))):\n    prediction = generator_g(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    im = PIL.Image.fromarray(prediction)\n    im.save(\"../images/\" + str(i) + \".jpg\")\n    i += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### **Next Step:** Implement from scratch","execution_count":null}],"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}