{"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 tensorflow as tf\nimport tensorflow.keras as keras\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport tensorflow_addons as tfa\nimport PIL\nfrom kaggle_datasets import KaggleDatasets\n\nfrom tensorflow.keras.initializers import RandomNormal\nfrom tensorflow.keras.layers import Input, Conv2D, Activation\nfrom tensorflow.keras.models import Sequential\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()\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-12-25T03:07:54.044338Z","iopub.execute_input":"2022-12-25T03:07:54.044678Z","iopub.status.idle":"2022-12-25T03:07:54.053431Z","shell.execute_reply.started":"2022-12-25T03:07:54.044651Z","shell.execute_reply":"2022-12-25T03:07:54.052076Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 1 **Variables**","metadata":{}},{"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path()\n\nMONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\nprint('Monet TFRecord Files:', len(MONET_FILENAMES))\nPHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\nprint('Monet TFRecord Files:', len(PHOTO_FILENAMES))\n\nimg_width = 256\nimg_height = 256\nbatch_size = 16\nchannels = 3","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:07:54.055534Z","iopub.execute_input":"2022-12-25T03:07:54.055800Z","iopub.status.idle":"2022-12-25T03:07:55.612914Z","shell.execute_reply.started":"2022-12-25T03:07:54.055775Z","shell.execute_reply":"2022-12-25T03:07:55.612204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 2 **Reading, Augmenting, and Displaying the images**","metadata":{}},{"cell_type":"markdown","source":"**Decoding**: Used tensorflow functions to decode the jpg and reshape the images","metadata":{}},{"cell_type":"markdown","source":"**Augmentation**: \n* 50% chance of training data to be reflected\n* crop out some parts of the image","metadata":{}},{"cell_type":"markdown","source":"**Displayment**: Using matplotlib to display the augmented image","metadata":{}},{"cell_type":"code","source":"def 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.reshape(image, [*[img_width, img_height], 3])\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\n\ndef data_augment(image):\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_crop(image, size=[img_width, img_height, 3])\n    return image\n\n\ndef load_dataset(filenames, labeled=True, ordered=False, augment=True):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    if augment:\n    \tdataset = dataset.map(data_augment, num_parallel_calls=AUTOTUNE)\n    \tdataset = dataset.repeat(count=2)\n    return dataset\n\nmonet_img = load_dataset(MONET_FILENAMES, labeled=True, augment=True).batch(1)\nphoto_img = load_dataset(PHOTO_FILENAMES, labeled=True, augment=True).batch(1)\n\ndef show_img():\n\texample_monet = next(iter(monet_img))\n\texample_photo = next(iter(photo_img))\n\n\tplt.subplot(121)\n\tplt.title('Photo')\n\tplt.imshow(example_photo[0] * 0.5 + 0.5)\n\n\tplt.subplot(122)\n\tplt.title('Monet')\n\tplt.imshow(example_monet[0] * 0.5 + 0.5)\n\n\tplt.show()\n    \nshow_img()","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:07:55.614319Z","iopub.execute_input":"2022-12-25T03:07:55.614709Z","iopub.status.idle":"2022-12-25T03:07:59.705145Z","shell.execute_reply.started":"2022-12-25T03:07:55.614685Z","shell.execute_reply":"2022-12-25T03:07:59.703706Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 3 **Generator**","metadata":{}},{"cell_type":"markdown","source":"I used the normal upsample and downsample method that everyone else used (which is great to get rid of the vanishing gradient), but I also added a resnet in the middle to see if that would help with the vanishing gradient even more","metadata":{}},{"cell_type":"code","source":"\ndef residual_block(filters, input_layer):\n    init = tf.keras.initializers.RandomNormal(stddev=0.02)\n\n    r = tf.keras.layers.Conv2D(filters, (3, 3), padding='same', kernel_initializer=init)(input_layer)\n    r = tfa.layers.InstanceNormalization(axis=-1)(r)\n    r = tf.keras.layers.LeakyReLU(alpha=0.2)(r)\n    \n    r = tf.keras.layers.Conv2D(filters, (3, 3), padding='same', kernel_initializer=init)(r)\n    r = tfa.layers.InstanceNormalization(axis =-1)(r)\n    \n    r = tf.keras.layers.Concatenate()([r,input_layer])\n    \n    return r\n\n\ndef downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    result.add(tf.keras.layers.LeakyReLU())\n\n    return result\n\n\ndef upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(tf.keras.layers.Conv2DTranspose(filters, size, strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n\n    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout:\n        result.add(tf.keras.layers.Dropout(0.5))\n\n    result.add(tf.keras.layers.ReLU())\n\n    return result\n\n\ndef Generator(image_shape=(img_width, img_height, 3), n_resnet=2):\n\tinput_layer = tf.keras.layers.Input(shape=image_shape)\n\n\tdown_stack = [\n\t    downsample(64, 4, apply_instancenorm=False),\n\t    downsample(128, 4),\n\t    downsample(256, 4),\n\t    downsample(512, 4),\n\t    downsample(512, 4),\n\t    downsample(512, 4),\n\t    downsample(512, 4),\n\t    downsample(512, 4),\n\t]\n\n\tup_stack = [\n\t    upsample(512, 4, apply_dropout=True),\n\t    upsample(512, 4, apply_dropout=True),\n\t    upsample(512, 4, apply_dropout=True),\n\t    upsample(512, 4),\n\t    upsample(256, 4),\n\t    upsample(128, 4),\n\t    upsample(64, 4),\n\t]\n\n\tinitializer = tf.random_normal_initializer(0., 0.02)\n\tlast = tf.keras.layers.Conv2DTranspose(channels, 4,\n\t                              strides=2,\n\t                              padding='same',\n\t                              kernel_initializer=initializer,\n\t                              activation='tanh')\n\n\tx = input_layer\n\n\tskips = []\n\tfor down in down_stack:\n\t    x = down(x)\n\t    skips.append(x)\n\n\tskips = reversed(skips[:-1])\n    \n\tfor l in range(n_resnet):\n\t\tx = residual_block(256, x)\n\n\tfor up, skip in zip(up_stack, skips):\n\t    x = up(x)\n\t    x = tf.keras.layers.Concatenate()([x, skip])\n\n\tx = last(x)\n\n\tmodel = tf.keras.models.Model(input_layer, x)\n\tmodel.summary()\n\treturn model","metadata":{"execution":{"iopub.status.busy":"2022-12-28T02:31:02.478598Z","iopub.execute_input":"2022-12-28T02:31:02.480173Z","iopub.status.idle":"2022-12-28T02:31:02.581667Z","shell.execute_reply.started":"2022-12-28T02:31:02.480116Z","shell.execute_reply":"2022-12-28T02:31:02.580152Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 4 **Discriminator**","metadata":{}},{"cell_type":"markdown","source":"Normally a CNN would work best but I've found that this model that everyone used worked even better.","metadata":{}},{"cell_type":"code","source":"def Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    inp = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')\n\n    x = inp\n\n    down1 = downsample(64, 4, False)(x)\n    down2 = downsample(128, 4)(down1)\n    down3 = downsample(256, 4)(down2)\n\n    zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)\n    conv = tf.keras.layers.Conv2D(512, 4, strides=1,\n                         kernel_initializer=initializer,\n                         use_bias=False)(zero_pad1)\n\n    norm1 = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(conv)\n\n    leaky_relu = tf.keras.layers.LeakyReLU()(norm1)\n\n    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)\n\n    last = tf.keras.layers.Conv2D(1, 4, strides=1,\n                         kernel_initializer=initializer)(zero_pad2)\n\n    return tf.keras.Model(inputs=inp, outputs=last)","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:07:59.727718Z","iopub.execute_input":"2022-12-25T03:07:59.728002Z","iopub.status.idle":"2022-12-25T03:07:59.746684Z","shell.execute_reply.started":"2022-12-25T03:07:59.727978Z","shell.execute_reply":"2022-12-25T03:07:59.745061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Visualize**","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    monet_generator = Generator()\n    photo_generator = Generator()\n\n    monet_discriminator = Discriminator()\n    photo_discriminator = Discriminator()\n\n\ndef show_generator():\n\texample_monet = next(iter(monet_img))\n\texample_photo = next(iter(photo_img))\n\n\tto_monet = monet_generator(example_photo)\n\n\tplt.subplot(1, 2, 1)\n\tplt.title(\"Original Photo\")\n\tplt.imshow(example_photo[0] * 0.5 + 0.5)\n\n\tplt.subplot(1, 2, 2)\n\tplt.title(\"Monet-esque Photo\")\n\tplt.imshow(to_monet[0] * 0.5 + 0.5)\n\tplt.show()\n\nshow_generator()","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:07:59.747768Z","iopub.execute_input":"2022-12-25T03:07:59.748151Z","iopub.status.idle":"2022-12-25T03:08:18.776456Z","shell.execute_reply.started":"2022-12-25T03:07:59.748078Z","shell.execute_reply":"2022-12-25T03:08:18.775012Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 5 **Combine**","metadata":{}},{"cell_type":"markdown","source":"The training loop consists of four basic steps:\n* Get the predictions\n* Calculate the loss\n* Backpropagation\n* Apply gradients to the optimizer","metadata":{}},{"cell_type":"code","source":"class CycleGan(keras.Model):\n    def __init__(\n        self,\n        monet_generator,\n        photo_generator,\n        monet_discriminator,\n        photo_discriminator,\n        lambda_cycle=10,\n    ):\n        super(CycleGan, self).__init__()\n        self.m_gen = monet_generator\n        self.p_gen = photo_generator\n        self.m_disc = monet_discriminator\n        self.p_disc = photo_discriminator\n        self.lambda_cycle = lambda_cycle\n        \n    def compile(\n        self,\n        m_gen_optimizer,\n        p_gen_optimizer,\n        m_disc_optimizer,\n        p_disc_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n        cycle_loss_fn,\n        identity_loss_fn\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.gen_loss_fn = gen_loss_fn\n        self.disc_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        real_monet, real_photo = batch_data\n        \n        with tf.GradientTape(persistent=True) as tape:\n            fake_monet = self.m_gen(real_photo, training=True)\n            cycled_photo = self.p_gen(fake_monet, training=True)\n\n            fake_photo = self.p_gen(real_monet, training=True)\n            cycled_monet = self.m_gen(fake_photo, training=True)\n\n            same_monet = self.m_gen(real_monet, training=True)\n            same_photo = self.p_gen(real_photo, training=True)\n\n            disc_real_monet = self.m_disc(real_monet, training=True)\n            disc_real_photo = self.p_disc(real_photo, training=True)\n\n            disc_fake_monet = self.m_disc(fake_monet, training=True)\n            disc_fake_photo = self.p_disc(fake_photo, training=True)\n\n            monet_gen_loss = self.gen_loss_fn(disc_fake_monet)\n            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)\n\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\n            total_monet_gen_loss = monet_gen_loss + total_cycle_loss + self.identity_loss_fn(real_monet, same_monet, self.lambda_cycle)\n            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo, self.lambda_cycle)\n\n            monet_disc_loss = self.disc_loss_fn(disc_real_monet, disc_fake_monet)\n            photo_disc_loss = self.disc_loss_fn(disc_real_photo, disc_fake_photo)\n\n        monet_generator_gradients = tape.gradient(total_monet_gen_loss,\n                                                  self.m_gen.trainable_variables)\n        photo_generator_gradients = tape.gradient(total_photo_gen_loss,\n                                                  self.p_gen.trainable_variables)\n\n        monet_discriminator_gradients = tape.gradient(monet_disc_loss,\n                                                      self.m_disc.trainable_variables)\n        photo_discriminator_gradients = tape.gradient(photo_disc_loss,\n                                                      self.p_disc.trainable_variables)\n\n        self.m_gen_optimizer.apply_gradients(zip(monet_generator_gradients,\n                                                 self.m_gen.trainable_variables))\n\n        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients,\n                                                 self.p_gen.trainable_variables))\n\n        self.m_disc_optimizer.apply_gradients(zip(monet_discriminator_gradients,\n                                                  self.m_disc.trainable_variables))\n\n        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients,\n                                                  self.p_disc.trainable_variables))\n        \n        return {\n            \"monet_gen_loss\": total_monet_gen_loss,\n            \"photo_gen_loss\": total_photo_gen_loss,\n            \"monet_disc_loss\": monet_disc_loss,\n            \"photo_disc_loss\": photo_disc_loss\n        }","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:08:18.778005Z","iopub.execute_input":"2022-12-25T03:08:18.778322Z","iopub.status.idle":"2022-12-25T03:08:18.795449Z","shell.execute_reply.started":"2022-12-25T03:08:18.778294Z","shell.execute_reply":"2022-12-25T03:08:18.794454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 6 **Loss Functions**","metadata":{}},{"cell_type":"markdown","source":"In CycleGAN, there is no paired data to train on, hence there is no guarantee that the input x and the target y pair are meaningful during training. Thus in order to enforce that the network learns the correct mapping, the authors propose the cycle consistency loss.","metadata":{}},{"cell_type":"markdown","source":"Cycle consistency means the result should be close to the original input. For example, if one translates a sentence from English to French, and then translates it back from French to English, then the resulting sentence should be the same as the original sentence.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    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\n    def generator_loss(generated):\n    \treturn tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated), generated)\n\n    def calc_cycle_loss(real_image, cycled_image, LAMBDA):\n    \tloss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n    \treturn LAMBDA * loss1\n\n    def identity_loss(real_image, same_image, LAMBDA):\n    \tloss = tf.reduce_mean(tf.abs(real_image - same_image))\n    \treturn LAMBDA * 0.5 * loss\n\n    monet_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\n    monet_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    \n    cycle_gan_model = CycleGan(\n        monet_generator, photo_generator, monet_discriminator, photo_discriminator\n    )\n    \n    cycle_gan_model.compile(\n        m_gen_optimizer = monet_generator_optimizer,\n        p_gen_optimizer = photo_generator_optimizer,\n        m_disc_optimizer = monet_discriminator_optimizer,\n        p_disc_optimizer = photo_discriminator_optimizer,\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    )","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:08:18.796607Z","iopub.execute_input":"2022-12-25T03:08:18.796920Z","iopub.status.idle":"2022-12-25T03:08:18.829777Z","shell.execute_reply.started":"2022-12-25T03:08:18.796894Z","shell.execute_reply":"2022-12-25T03:08:18.828560Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Fitting**","metadata":{}},{"cell_type":"code","source":"cycle_gan_model.fit(\n    tf.data.Dataset.zip((monet_img, photo_img)),\n    epochs=25\n)","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:08:18.831791Z","iopub.execute_input":"2022-12-25T03:08:18.832289Z","iopub.status.idle":"2022-12-25T03:09:35.054394Z","shell.execute_reply.started":"2022-12-25T03:08:18.832250Z","shell.execute_reply":"2022-12-25T03:09:35.052421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"STEP 7 **Plot and Submit**","metadata":{}},{"cell_type":"code","source":"! mkdir ../images","metadata":{"execution":{"iopub.status.busy":"2022-12-28T21:37:09.556739Z","iopub.execute_input":"2022-12-28T21:37:09.557269Z","iopub.status.idle":"2022-12-28T21:37:10.693763Z","shell.execute_reply.started":"2022-12-28T21:37:09.557189Z","shell.execute_reply":"2022-12-28T21:37:10.692182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_, ax = plt.subplots(5, 2, figsize=(12, 12))\nfor i, img in enumerate(photo_img.take(5)):\n    prediction = monet_generator(img, training=False)[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 Photo\")\n    ax[i, 1].set_title(\"Monet-esque\")\n    ax[i, 0].axis(\"off\")\n    ax[i, 1].axis(\"off\")\nplt.show()\n\ni = 1\n\nfor img in photo_img:\n    prediction = monet_generator(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","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:09:35.055913Z","iopub.status.idle":"2022-12-25T03:09:35.056328Z","shell.execute_reply.started":"2022-12-25T03:09:35.056143Z","shell.execute_reply":"2022-12-25T03:09:35.056161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2022-12-25T03:09:35.057648Z","iopub.status.idle":"2022-12-25T03:09:35.058054Z","shell.execute_reply.started":"2022-12-25T03:09:35.057834Z","shell.execute_reply":"2022-12-25T03:09:35.057887Z"},"trusted":true},"execution_count":null,"outputs":[]}]}