{"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":"# 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)\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport matplotlib.pyplot as plt\nimport os\nimport time\nfrom IPython.display import clear_output\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# 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","execution":{"iopub.status.busy":"2022-10-20T04:30:34.159448Z","iopub.execute_input":"2022-10-20T04:30:34.160026Z","iopub.status.idle":"2022-10-20T04:30:44.895721Z","shell.execute_reply.started":"2022-10-20T04:30:34.159868Z","shell.execute_reply":"2022-10-20T04:30:44.894679Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loading, Decoding tfrecords and using data augmentation on images","metadata":{}},{"cell_type":"code","source":"class dataset_pipeline():\n    def __init__(self) -> None:\n        self.image_feature_description = {\n            'image_name': tf.io.FixedLenFeature([], tf.string),\n            'image': tf.io.FixedLenFeature([], tf.string),\n            'target' : tf.io.FixedLenFeature([], tf.string)\n        }\n        self.IMAGE_SIZE= [256,256]\n        self.IMG_HEIGHT = 256\n        self.IMG_WIDTH = 256\n\n    def decode_image(self, image):\n        \"\"\"Decoding image from tf.string form representation from tfrecord\"\"\"\n        image = tf.image.decode_jpeg(image, channels=3)\n        image = (tf.cast(image, tf.float32)/127.5) -1\n        image = tf.reshape(image, [*self.IMAGE_SIZE,3])\n        return image\n    \n    def random_augment(self, image):\n        image = tf.image.resize(image, [286,286], method= tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n        image = tf.image.random_crop(image, size=[self.IMG_HEIGHT, self.IMG_WIDTH,3])\n        image = tf.image.random_flip_left_right(image)\n        return image\n    \n    def parse_image(self, example):        \n        parsed_example = tf.io.parse_single_example(example, self.image_feature_description)\n        image = self.decode_image(parsed_example[\"image\"])\n        image = self.random_augment(image)\n        return image\n\n    def get_paths(self, directory_path):\n        \"\"\"directory_path : path to the directory containing all tfrecords\n        returns the list of all tfrecords\"\"\"\n        path_list = []\n        for path in os.listdir(directory_path):\n            path_list.append(os.path.join(directory_path, path))\n        return path_list\n\n    def load_tfrecord_dataset(self, directory_path):\n        \"\"\"directory_path : list of paths of all tfrecords.\"\"\"\n        paths =  self.get_paths(directory_path)\n        dataset = tf.data.TFRecordDataset(paths)\n        dataset = dataset.map(self.parse_image, num_parallel_calls= tf.data.AUTOTUNE)\n        dataset = dataset.map(self.random_augment, num_parallel_calls= tf.data.AUTOTUNE)\n        return dataset\n\n    def plot_sample(self, dataset):\n        sample= next(iter(dataset))\n        plt.imshow(sample[0]*0.5 + 0.5)\n        plt.axis('off')\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:35:30.051600Z","iopub.execute_input":"2022-10-20T04:35:30.052249Z","iopub.status.idle":"2022-10-20T04:35:30.076084Z","shell.execute_reply.started":"2022-10-20T04:35:30.052213Z","shell.execute_reply":"2022-10-20T04:35:30.074243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_tf = \"../input/gan-getting-started/monet_tfrec\"\nphoto_tf = \"../input/gan-getting-started/photo_tfrec\"","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:31:22.599775Z","iopub.execute_input":"2022-10-20T04:31:22.600478Z","iopub.status.idle":"2022-10-20T04:31:22.604964Z","shell.execute_reply.started":"2022-10-20T04:31:22.600444Z","shell.execute_reply":"2022-10-20T04:31:22.603949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_pipe = dataset_pipeline()\nMONET_DS = data_pipe.load_tfrecord_dataset(monet_tf).batch(1)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:36:47.089856Z","iopub.execute_input":"2022-10-20T04:36:47.090297Z","iopub.status.idle":"2022-10-20T04:36:47.171395Z","shell.execute_reply.started":"2022-10-20T04:36:47.090263Z","shell.execute_reply":"2022-10-20T04:36:47.170341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_pipe.plot_sample(MONET_DS)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:36:48.573999Z","iopub.execute_input":"2022-10-20T04:36:48.574375Z","iopub.status.idle":"2022-10-20T04:36:48.837165Z","shell.execute_reply.started":"2022-10-20T04:36:48.574345Z","shell.execute_reply":"2022-10-20T04:36:48.836038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PHOTO_DS = data_pipe.load_tfrecord_dataset(photo_tf).batch(1)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:37:29.159515Z","iopub.execute_input":"2022-10-20T04:37:29.159945Z","iopub.status.idle":"2022-10-20T04:37:29.235967Z","shell.execute_reply.started":"2022-10-20T04:37:29.159906Z","shell.execute_reply":"2022-10-20T04:37:29.234932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_pipe.plot_sample(PHOTO_DS)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:37:30.904763Z","iopub.execute_input":"2022-10-20T04:37:30.905693Z","iopub.status.idle":"2022-10-20T04:37:31.139997Z","shell.execute_reply.started":"2022-10-20T04:37:30.905660Z","shell.execute_reply":"2022-10-20T04:37:31.138747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"markdown","source":"## Instance Normalization","metadata":{}},{"cell_type":"code","source":"class InstanceNormalization(layers.Layer):\n    \"\"\"Instance Normalization layer\"\"\"\n    def __init__(self, epsilon= 1e-5):\n        super(InstanceNormalization, self).__init__()\n        self.epsilon= epsilon\n\n    def build(self, input_shape):\n        self.scale = self.add_weight(name= 'scale', shape= input_shape[-1:], initializer= tf.random_normal_initializer(1, 0.02), trainable= True)\n        self.offset = self.add_weight(name='offset', shape= input_shape[-1:], initializer= 'zeros', trainable= True)\n\n    def call(self, x):\n        mean, variance = tf.nn.moments(x, axes=[1,2], keepdims=True)\n        inv = tf.math.rsqrt(variance + self.epsilon)\n        normalized= (x-mean)*inv\n        return self.scale*normalized + self.offset","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:44:50.759963Z","iopub.execute_input":"2022-10-20T04:44:50.760359Z","iopub.status.idle":"2022-10-20T04:44:50.769270Z","shell.execute_reply.started":"2022-10-20T04:44:50.760328Z","shell.execute_reply":"2022-10-20T04:44:50.768157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"class Model:\n\n    def __init__(self, output_channels=3):                     \n        self.cross_entropy_loss = keras.losses.BinaryCrossentropy(from_logits=True)\n\n        self.OUTPUT_CHANNELS= output_channels\n        self.LAMBDA = 10\n\n        self.generator_g = self.Generator()\n        self.generator_f = self.Generator()\n\n        self.discriminator_x = self.Discriminator(target= False)\n        self.discriminator_y = self.Discriminator(target= False)\n\n        self.generator_g_optimizer = keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.generator_f_optimizer = keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.discriminator_x_optimizer = keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.discriminator_y_optimizer = keras.optimizers.Adam(2e-4, beta_1=0.5)\n        \n\n    def upsample(self, filters, size, apply_dropout=False):\n        initializer = tf.random_normal_initializer(0., 0.02)\n        result = keras.Sequential()\n        result.add(layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer= initializer, use_bias= False))\n\n        if apply_dropout:            \n            result.add(layers.Dropout(0.5))\n\n        result.add(layers.LeakyReLU())\n        return result\n\n\n    def downsample(self, filters, size, apply_norm = True):\n        initializer = tf.random_normal_initializer(0., 0.02)\n        result = keras.Sequential()\n        result.add(layers.Conv2D(filters, size, strides=2, padding='same', kernel_initializer= initializer, use_bias= False))\n\n        if apply_norm:\n            result.add(InstanceNormalization())\n        result.add(layers.LeakyReLU())\n        return result\n\n    def discriminator_loss(self, real, generated):\n        real_loss = self.cross_entropy_loss(tf.ones_like(real), real)\n        gen_loss = self.cross_entropy_loss(tf.zeros_like(generated), generated)\n        total_loss = real_loss + gen_loss\n        return total_loss*0.5\n    \n    def generator_loss(self, preds_generated):  \t\t\n        return self.cross_entropy_loss(tf.ones_like(preds_generated), preds_generated)\n    \n    def cycle_loss(self, real_image, cycled_image):\n        loss = tf.reduce_mean(tf.abs(real_image-cycled_image))\n        return loss*self.LAMBDA\n\n    def identity_loss(self, real_image, same_image):\n\n        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n        return loss*self.LAMBDA*0.5\n\n\n    def Generator(self):\n\n        down_stack=[\n          self.downsample(64,4, apply_norm=False),\n          self.downsample(128,4),\n          self.downsample(256,4),\n          self.downsample(512,4),\n          self.downsample(512,4),\n          self.downsample(512,4),\n          self.downsample(512,4),\n          self.downsample(512,4),\n          ]\n\n        up_stack = [\n          self.upsample(512, 4, apply_dropout=True),\n          self.upsample(512, 4, apply_dropout=True),\n          self.upsample(512, 4, apply_dropout=True),\n          self.upsample(512, 4),\n          self.upsample(256, 4),\n          self.upsample(128, 4),\n          self.upsample(64, 4),\n          ]\n        \n        initializer = tf.random_normal_initializer(0., 0.02)\n        last = layers.Conv2DTranspose(self.OUTPUT_CHANNELS, 4, strides=2, padding='same', kernel_initializer= initializer, activation= 'tanh')\n\n        concat = layers.Concatenate()\n        inputs= layers.Input(shape=[None, None, 3])\n\n        x= inputs\n\n        skips= []\n\n        for down in down_stack:            \n            x = down(x)\n            skips.append(x)\n        \n        skips = reversed(skips[:-1])\n\n        for up,skip in zip(up_stack, skips):\n            x= up(x)\n            x= concat([x, skip])\n\n        x= last(x)\n        return keras.Model(inputs= inputs, outputs= x)\n\n    def Discriminator(self, target=True):\n\n        initializer = tf.random_normal_initializer(0., 0.02)\n        inp= layers.Input(shape=[None, None, 3], name='input_image')\n        x= inp\n\n        if target:            \n            tar = layers.Input(shape=[None, None, 3], name='target_image')\n            x = layers.Concatenate([inp,tar])\n\n        down1= self.downsample(64,4, apply_norm=False)(x)\n        down2= self.downsample(128,4)(down1)\n        down3= self.downsample(256, 4)(down2)\n\n        zero_pad1= layers.ZeroPadding2D()(down3)\n        conv= layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1)\n        norm1= InstanceNormalization()(conv)\n        leaky_relu= layers.LeakyReLU()(norm1)\n        zero_pad2= layers.ZeroPadding2D()(leaky_relu)\n        last = layers.Conv2D(1, 4, strides= 1, kernel_initializer=initializer)(zero_pad2)\n\n        if target:\n            return keras.Model(inputs= [inp, last], outputs= last)\n        else:            \n            return keras.Model(inputs= inp, outputs= last)\n\n        \n    @tf.function\n    def train_step(self, real_x, real_y):\n\n\n        with tf.GradientTape(persistent=True) as tape:\n            \n            # Cycled predictions, output of cycle should be same as input provided\n            # Let generator_g transform x->y , and generator_f transform y->x\n            # discriminator_x takes x and classify, and discriminator_y takes y and classify\n\n            # 1. generator_g takes x, generate fake_y, then generator_f takes fake_y and generated, cycle_x (x,cycle_x should be similar)\n\n            fake_y = self.generator_g(real_x, training=True)\n            cycled_x = self.generator_f(fake_y, training= True)\n\n            # 2. Same cycle but starting from generator_f\n\n            fake_x = self.generator_f(real_y, training=True)\n            cycled_y = self.generator_g(fake_x, training= True)\n\n            # 3. Case when input given to generator is same as expected output\n\n            same_y = self.generator_g(real_y, training=True)\n            same_x = self.generator_f(real_x, training= True)\n\n            disc_pred_real_x = self.discriminator_x(real_x, training=True)\n            disc_pred_fake_x = self.discriminator_x(fake_x, training=True)\n\n            disc_pred_real_y = self.discriminator_y(real_y, training=True)\n            disc_pred_fake_y = self.discriminator_y(fake_y, training=True)\n\n            # All outputs generated, calculating lossess from these\n            total_cycle_loss = self.cycle_loss(real_x, cycled_x) + self.cycle_loss(real_y, cycled_y)\n\n            gen_g_identity_loss = self.identity_loss(real_y, same_y)\n            gen_f_identity_loss = self.identity_loss(real_x, same_x)\n\n            gen_g_loss = self.generator_loss(disc_pred_fake_y)\n            gen_f_loss = self.generator_loss(disc_pred_fake_x)\n\n            total_gen_g_loss = gen_g_loss + gen_g_identity_loss + total_cycle_loss\n            total_gen_f_loss = gen_f_loss + gen_f_identity_loss + total_cycle_loss\n\n            disc_x_loss = self.discriminator_loss(disc_pred_real_x, disc_pred_fake_x)\n            disc_y_loss = self.discriminator_loss(disc_pred_real_y, disc_pred_fake_y)\n                                                    \n        gen_g_grads = tape.gradient(total_gen_g_loss, self.generator_g.trainable_variables)\n        gen_f_grads = tape.gradient(total_gen_f_loss, self.generator_f.trainable_variables)\n\n        disc_x_grads = tape.gradient(disc_x_loss, self.discriminator_x.trainable_variables)\n        disc_y_grads = tape.gradient(disc_y_loss, self.discriminator_y.trainable_variables)\n\n        self.generator_g_optimizer.apply_gradients(zip(gen_g_grads, self.generator_g.trainable_variables ))\n        self.generator_f_optimizer.apply_gradients(zip(gen_f_grads, self.generator_f.trainable_variables ))\n\n        self.discriminator_x_optimizer.apply_gradients(zip(disc_x_grads, self.discriminator_x.trainable_variables ))\n        self.discriminator_y_optimizer.apply_gradients(zip(disc_y_grads, self.discriminator_y.trainable_variables ))\n        \n\n    def train_image(self, photo, gen_model,fig_size=(12,12)):\n        \"\"\"Plotting image and prediction during training\"\"\"\n        predicted_image = gen_model(photo)\n        titles = [\"Input image\", \"Predicted Image\"]\n        Images = [photo[0], predicted_image[0]]\n        plt.figure(figsize=fig_size)\n\n        for i in range(2):\n            plt.subplot(1,2,i+1)\n            plt.title(titles[i])\n            plt.imshow(Images[i]*0.5 + 0.5)\n            plt.axis('off')\n        plt.show()\n\n    def generate_images(self, dataset, gen_model, fig_size=(12,12),num=5):\n        \"\"\"saving and plotting predictions in bulk after training\"\"\"\n        plt.figure(figsize=fig_size)\n        for img in dataset.take(num):            \n            self.train_image(img,gen_model)\n\n    def train(self, train_x, train_y, vis=\"g\", epochs=10):\n        if(vis.lower() == \"g\"):\n            photo = next(iter(train_x))\n            gen_model = self.generator_g\n        elif(vis.lower()==\"f\"):\n            photo = next(iter(train_y))\n            gen_model = self.generator_f\n            \n        print(\"Before Training\")\n        self.train_image(photo, gen_model)\n        for epoch in range(epochs):            \n            start = time.time()\n            for image_x, image_y in tf.data.Dataset.zip((train_x, train_y)):\n                self.train_step(image_x, image_y)\n            clear_output(wait=True)\n            self.train_image(photo,gen_model)\n            print(f\"Time for epoch {epoch+1} is : {time.time()- start}\")","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:59:14.487463Z","iopub.execute_input":"2022-10-20T04:59:14.487862Z","iopub.status.idle":"2022-10-20T04:59:14.528144Z","shell.execute_reply.started":"2022-10-20T04:59:14.487813Z","shell.execute_reply":"2022-10-20T04:59:14.527141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = Model()","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:59:17.565872Z","iopub.execute_input":"2022-10-20T04:59:17.566252Z","iopub.status.idle":"2022-10-20T04:59:18.814409Z","shell.execute_reply.started":"2022-10-20T04:59:17.566223Z","shell.execute_reply":"2022-10-20T04:59:18.813321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample = next(iter(PHOTO_DS))","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:54:54.926424Z","iopub.execute_input":"2022-10-20T04:54:54.926784Z","iopub.status.idle":"2022-10-20T04:54:55.037140Z","shell.execute_reply.started":"2022-10-20T04:54:54.926752Z","shell.execute_reply":"2022-10-20T04:54:55.035980Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.train_image(sample, model.generator_g)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:57:29.824738Z","iopub.execute_input":"2022-10-20T04:57:29.825221Z","iopub.status.idle":"2022-10-20T04:57:30.184236Z","shell.execute_reply.started":"2022-10-20T04:57:29.825184Z","shell.execute_reply":"2022-10-20T04:57:30.182801Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.train(PHOTO_DS, MONET_DS)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T04:59:40.456623Z","iopub.execute_input":"2022-10-20T04:59:40.457005Z","iopub.status.idle":"2022-10-20T05:20:29.357955Z","shell.execute_reply.started":"2022-10-20T04:59:40.456973Z","shell.execute_reply":"2022-10-20T05:20:29.357017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.generate_images(PHOTO_DS, model.generator_g)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T05:20:46.953356Z","iopub.execute_input":"2022-10-20T05:20:46.953711Z","iopub.status.idle":"2022-10-20T05:20:48.865902Z","shell.execute_reply.started":"2022-10-20T05:20:46.953680Z","shell.execute_reply":"2022-10-20T05:20:48.864840Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.generate_images(MONET_DS, model.generator_f)","metadata":{"execution":{"iopub.status.busy":"2022-10-20T05:21:10.624728Z","iopub.execute_input":"2022-10-20T05:21:10.625430Z","iopub.status.idle":"2022-10-20T05:21:12.472621Z","shell.execute_reply.started":"2022-10-20T05:21:10.625394Z","shell.execute_reply":"2022-10-20T05:21:12.471426Z"},"trusted":true},"execution_count":null,"outputs":[]}]}