{"cells":[{"metadata":{},"cell_type":"markdown","source":"The original paper for CycleGan was a huge help. It can help even if you are not a math wizard:\nhttps://arxiv.org/abs/1703.10593 \"Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks\""},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport pandas as pd  \nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\nfrom functools import partial\nfrom tqdm import tqdm\nfrom PIL import Image\n\nfrom kaggle_datasets import KaggleDatasets","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Setting up the random seeds for repruducibility"},{"metadata":{"trusted":true},"cell_type":"code","source":"seed = 159\n\ntf.random.set_seed(seed)\nnp.random.seed(seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"the main model parameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 1\nimg_size =  256","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Getting the file directories"},{"metadata":{"trusted":true},"cell_type":"code","source":"autotune = tf.data.experimental.AUTOTUNE\nGCS_PATH = KaggleDatasets().get_gcs_path()\n\n#file directories of the tfrec files\nmonet_imgs_fnames = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\nphoto_imgs_fnames = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef read_tfr(example):\n    \n    \"\"\"\n    function to read the TFrecords files, as can be seen\n    at the end only the image is returned leaving behind the label and te name\n    \"\"\"\n    \n    feature_map = {\n        \"image\" : tf.io.FixedLenFeature([], tf.string),\n        \"image_name\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.string)\n        }\n    \n    example = tf.io.parse_single_example(example, feature_map)\n    \n    image = tf.io.decode_jpeg(example['image'])\n    image = tf.image.resize(image, (256,256))\n    \n    image = image / 127.5 - 1\n    \n    \n    return image\n\n\n\ndef load_ds(filenames, aug):\n    \n    tfrecords = tf.data.TFRecordDataset(filenames)\n    tfrecords = tfrecords.map(read_tfr, num_parallel_calls = autotune)\n    if aug:\n        tfrecords = tfrecords.map(augment, num_parallel_calls = autotune)\n    return tfrecords\n\n\ndef get_ds(filenames, aug = True):\n    \n    \n    ds = load_ds(filenames, aug)\n    ds = ds.shuffle(2048)\n    ds = ds.batch(batch_size) #setting the dataset in batches\n    ds = ds.prefetch(autotune) # prefetch is used to improve performance\n    return ds\n\ndef augment(image):\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_adjust_brightness = tf.random.uniform([], 0, 1.0, dtype = tf.float32)\n    \n    image = tf.image.resize(image, (286,286))\n    \n    if p_rotate > .9:\n        image = tf.image.rot90(image, k=3) # rotate 270º\n    elif p_rotate > .7:\n        image = tf.image.rot90(image, k=2) # rotate 180º\n    elif p_rotate > .5:\n        image = tf.image.rot90(image, k=1) # rotate 90º\n        \n    if p_adjust_brightness > .9:\n        image = tf.image.adjust_brightness(image, 0.12)\n    elif p_adjust_brightness > .7:\n        image = tf.image.adjust_brightness(image, 0.08)\n    elif p_adjust_brightness > .5:\n        image = tf.image.adjust_brightness(image, 0.05)\n    \n    image = tf.image.random_crop(image, [256,256,3])\n    \n    \n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"monet_ds = get_ds(monet_imgs_fnames)\nphoto_ds = get_ds(photo_imgs_fnames)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Getting 5 examples from each dataset, the photo dataset and the monet Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"monet_iter = iter(monet_ds)\nphoto_iter = iter(photo_ds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_ = plt.figure(figsize=(10,30))\n\nfor i in range(8):\n\n    plt.subplot(8,2,i*2 +1)\n    plt.imshow(tf.cast(next(monet_iter)[0] * 127.5 + 127.5, tf.uint8))\n    plt.title(\"Monet image\")\n\n    plt.subplot(8,2,i*2 +2)\n    plt.imshow(tf.cast(next(photo_iter)[0] * 127.5 + 127.5, tf.uint8))\n    plt.title(\"Photo image\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The following block of code will be creating the upsample and downsample blocks of the modified Unet (the cycle gan model is based on the Unet model)"},{"metadata":{"trusted":true},"cell_type":"code","source":"# starting with the generator model\n\ndef downsample(filters, size, apply_batchnorm = True):\n    \n    \n    \n    initializer = tf.random_normal_initializer()\n\n        # the use of sequential is preferred to make the code simpler\n\n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2D(filters,\n                                         size,\n                                         strides = 2,\n                                         padding = 'same',\n                                         kernel_initializer = initializer,\n                                         use_bias = False))\n\n    if apply_batchnorm:\n        result.add(tf.keras.layers.BatchNormalization())\n\n    result.add(tf.keras.layers.LeakyReLU())\n\n    return result\n\n\ndef upsample(filters, size, apply_dropout = False):\n    \n   \n    \n    initialzer = tf.random_normal_initializer()\n\n    result = tf.keras.Sequential()\n\n    result.add(tf.keras.layers.Conv2DTranspose(filters,\n                                                  size,\n                                                  strides = 2,\n                                                  padding = 'same',\n                                                  kernel_initializer = initialzer,\n                                                  use_bias = False))\n\n    result.add(tf.keras.layers.BatchNormalization())\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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Generator model, which is the unet, the use of the blocks previousled created can be seen here."},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_generator():\n    \n\n    \n    inputs = tf.keras.Input(shape = (256,256,3,))\n\n    down_stack = [\n            downsample(64,4, apply_batchnorm= False),\n            downsample(128,4),\n            downsample(256,4),\n            downsample(512,4),\n            downsample(512,4),\n            downsample(512,4),\n            downsample(512,4),\n            downsample(512,4)\n        ]\n\n    up_stack = [\n            upsample(512, 4, apply_dropout= True),\n            upsample(512, 4, apply_dropout= True),\n            upsample(512, 4, apply_dropout= True),\n            upsample(512, 4),\n            upsample(512, 4),\n            upsample(256, 4),\n            upsample(128, 4),\n            upsample(64, 4)\n        ]\n\n\n    initializer = tf.random_normal_initializer()\n\n    x = inputs\n\n        #downsampling stage\n\n    skips = [] #skips used to concatenate with the model in the upsampling stage\n\n    for down in down_stack:\n\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\n        x = up(x)\n        x = tf.keras.layers.concatenate([x, skip])\n\n    x = tf.keras.layers.Conv2DTranspose(3, 4, strides = 2,\n                                           padding = 'same',\n                                           kernel_initializer = initializer, activation = 'tanh')(x)\n\n    return tf.keras.models.Model(inputs = inputs, outputs = x)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"generator_monet = create_generator()\ngenerator_photo = create_generator()\ngenerator_monet.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# picture generated from the random weights of the generator\nrandom_img = generator_monet(next(photo_iter) *127.5 + 127.5)\nrandom_img = random_img[0]\nplt.imshow(random_img)\nplt.axis('off')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The discriminator model is a PatchGan which is basically a regular discriminator that is more speficific tor parts of the photo rather than overall"},{"metadata":{"trusted":true},"cell_type":"code","source":"# building the discriminator model which is a PatchGan\n\ndef create_discriminator():\n    \n    \n\n\n    \n    initializer = tf.random_normal_initializer()\n\n    x = tf.keras.Input(shape = (256, 256, 3,))\n\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\n    conv = tf.keras.layers.Conv2D(512, 4, strides = 1,\n                                     kernel_initializer = initializer,\n                                     use_bias = False)(zero_pad1)\n\n    batchnorm1 = tf.keras.layers.BatchNormalization()(conv)\n\n    leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)\n\n    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)\n\n    outputs = tf.keras.layers.Conv2D(1,4, strides = 1,\n                                        kernel_initializer = initializer)(zero_pad2)\n\n    return tf.keras.Model(inputs = x, outputs = outputs)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"discriminator_monet = create_discriminator()\ndiscriminator_photo = create_discriminator()\ndiscriminator_monet.summary()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"monet_example = discriminator_monet(next(monet_iter))\nphoto_example = discriminator_monet(next(photo_iter))\nmonet_example = monet_example[0]\nphoto_example = photo_example[0]\n\n_ = plt.figure(figsize= (10,10))\nplt.subplot(1,2,1)\nplt.imshow(tf.squeeze(monet_example, 2))\nplt.title(\"Monet image\")\n\nplt.subplot(1,2,2)\nplt.imshow(tf.squeeze(photo_example, 2))\nplt.title(\"Photo image\")\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# defining the losses of the models \n\nloss_object = tf.keras.losses.BinaryCrossentropy(from_logits = True) # the from logits argument is very important since the \n#discriminator model will return a grid of probabilities \n\ndef discriminator_loss(real_image, generated):\n    \n    \"\"\"\n    The loss of the discriminator is simple, just compare two inputs of the same label\n    one is synthetic and the other one is real\n    \"\"\"\n\n    real_loss = loss_object(tf.ones_like(real_image), real_image)\n    generated_loss = loss_object(tf.zeros_like(generated), generated)\n\n    total_loss = real_loss + generated_loss\n\n    return total_loss\n\ndef generator_loss(generated):\n    \n    \"\"\"\n    The generator loss is also simple, basically does better when the discriminator does worse, this is the \n    adversarial component of the model\n    \"\"\"\n\n    gen_loss = loss_object(tf.ones_like(generated), generated)\n\n    return gen_loss\n\ndef consistency_loss(original, returned, LAMBDA):\n    \n    \"\"\"\n    this is also known as the cycle loss\n    \"\"\"\n    \n    loss = tf.reduce_mean(tf.abs(original-returned))\n    \n    return LAMBDA * loss\n\ndef identity_loss(real_image, same_image, LAMBDA):\n    \n    \"\"\"\n    this function helps color preservation on the output of the generators\n    \"\"\"\n    loss = tf.reduce_mean(tf.abs(real_image - same_image))\n    return 0.5 * LAMBDA * loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"One optimizer of each model"},{"metadata":{"trusted":true},"cell_type":"code","source":"m_gen_opt = tf.keras.optimizers.Adam(2e-4)\np_gen_opt = tf.keras.optimizers.Adam(2e-4)\nm_disc_opt = tf.keras.optimizers.Adam(2e-4)\np_disc_opt = tf.keras.optimizers.Adam(2e-4)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_step(photo, monet):\n    \n    with tf.GradientTape() as monet_gen_tape, tf.GradientTape() as monet_disc_tape, tf.GradientTape() as photo_gen_tape, tf.GradientTape() as photo_disc_tape:\n\n            #monet generator\n        fake_monet = generator_monet(photo, training =True)\n        cycled_photo = generator_photo(fake_monet, training = True)\n            \n            #photo generator\n            \n        fake_photo = generator_photo(monet, training = True)\n        cycled_monet = generator_monet(fake_photo, training = True)\n            \n            #monet discriminator \n            \n        fake_monet_disc = discriminator_monet(fake_monet, training = True)\n        real_monet_disc = discriminator_monet(monet, training = True)\n            \n            #photo discriminator\n            \n        fake_photo_disc = discriminator_photo(fake_photo, training = True)\n        real_photo_disc = discriminator_photo(photo, training = True)\n        \n            # for identity loss\n        same_monet = generator_monet(monet, training = True)\n        same_photo = generator_photo(photo, training = True)\n            \n            \n            #generator loss\n        monet_gen_loss = generator_loss(fake_monet_disc)\n        photo_gen_loss = generator_loss(fake_photo_disc)\n            \n            #discrimintor loss\n            \n        monet_disc_loss = discriminator_loss(real_monet_disc, fake_monet_disc)\n        photo_disc_loss = discriminator_loss(real_photo_disc, fake_photo_disc)\n            \n            #consistency/cycle loss\n            \n        monet_cycle_loss = consistency_loss(monet, cycled_monet, 0.25)\n        photo_cycle_loss = consistency_loss(photo, cycled_photo, 0.25)\n        total_cycle_loss = monet_cycle_loss + photo_cycle_loss\n        \n            #identity loss\n            \n        monet_id_loss = identity_loss(monet, same_monet, 0.25)\n        photo_id_loss = identity_loss(photo, same_photo, 0.25)\n            \n            \n            #generator total losses\n            \n        monet_generator_total_loss = monet_gen_loss + total_cycle_loss + monet_id_loss\n        photo_generator_total_loss = photo_gen_loss + total_cycle_loss + photo_id_loss\n            \n\n    monet_generator_gradients = monet_gen_tape.gradient(monet_generator_total_loss, generator_monet.trainable_variables)\n    photo_generator_gradients = photo_gen_tape.gradient(photo_generator_total_loss, generator_photo.trainable_variables)\n        \n    monet_discriminator_gradients = monet_disc_tape.gradient(monet_disc_loss, discriminator_monet.trainable_variables)\n    photo_discriminator_gradients = photo_disc_tape.gradient(photo_disc_loss, discriminator_photo.trainable_variables)\n\n    m_gen_opt.apply_gradients(zip(monet_generator_gradients, generator_monet.trainable_variables))\n    p_gen_opt.apply_gradients(zip(photo_generator_gradients, generator_photo.trainable_variables))\n        \n    m_disc_opt.apply_gradients(zip(monet_discriminator_gradients, discriminator_monet.trainable_variables))\n    p_disc_opt.apply_gradients(zip(photo_discriminator_gradients, discriminator_photo.trainable_variables))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_progress():\n    \n    \"\"\"\n    function that takes 4 examples of each sample dataset to show its progress while training\n    \"\"\"\n\n    photo_iter = iter(photo_ds)\n        \n    _ = plt.figure(figsize = (15,5))\n        \n    for i in range(4):\n        \n        fake_monets = generator_monet(next(photo_iter))[0]\n            \n        plt.subplot(1,4, i+1)\n        plt.imshow(tf.cast(fake_monets * 127.5 + 127.5, tf.uint8))\n        plt.title(\"monet generator\")\n            \n    plt.show()\n\ndef fit(monets, photos, epochs):\n\n    for epoch in range(epochs):\n\n        show_progress()\n\n    # this is the actual training bit\n\n        for mon, pho in tqdm(zip(monets, photos)):\n\n            train_step(pho, mon)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"monet_iter = None\nepochs = 16\n\nfit(monet_ds, photo_ds, epochs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not(os.path.exists('images')):\n    os.makedirs('images') # Create folder to save generated images\n\ndef predict_and_save(input_ds, generator_model):\n    i = 1\n    for img in tqdm(input_ds):\n        prediction = generator_model(img, training=False) # make predition\n        prediction = tf.cast((prediction * 127.5 + 127.5), tf.uint8)   # re-scale\n        im = tf.squeeze(prediction, 0)\n        im = im.numpy()\n        im = Image.fromarray(im)\n        im.save(\"images/\" + str(i) + '.jpg')\n        \n        i += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 1 # this is for the output only, taking an image one by one\nphoto_ds = get_ds(photo_imgs_fnames, aug = False)\n\npredict_and_save(photo_ds, generator_monet)\n    \n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive('/kaggle/working/images/', 'zip', 'images')\nshutil.rmtree(\"./images\")","execution_count":null,"outputs":[]}],"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}