{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport tensorflow_addons as tfa\nimport numpy as np\n\n\nstrategy = tf.distribute.get_strategy()\nAUTOTUNE = tf.data.experimental.AUTOTUNE\n\nINPUT_DIR = \"../input/gan-getting-started\"\n# Load data\nMONET_FILENAMES = tf.io.gfile.glob(INPUT_DIR + '/monet_tfrec/*.tfrec')\nPHOTO_FILENAMES = tf.io.gfile.glob(INPUT_DIR + '/photo_tfrec/*.tfrec')\n\n# Data augment\ndef data_augment(image):\n    p_spatial = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n\n    # rotation\n    if p_rotate > .8:\n        image = tf.image.rot90(image)\n    elif p_rotate > .6:\n        image = tf.image.rot90(image, 2)\n    elif p_rotate > .4:\n        image = tf.image.rot90(image, 3)\n\n    # flip\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_flip_up_down(image)\n    if p_spatial > .75:\n        image = tf.image.transpose(image)\n\n    # crop\n    # image = tf.image.random_crop(image, size=[128, 128, 3])\n\n    return image\n\n# Normalize\ndef normalize_img(image):\n    img = (tf.cast(image, dtype=tf.float32) / 127.5) - 1.\n    return img\n\n# Decode img\ndef decode_img(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    img = tf.reshape(image, [256, 256, 3])\n    return img\n\n# Record data\ndef read_tfrecord(example):\n    tfrecord_format = {\n        'image': tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_img(example['image'])\n    return image\n\n# Load dataset\ndef load_dataset(filename):\n    dataset = tf.data.TFRecordDataset(filename)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    return dataset\n\n# Get dataset\ndef get_dataset(filenames, augment=None, batch_size=1):\n    dataset = load_dataset(filenames)\n\n    dataset = dataset.repeat(2)\n\n    if augment:\n        dataset = dataset.map(augment, num_parallel_calls=AUTOTUNE)\n    dataset = dataset.map(normalize_img, num_parallel_calls=AUTOTUNE)\n\n    dataset = dataset.batch(batch_size)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTOTUNE)\n\n    return dataset\n\nmonet_ds = get_dataset(MONET_FILENAMES, augment=data_augment)\nphoto_ds = get_dataset(PHOTO_FILENAMES)\n\n# Model functions\ndef encoder_block(filters, size=3, strides=2, 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    block = keras.Sequential()\n    # Convolution\n    block.add(layers.Conv2D(filters, size, strides, padding='same', kernel_initializer=initializer, use_bias=False))\n\n    # Normalization\n    if apply_instancenorm:\n        block.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    # Activation\n    block.add(layers.ReLU())\n\n    return block\n\n\ndef decoder_block(filters, size=3, strides=2, 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    block = keras.Sequential()\n    # Transposed convolutional layer\n    block.add(\n        layers.Conv2DTranspose(filters, size, strides, padding='same', kernel_initializer=initializer, use_bias=False))\n\n    # Normalization\n    if apply_instancenorm:\n        block.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    # Activation\n    block.add(layers.LeakyReLU())\n    return block\n\n\ndef resnet_block(input_layer, size=3, strides=1):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    filters = input_layer.shape[-1]\n\n    block = layers.Conv2D(filters, size, strides=strides, padding='same', use_bias=False,\n                          kernel_initializer=initializer)(input_layer)\n    block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(block)\n    block = layers.ReLU()(block)\n\n    block = layers.Conv2D(filters, size, strides=strides, padding='same', use_bias=False,\n                          kernel_initializer=initializer)(block)\n    block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(block)\n\n    block = layers.Add()([block, input_layer])\n\n    return block\n\n# Generator\nOUTPUT_CHANNELS = 3\n\n\ndef Generator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    inputs = layers.Input(shape=[256, 256, 3])\n    x = inputs\n\n    down_layers = [encoder_block(64, 7, 1, apply_instancenorm=False),\n                   encoder_block(128, 3, 2, apply_instancenorm=True),\n                   encoder_block(256, 3, 2, apply_instancenorm=True)]\n\n    skips = []\n    for layer in down_layers:\n        x = layer(x)\n        skips.append(x)\n\n    skips = reversed(skips[:-1])\n\n    for n in range(5):\n        x = resnet_block(x, 3, 1)\n\n    up_layers = [decoder_block(256, 3, 2),\n                 decoder_block(128, 3, 2),\n                 decoder_block(64, 3, 2), ]\n\n    for layer, skip in zip(up_layers, skips):\n        x = layer(x)\n        x = layers.Concatenate()([x, skip])\n\n    last = layers.Conv2D(OUTPUT_CHANNELS, 7,\n                         strides=1, padding='same',\n                         kernel_initializer=initializer,\n                         use_bias=False,\n                         activation='tanh')\n\n    outputs = last(x)\n\n    return keras.Model(inputs=inputs, outputs=outputs)\n\n# discriminator\ndef Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    inputs = layers.Input(shape=[256, 256, 3])\n    x = inputs\n    x = encoder_block(64, 4, 2, apply_instancenorm=False)(x)  # (bs, 128, 128, 64)\n    x = encoder_block(128, 4, 2, apply_instancenorm=True)(x)  # (bs, 64, 64, 128)\n    x = encoder_block(256, 4, 2, apply_instancenorm=True)(x)  # (bs, 32, 32, 256)\n    x = encoder_block(512, 4, 1, apply_instancenorm=True)(x)\n\n    outputs = layers.Conv2D(1, 4, strides=1, padding='valid', kernel_initializer=initializer)(x)\n\n    return keras.Model(inputs=inputs, outputs=outputs)\n\n# Build model\nwith strategy.scope():\n    monet_generator = Generator()  # transforms photos to Monet-esque paintings\n    photo_generator = Generator()  # transforms Monet paintings to be more like photos\n\n    monet_discriminator = Discriminator()  # differentiates real Monet paintings and generated Monet paintings\n    photo_discriminator = Discriminator()  # differentiates real photos and generated photos\n\n\nclass 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            # photo to monet back to photo\n            fake_monet = self.m_gen(real_photo, training=True)\n            cycled_photo = self.p_gen(fake_monet, training=True)\n\n            # monet to photo back to monet\n            fake_photo = self.p_gen(real_monet, training=True)\n            cycled_monet = self.m_gen(fake_photo, training=True)\n\n            # generating itself\n            same_monet = self.m_gen(real_monet, training=True)\n            same_photo = self.p_gen(real_photo, training=True)\n\n            # discriminator used to check, inputing real images\n            disc_real_monet = self.m_disc(real_monet, training=True)\n            disc_real_photo = self.p_disc(real_photo, training=True)\n\n            # discriminator used to check, inputing fake images\n            disc_fake_monet = self.m_disc(fake_monet, training=True)\n            disc_fake_photo = self.p_disc(fake_photo, training=True)\n\n            # evaluates generator loss\n            monet_gen_loss = self.gen_loss_fn(disc_fake_monet)\n            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)\n\n            # evaluates total cycle consistency loss\n            total_cycle_loss = self.cycle_loss_fn(real_monet, cycled_monet, self.lambda_cycle) + self.cycle_loss_fn(\n                real_photo, cycled_photo, self.lambda_cycle)\n\n            # evaluates total generator loss\n            total_monet_gen_loss = monet_gen_loss + total_cycle_loss + self.identity_loss_fn(real_monet, same_monet,\n                                                                                             self.lambda_cycle)\n            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo,\n                                                                                             self.lambda_cycle)\n\n            # evaluates discriminator loss\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        # Calculate the gradients for generator and discriminator\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        # Apply the gradients to the optimizer\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        }\n\n# Loss Function\nwith 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        return 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        loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n        return LAMBDA * loss1\n\n    def identity_loss(real_image, same_image, LAMBDA):\n        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n        return LAMBDA * 0.5 * loss\n\n# #optimizer\nwith strategy.scope():\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#compile\nwith strategy.scope():\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    )\n\n#train\ncycle_gan_model.fit(tf.data.Dataset.zip((monet_ds, photo_ds)), epochs=10)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"os.makedirs('/kaggle/working/images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\ni = 1\nfor img in photo_ds:\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(os.getcwd() + '/images/' + str(i) + \".jpg\")\n    i += 1\n    if i==7038:\n        break\nprint('images saved')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive('/kaggle/working/images/', 'zip', '/kaggle/working/images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"Generated samples: {len([name for name in os.listdir('/kaggle/working/images/') if os.path.isfile(os.path.join('/kaggle/working/images/', name))])}\")","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}