{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Imports\nimport os\nimport numpy as np\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers, Model\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.preprocessing.image import array_to_img\nimport zipfile\n\n# --- TPU / GPU Setup ---\nAUTO = tf.data.AUTOTUNE\n\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\n# Mixed precision for memory savings\nfrom tensorflow.keras import mixed_precision\nmixed_precision.set_global_policy('mixed_float16')\n\n# Paths for TFRecords\nmonet_tfrec = '/kaggle/input/gan-getting-started/monet_tfrec'\nphoto_tfrec = '/kaggle/input/gan-getting-started/photo_tfrec'\n\n# Constants\nIMAGE_SIZE = (256, 256)\nIMAGE_SHAPE = (256, 256, 3)\nBATCH_SIZE = 1\nEPOCHS = 100\nLAMBDA = 10\n\n# --- Custom InstanceNorm (instead of tfa.layers.InstanceNormalization) ---\nclass InstanceNormalization(layers.Layer):\n    def __init__(self, epsilon=1e-5, **kwargs):\n        super(InstanceNormalization, self).__init__(**kwargs)\n        self.epsilon = epsilon\n\n    def build(self, input_shape):\n        self.gamma = self.add_weight(\n            shape=(input_shape[-1],),\n            initializer=\"ones\",\n            trainable=True\n        )\n        self.beta = self.add_weight(\n            shape=(input_shape[-1],),\n            initializer=\"zeros\",\n            trainable=True\n        )\n        super(InstanceNormalization, self).build(input_shape)\n\n    def call(self, x):\n        mean, var = tf.nn.moments(x, axes=[1, 2], keepdims=True)\n        return self.gamma * (x - mean) / tf.sqrt(var + self.epsilon) + self.beta\n\n# Function to parse TFRecords\ndef parse_tfrecord(example_proto):\n    feature_description = {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n    }\n    example = tf.io.parse_single_example(example_proto, feature_description)\n    img = tf.io.decode_jpeg(example['image'], channels=3)\n    img = tf.image.resize(img, IMAGE_SIZE)\n    img = (tf.cast(img, tf.float32) / 127.5) - 1.0  # normalize [-1,1]\n    return img\n\n# Load datasets\nmonet_files = tf.io.gfile.glob(os.path.join(monet_tfrec, \"*.tfrec\"))\nphoto_files = tf.io.gfile.glob(os.path.join(photo_tfrec, \"*.tfrec\"))\n\nmonet_dataset = tf.data.TFRecordDataset(monet_files).cache()\nmonet_dataset = monet_dataset.map(parse_tfrecord, num_parallel_calls=AUTO)\nmonet_dataset = monet_dataset.shuffle(1000).batch(BATCH_SIZE).prefetch(AUTO)\n\nphoto_dataset = tf.data.TFRecordDataset(photo_files).cache()\nphoto_dataset = photo_dataset.map(parse_tfrecord, num_parallel_calls=AUTO)\nphoto_dataset = photo_dataset.shuffle(1000).batch(BATCH_SIZE).prefetch(AUTO)\n\n# --- Define Models ---\ndef downsample(filters, size, apply_instancenorm=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',\n                             kernel_initializer=initializer, use_bias=False))\n    if apply_instancenorm:\n        result.add(InstanceNormalization())\n    result.add(layers.LeakyReLU())\n    return result\n\ndef upsample(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,\n                                      padding='same', kernel_initializer=initializer,\n                                      use_bias=False))\n    result.add(InstanceNormalization())\n    if apply_dropout:\n        result.add(layers.Dropout(0.5))\n    result.add(layers.ReLU())\n    return result\n\ndef Generator():\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=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    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(256, 4),\n        upsample(128, 4),\n        upsample(64, 4),\n    ]\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(IMAGE_SHAPE[2], 4,\n                                   strides=2,\n                                   padding='same',\n                                   kernel_initializer=initializer,\n                                   activation='tanh')\n    concat = layers.Concatenate()\n    inputs = layers.Input(shape=IMAGE_SHAPE)\n    x = inputs\n    skips = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n    skips = reversed(skips[:-1])\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = concat([x, skip])\n    x = last(x)\n    return Model(inputs=inputs, outputs=x, name=\"Generator\")\n\ndef Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    inputs = layers.Input(shape=IMAGE_SHAPE, name='discriminator_input')\n    x = inputs\n    down1 = downsample(64, 4, False)(x)\n    down2 = downsample(128, 4)(down1)\n    down3 = downsample(256, 4)(down2)\n    zero_pad1 = layers.ZeroPadding2D()(down3)\n    conv = layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer,\n                         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    return Model(inputs=inputs, outputs=last, name=\"Discriminator\")\n\n# --- Loss and Optimizers ---\nwith strategy.scope():\n    loss_obj = keras.losses.BinaryCrossentropy(from_logits=True)\n\n    def discriminator_loss(real_img, fake_img):\n        real_loss = loss_obj(tf.ones_like(real_img), real_img)\n        fake_loss = loss_obj(tf.zeros_like(fake_img), fake_img)\n        # Cast to float32 before adding to avoid type mismatch\n        return tf.cast(real_loss, tf.float32) + tf.cast(fake_loss, tf.float32)\n\n    def generator_loss(fake_img):\n        # Cast to float32 to be consistent with other losses\n        return tf.cast(loss_obj(tf.ones_like(fake_img), fake_img), tf.float32)\n\n    def cycle_loss(real_img, cycled_img):\n        # Cast to float32 before adding to avoid type mismatch\n        return LAMBDA * tf.reduce_mean(tf.abs(tf.cast(real_img, tf.float32) - tf.cast(cycled_img, tf.float32)))\n\n    def identity_loss(real_img, same_img):\n        # Cast to float32 before adding to avoid type mismatch\n        return LAMBDA * 0.5 * tf.reduce_mean(tf.abs(tf.cast(real_img, tf.float32) - tf.cast(same_img, tf.float32)))\n\n    monet_generator_optimizer = Adam(2e-4, beta_1=0.5)\n    photo_generator_optimizer = Adam(2e-4, beta_1=0.5)\n    monet_discriminator_optimizer = Adam(2e-4, beta_1=0.5)\n    photo_discriminator_optimizer = Adam(2e-4, beta_1=0.5)\n\n    # Instantiate models inside the strategy scope\n    monet_generator = Generator()\n    photo_generator = Generator()\n    monet_discriminator = Discriminator()\n    photo_discriminator = Discriminator()\n\n\n# --- Training Step ---\n@tf.function\ndef train_step(real_monet, real_photo):\n    # Cast input data to the mixed precision policy's dtype (float16)\n    real_monet = tf.cast(real_monet, tf.float16)\n    real_photo = tf.cast(real_photo, tf.float16)\n    \n    with tf.GradientTape(persistent=True) as tape:\n        # Generate images from both domains\n        fake_monet = monet_generator(real_photo, training=True)\n        cycled_photo = photo_generator(fake_monet, training=True)\n        fake_photo = photo_generator(real_monet, training=True)\n        cycled_monet = monet_generator(fake_photo, training=True)\n\n        # Apply generators to same-domain images for identity loss\n        same_monet = monet_generator(real_monet, training=True)\n        same_photo = photo_generator(real_photo, training=True)\n\n        # Get discriminator outputs for real and fake images\n        disc_fake_monet = monet_discriminator(fake_monet, training=True)\n        disc_real_monet = monet_discriminator(real_monet, training=True)\n        disc_fake_photo = photo_discriminator(fake_photo, training=True)\n        disc_real_photo = photo_discriminator(real_photo, training=True)\n\n        # --- Compute losses ---\n        # Generator losses\n        monet_gen_loss = generator_loss(disc_fake_monet)\n        photo_gen_loss = generator_loss(disc_fake_photo)\n        # Cycle losses\n        total_cycle_loss = cycle_loss(real_monet, cycled_monet) + cycle_loss(real_photo, cycled_photo)\n        # Identity losses\n        id_loss_monet = identity_loss(real_monet, same_monet)\n        id_loss_photo = identity_loss(real_photo, same_photo)\n\n        # Total generator losses (all are now float32)\n        total_monet_gen_loss = monet_gen_loss + total_cycle_loss + id_loss_monet\n        total_photo_gen_loss = photo_gen_loss + total_cycle_loss + id_loss_photo\n\n        # Discriminator losses (all are now float32)\n        monet_disc_loss = discriminator_loss(disc_real_monet, disc_fake_monet)\n        photo_disc_loss = discriminator_loss(disc_real_photo, disc_fake_photo)\n\n    # --- Compute and apply gradients ---\n    monet_gen_grads = tape.gradient(total_monet_gen_loss, monet_generator.trainable_variables)\n    photo_gen_grads = tape.gradient(total_photo_gen_loss, photo_generator.trainable_variables)\n    monet_disc_grads = tape.gradient(monet_disc_loss, monet_discriminator.trainable_variables)\n    photo_disc_grads = tape.gradient(photo_disc_loss, photo_discriminator.trainable_variables)\n\n    monet_generator_optimizer.apply_gradients(zip(monet_gen_grads, monet_generator.trainable_variables))\n    photo_generator_optimizer.apply_gradients(zip(photo_gen_grads, photo_generator.trainable_variables))\n    monet_discriminator_optimizer.apply_gradients(zip(monet_disc_grads, monet_discriminator.trainable_variables))\n    photo_discriminator_optimizer.apply_gradients(zip(photo_disc_grads, photo_discriminator.trainable_variables))\n\n    return total_monet_gen_loss, total_photo_gen_loss, monet_disc_loss, photo_disc_loss\n\n# --- Training loop ---\nwith strategy.scope():\n    monet_iterator = iter(monet_dataset)\n    photo_iterator = iter(photo_dataset)\n\n    for epoch in range(EPOCHS):\n        num_monet_samples = len(tf.io.gfile.glob(os.path.join(monet_tfrec, \"*.tfrec\")))\n        num_photo_samples = len(tf.io.gfile.glob(os.path.join(photo_tfrec, \"*.tfrec\")))\n        num_steps = min(num_monet_samples, num_photo_samples) // BATCH_SIZE\n        \n        for step in range(num_steps):\n            try:\n                monet_batch = next(monet_iterator)\n                photo_batch = next(photo_iterator)\n            except StopIteration:\n                monet_iterator = iter(monet_dataset)\n                photo_iterator = iter(photo_dataset)\n                monet_batch = next(monet_iterator)\n                photo_batch = next(photo_iterator)\n\n            monet_gen_loss, photo_gen_loss, monet_disc_loss, photo_disc_loss = train_step(monet_batch, photo_batch)\n\n        print(f\"Epoch {epoch}/{EPOCHS}, Monet Gen Loss: {monet_gen_loss:.4f}, Photo Gen Loss: {photo_gen_loss:.4f}, Monet Disc Loss: {monet_disc_loss:.4f}, Photo Disc Loss: {photo_disc_loss:.4f}\")\n\n# --- Generate and save images ---\noutput_zip = '/kaggle/working/images.zip'\nwith zipfile.ZipFile(output_zip, 'w') as zipf:\n    img_count = 0\n    \n    for batch in photo_dataset.repeat():\n        if img_count >= 7000:\n            break\n            \n        gen_imgs = monet_generator.predict(batch)\n        for i in range(gen_imgs.shape[0]):\n            img = (0.5 * gen_imgs[i] + 0.5) * 255.0\n            img = array_to_img(img)\n            img_path = f'monet_generated_{img_count}.jpg'\n            img.save(img_path)\n            zipf.write(img_path)\n            os.remove(img_path)\n            img_count += 1\n\nprint(f\"Generated {img_count} images into images.zip\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-09-24T21:59:19.064336Z","iopub.execute_input":"2025-09-24T21:59:19.065078Z"}},"outputs":[],"execution_count":null}]}