{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nfrom matplotlib import pyplot as plt\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\nfrom tensorflow.keras import layers, Sequential, Model\nfrom IPython.display import clear_output\nfrom tqdm import tqdm\nimport time\nAUTOTUNE = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# HYPER PARAMETERS\n\nWIDTH, HEIGHT = 256, 256\nEPOCHS = 20\nBUFFER_SIZE = 400\nBATCH_SIZE = 1\nOUTPUT_CHANNEL = 3\n\n# LOSSES\nLAYER_COUNT = 0\nLAMBDA = 100","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH_X = \"/kaggle/input/gan-getting-started/photo_jpg\"\nPATH_Y = \"/kaggle/input/gan-getting-started/monet_jpg\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_preprocess(img_path):\n\n    img = tf.io.read_file(img_path)\n    img = tf.image.decode_jpeg(img)    \n    img = tf.image.resize(img, [HEIGHT, WIDTH], method = tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n    # normalising\n    img = tf.cast(img, tf.float32)\n    img = img/127.5 -1\n    \n    # random flipping    \n    img = tf.image.random_flip_left_right(img)\n\n    return img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_x = tf.data.Dataset.list_files(PATH_X + '/*.jpg').map(train_preprocess).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\ntrain_y = tf.data.Dataset.list_files(PATH_Y + '/*.jpg').map(train_preprocess).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_x = next(iter(train_x))[0]\nsample_x = tf.reshape(sample_x, [-1, 256, 256, 3])\nsample_y = next(iter(train_y))[0]\nsample_y = tf.reshape(sample_y, [-1, 256, 256, 3])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.subplot(121)\nplt.imshow((sample_x[0] + 1.)/2.)\nplt.subplot(122)\nplt.imshow((sample_y[0] + 1.)/2.)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def add_layer(filters, kernel_size,batchnorm = True):\n    init = tf.keras.initializers.random_normal(0., 0.02)\n    blocks = Sequential()\n    blocks.add(layers.Conv2D(filters, kernel_size=kernel_size, strides = 2, kernel_initializer=init,padding='same', use_bias=False))\n    if batchnorm:\n        blocks.add(layers.BatchNormalization())\n    blocks.add(layers.LeakyReLU())\n    return blocks\n\ndef add_trans_layer(filters, kernel_size, dropout=True):\n    init = tf.keras.initializers.random_normal(0., 0.02)\n    blocks = Sequential()\n    blocks.add(layers.Conv2DTranspose(filters, kernel_size, strides=2, kernel_initializer=init, use_bias=False, padding='same'))\n    blocks.add(layers.BatchNormalization())\n    if dropout:\n        blocks.add(layers.Dropout(0.4))\n    blocks.add(layers.LeakyReLU())\n\n    return blocks","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef make_gen():\n    inputs = layers.Input(shape = [256, 256, 3])\n\n    down = [\n        add_layer(64, 5, False),\n        add_layer(128, 5),\n        add_layer(256, 5),\n        add_layer(512, 5),\n        add_layer(512, 5),\n        add_layer(512, 5),\n        add_layer(512, 5),\n        add_layer(512, 5)\n    ]\n\n    up = [\n          add_trans_layer(512, 5),\n          add_trans_layer(512, 5),\n          add_trans_layer(512, 5),\n          add_trans_layer(512, 5, dropout=False),\n          add_trans_layer(256, 5, dropout=False),\n          add_trans_layer(128, 5, dropout=False),\n          add_trans_layer(64, 5, dropout=False)\n    ]\n\n    init = tf.random_normal_initializer(0., 0.02)\n\n    output = layers.Conv2DTranspose(OUTPUT_CHANNEL, kernel_size=5, strides=(2,2), padding='same', kernel_initializer=init, activation='tanh')\n\n    x = inputs\n\n    stack = []\n\n    for d in down:\n        x = d(x)\n        stack.append(x)\n  \n    stack.pop(-1)\n\n    for a in up:\n        x = a(x)\n        x = layers.Concatenate()([x, stack[-1]])\n        stack.pop(-1)\n\n    x = output(x)\n  \n    return Model(inputs = inputs, outputs=x)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Gx = make_gen()\nGy = make_gen()\ntf.keras.utils.plot_model(Gx, show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def gen_loss(fake_output):\n    return cross_entropy(tf.ones_like(fake_output), fake_output)\n\nGx_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\nGy_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_disc():\n    init = tf.random_normal_initializer(0., 0.05)\n\n    input_shape = [256, 256, 3]\n\n    down = [\n      add_layer(64, 5, batchnorm=False),\n      add_layer(128, 5),\n      add_layer(256, 5),\n      layers.ZeroPadding2D(),\n      layers.Conv2D(512, kernel_size=5, strides=1, kernel_initializer=init, use_bias=False),\n      layers.BatchNormalization(),\n      layers.LeakyReLU(),\n      layers.ZeroPadding2D(),\n      layers.Conv2D(1, kernel_size=5, strides=1, kernel_initializer=init)\n    ]\n\n    inp = layers.Input(shape=input_shape)\n\n    x = inp\n    for layer in down:\n        x = layer(x)\n    \n    return Model(inputs=inp, outputs = x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Dx = make_disc()\nDy = make_disc()\n\ntf.keras.utils.plot_model(Dx, show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def disc_loss(real_output, fake_output):\n    real_loss = cross_entropy(tf.ones_like(real_output), real_output)\n    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)\n    return fake_loss + real_loss\n\nDx_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\nDy_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def cyclic_loss(real_img, cycled_img):\n    return LAMBDA*tf.reduce_mean(abs(real_img-cycled_img))\n\ndef identity_loss(real_img, same_img):\n    return 0.5*LAMBDA*tf.reduce_mean(abs(real_img-same_img))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CycleGAN(Model):\n    \n    def __init__(self):\n        super(CycleGAN, self).__init__()\n        self.Gx = make_gen() # converts X -> Y => fake_y = Gx(input_x)\n        self.Gy = make_gen() # converts Y -> X\n        self.Dx = make_disc() # checks wheather the input is X or not\n        self.Dy = make_disc() # checks wheather the input is Y or not\n    \n    def compile(self, gen_loss, disc_loss, cyclic_loss, identity_loss):\n        super(CycleGAN, self).compile()\n        self.Gx_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\n        self.Gy_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\n        self.Dx_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\n        self.Dy_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\n        self.gen_loss = gen_loss\n        self.disc_loss = disc_loss\n        self.cyclic_loss = cyclic_loss\n        self.identity_loss = identity_loss\n    \n    \n    def train_step(self, batch):\n        real_x, real_y = batch\n        \n        with tf.GradientTape(persistent=True) as tape:\n            fake_y = self.Gx(real_x, training=True)\n            fake_x = self.Gy(real_y, training=True)\n            cycled_x = self.Gy(fake_y, training=True)\n            cycled_y = self.Gx(fake_x, training=True)\n            same_x = self.Gy(real_x, training=True)\n            same_y = self.Gx(real_y, training=True)\n            \n            Dx_real = self.Dx(real_x, training=True)\n            Dx_fake = self.Dx(fake_x, training=True)\n            \n            Dy_real = self.Dy(real_y, training=True)\n            Dy_fake = self.Dy(fake_y, training=True)\n            \n            cycled_loss = self.cyclic_loss(real_x, cycled_x) + self.cyclic_loss(real_y, cycled_y)\n            \n            Gx_loss = self.gen_loss(Dy_fake) + cycled_loss + self.identity_loss(real_x, same_x)\n            Gy_loss = self.gen_loss(Dx_fake) + cycled_loss + self.identity_loss(real_y, same_y)\n            \n            Dx_loss = self.disc_loss(Dx_real, Dx_fake)\n            Dy_loss = self.disc_loss(Dy_real, Dy_fake)\n        \n        Gx_grad = tape.gradient(Gx_loss, self.Gx.trainable_variables)\n        Gy_grad = tape.gradient(Gy_loss, self.Gy.trainable_variables)\n        Dx_grad = tape.gradient(Dx_loss, self.Dx.trainable_variables)\n        Dy_grad = tape.gradient(Dy_loss, self.Dy.trainable_variables)\n        \n        self.Gx_opt.apply_gradients(zip(Gx_grad, self.Gx.trainable_variables))\n        self.Gy_opt.apply_gradients(zip(Gy_grad, self.Gy.trainable_variables))\n        self.Dx_opt.apply_gradients(zip(Dx_grad, self.Dx.trainable_variables))\n        self.Dy_opt.apply_gradients(zip(Dy_grad, self.Dy.trainable_variables))\n        \n        return {\n            \"Gx_loss\" : Gx_loss,\n            \"Gy_loss\" : Gy_loss,\n            \"Dx_loss\" : Dx_loss,\n            \"Dy_loss\" : Dy_loss,\n            \"cycled_loss\" : cycled_loss\n        }\n    \n    \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cyclegan = CycleGAN()\ncyclegan.compile(gen_loss=gen_loss, disc_loss=disc_loss, cyclic_loss=cyclic_loss, identity_loss=identity_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"history = cyclegan.fit(tf.data.Dataset.zip((train_x, train_y)), epochs=EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.subplot(121)\nplt.imshow(sample_x[0])\nplt.subplot(122)\nplt.imshow(cyclegan.Gx(sample_x)[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!mkdir -p saved_model\ncyclegan.save(\"saved_model/cyclegan-model\")","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}