{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport tensorflow_addons as tfa\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nfrom PIL import Image\nfrom kaggle_datasets import KaggleDatasets\nimport matplotlib.pyplot as plt","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_path = KaggleDatasets().get_gcs_path()\nmonet_files = tf.io.gfile.glob(str(test_path + '/monet_tfrec/*.tfrec'))\nphoto_files = tf.io.gfile.glob(test_path + '/photo_tfrec/*.tfrec')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nstrategy = tf.distribute.experimental.TPUStrategy(tpu)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"AUTOTUNE = tf.data.experimental.AUTOTUNE\nIMAGE_SIZE = [256, 256]\n\nfeatures = {'image': tf.io.FixedLenFeature([], tf.string)}\n\ndef read_tfrecord(example):\n    image_data = tf.io.parse_single_example(example, features)\n    image = image_data['image']\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_dataset = tf.data.TFRecordDataset(monet_files).map(read_tfrecord, num_parallel_calls=AUTOTUNE).batch(1)\nphoto_dataset = tf.data.TFRecordDataset(photo_files).map(read_tfrecord, num_parallel_calls=AUTOTUNE).batch(1)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example_photo = next(iter(photo_dataset))\nexample_monet = next(iter(monet_dataset))\n\nplt.subplot(122)\nplt.title('Monet')\nplt.imshow(example_monet[0] * 0.5 + 0.5)\n\nplt.subplot(121)\nplt.title('Photo')\nplt.imshow(example_photo[0] * 0.5 + 0.5)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\n\n\nclass CycleGan(keras.Model):\n    def __init__(self):\n        super(CycleGan, self).__init__()\n        self.monet_generator = self.create_generator()\n        self.photo_generator = self.create_generator()\n        self.monet_discriminator = self.create_discriminator()\n        self.photo_discriminator = self.create_discriminator()\n        self.lambda_cycle = 10\n\n    def compile(self):\n        super(CycleGan, self).compile()\n        self.monet_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.monet_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.photo_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.photo_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        \n    def create_downsampler(self, filters, size, apply_instance_norm=True):\n        model = keras.Sequential()\n        model.add(\n            layers.Conv2D(\n                filters,\n                size,\n                strides=2,\n                padding=\"same\",\n                use_bias=False,\n                kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n            )\n        )\n        if apply_instance_norm:\n            model.add(\n                tfa.layers.InstanceNormalization(\n                    gamma_initializer=keras.initializers.RandomNormal(\n                        mean=0.0, stddev=0.02\n                    )\n                )\n            )\n        model.add(layers.LeakyReLU())\n        return model\n\n    \n    def create_upsampler(self, filters, size, apply_dropout=False):\n        model = keras.Sequential()\n        model.add(\n            layers.Conv2DTranspose(\n                filters,\n                size,\n                strides=2,\n                padding=\"same\",\n                use_bias=False,\n                kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n            )\n        )\n        model.add(\n            tfa.layers.InstanceNormalization(\n                gamma_initializer=tf.random_normal_initializer(0.0, 0.02)\n            )\n        )\n        if apply_dropout:\n            model.add(layers.Dropout(0.5))\n        model.add(layers.ReLU())\n        return model\n\n    def create_generator(self):\n        downsampler_stack = [\n            self.create_downsampler(64, 4, apply_instance_norm=False),\n            self.create_downsampler(128, 4),\n            self.create_downsampler(256, 4),\n        ] + [self.create_downsampler(512, 4) for i in range(5)]\n        upsampler_stack = [\n            self.create_upsampler(512, 4, apply_dropout=True) for i in range(3)\n        ] + [\n            self.create_upsampler(512, 4),\n            self.create_upsampler(256, 4),\n            self.create_upsampler(128, 4),\n            self.create_upsampler(64, 4),\n        ]\n        input_layer = layers.Input(shape=[256, 256, 3])\n        x = input_layer\n        skips = []\n        for downsampler in downsampler_stack:\n            x = downsampler(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n        for upsampler, skip_layer in zip(upsampler_stack, skips):\n            x = upsampler(x)\n            x = layers.Concatenate()([x, skip_layer])\n        last_layer = layers.Conv2DTranspose(\n            OUTPUT_CHANNELS,\n            4,\n            strides=2,\n            padding=\"same\",\n            kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n            activation=\"tanh\",\n        )\n        x = last_layer(x)\n        return keras.Model(inputs=input_layer, outputs=x)\n\n    def create_discriminator(self):\n        input_layer = layers.Input(shape=[256, 256, 3], name=\"input_image\")\n        x = input_layer\n        downsampler1 = self.create_downsampler(64, 4, False)(x)\n        downsampler2 = self.create_downsampler(128, 4)(downsampler1)\n        downsampler3 = self.create_downsampler(256, 4)(downsampler2)\n        zero_pad1 = layers.ZeroPadding2D()(downsampler3)\n        conv_layer = layers.Conv2D(\n            512,\n            4,\n            strides=1,\n            use_bias=False,\n            kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n        )(zero_pad1)\n        normalization_layer1 = tfa.layers.InstanceNormalization(\n            gamma_initializer=tf.random_normal_initializer(0.0, 0.02)\n        )(conv_layer)\n        leaky_relu_layer = layers.LeakyReLU()(normalization_layer1)\n        zero_pad2 = layers.ZeroPadding2D()(leaky_relu_layer)\n        last_layer = layers.Conv2D(\n            1, 4, strides=1, kernel_initializer=tf.random_normal_initializer(0.0, 0.02)\n        )(zero_pad2)\n        return tf.keras.Model(inputs=input_layer, outputs=last_layer)\n\n    def discriminator_loss_fn(self, real, fake):\n        real_loss = tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )(tf.ones_like(real), real)\n        fake_loss = tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )(tf.zeros_like(fake), fake)\n        return (real_loss + fake_loss) / 2\n\n    def generator_loss_fn(self, generated_image):\n        return tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )(tf.ones_like(generated_image), generated_image)\n\n    def cycle_loss_fn(self, image, cycled_image, lambda_cycle):\n        return tf.reduce_mean(tf.abs(image - cycled_image)) * lambda_cycle\n\n    def identity_loss_fn(self, real_photo, photo, lambda_cycle):\n        return tf.reduce_mean(tf.abs(real_photo - photo)) * lambda_cycle / 2\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            fake_monet = self.monet_generator(real_photo, training=True)\n            cycled_photo = self.photo_generator(fake_monet, training=True)\n            fake_photo = self.photo_generator(real_monet, training=True)\n            cycled_monet = self.monet_generator(fake_photo, training=True)\n\n            monet1 = self.monet_generator(real_monet, training=True)\n            photo1 = self.photo_generator(real_photo, training=True)\n\n            monet_real_discriminated = self.monet_discriminator(\n                real_monet, training=True\n            )\n            monet_fake_discriminated = self.monet_discriminator(\n                fake_monet, training=True\n            )\n            photo_real_discriminated = self.photo_discriminator(\n                real_photo, training=True\n            )\n            photo_fake_discriminated = self.photo_discriminator(\n                fake_photo, training=True\n            )\n\n            monet_generator_loss = self.generator_loss_fn(monet_fake_discriminated)\n            photo_generator_loss = self.generator_loss_fn(photo_fake_discriminated)\n            cycle_loss = self.cycle_loss_fn(\n                real_monet, cycled_monet, self.lambda_cycle\n            ) + self.cycle_loss_fn(real_photo, cycled_photo, self.lambda_cycle)\n            total_monet_generator_loss = (\n                monet_generator_loss\n                + cycle_loss\n                + self.identity_loss_fn(real_monet, monet1, self.lambda_cycle)\n            )\n            total_photo_generator_loss = (\n                photo_generator_loss\n                + cycle_loss\n                + self.identity_loss_fn(real_photo, photo1, self.lambda_cycle)\n            )\n            monet_discriminator_loss = self.discriminator_loss_fn(\n                monet_real_discriminated, monet_fake_discriminated\n            )\n            photo_discriminator_loss = self.discriminator_loss_fn(\n                photo_real_discriminated, photo_fake_discriminated\n            )\n        monet_generator_gradients = tape.gradient(\n            total_monet_generator_loss, self.monet_generator.trainable_variables\n        )\n        photo_generator_gradients = tape.gradient(\n            total_photo_generator_loss, self.photo_generator.trainable_variables\n        )\n        monet_discriminator_gradients = tape.gradient(\n            monet_discriminator_loss, self.monet_discriminator.trainable_variables\n        )\n        photo_discriminator_gradients = tape.gradient(\n            photo_discriminator_loss, self.photo_discriminator.trainable_variables\n        )\n        self.monet_generator_optimizer.apply_gradients(\n            zip(monet_generator_gradients, self.monet_generator.trainable_variables)\n        )\n        self.photo_generator_optimizer.apply_gradients(\n            zip(photo_generator_gradients, self.photo_generator.trainable_variables)\n        )\n        self.monet_discriminator_optimizer.apply_gradients(\n            zip(\n                monet_discriminator_gradients,\n                self.monet_discriminator.trainable_variables,\n            )\n        )\n        self.photo_discriminator_optimizer.apply_gradients(\n            zip(\n                photo_discriminator_gradients,\n                self.photo_discriminator.trainable_variables,\n            )\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_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    cycle_gan_model = CycleGan()\n    cycle_gan_model.compile()\n    ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cycle_gan_model.fit(\n    tf.data.Dataset.zip((monet_dataset.repeat(-1), photo_dataset.repeat(-1))),\n    steps_per_epoch=300,\n    epochs=25\n)\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from io import BytesIO\nfrom PIL import Image\nfrom zipfile import ZipFile\n\n\nwith ZipFile('images.zip', mode='w') as zip_file:\n    i = 1\n    for img in photo_dataset:\n        generated_image_data = cycle_gan_model.monet_generator(img, training=False)[0].numpy()\n        scaled_generated_image_data = (generated_image_data * 127.5 + 127.5).astype(np.uint8)\n        with BytesIO() as image_bytes_io:\n            Image.fromarray(scaled_generated_image_data).save(image_bytes_io, 'JPEG')\n            image_bytes_io.seek(0)\n            zip_file.writestr('{}.jpg'.format(i), image_bytes_io.read())\n            i += 1","metadata":{},"execution_count":null,"outputs":[]}]}