{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.16","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30920,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nfrom tensorflow.keras.layers import GroupNormalization\nimport tensorflow_probability as tfp\nimport matplotlib.pyplot as plt\nfrom tensorflow.keras.preprocessing.image import array_to_img\nfrom tensorflow.keras.callbacks import Callback\nimport numpy as np\nimport re\nimport os\nimport random\nfrom tensorflow.keras.applications import InceptionV3\nimport PIL\nfrom PIL import Image\nimport zipfile\nimport io\nimport time\nfrom tqdm.notebook import tqdm\ntry:\n    from kaggle_datasets import KaggleDatasets\nexcept:\n    pass","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.304578Z","iopub.execute_input":"2025-03-10T21:24:41.304896Z","iopub.status.idle":"2025-03-10T21:24:41.309710Z","shell.execute_reply.started":"2025-03-10T21:24:41.304870Z","shell.execute_reply":"2025-03-10T21:24:41.308909Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# CONSTANTS\nBATCH_SIZE = 32\nIMAGE_SIZE = [256, 256]\nEPOCHS = 50\nSTEPS_PER_EPOCH = 100\nFID_INTERVAL = 5  ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.310458Z","iopub.execute_input":"2025-03-10T21:24:41.310658Z","iopub.status.idle":"2025-03-10T21:24:41.323222Z","shell.execute_reply.started":"2025-03-10T21:24:41.310638Z","shell.execute_reply":"2025-03-10T21:24:41.322200Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create directories\n\n# Create all required directories\ndirs_to_create = [\n    'cache',               # For dataset caching\n    'cache/monet',         # Specific cache directories \n    'cache/photo',         # Specific cache directories\n    'cyclegan_checkpoints',# For model checkpoints\n    'generated_samples',   # For saving generated images\n    'logs'                 # For tensorboard logs if needed\n]\n\n# Create each directory\nfor directory in dirs_to_create:\n    os.makedirs(directory, exist_ok=True)\n    print(f\"Directory {directory} is ready\")\n\n# Function to safely use cache paths\ndef get_cache_path(name):\n    \"\"\"Returns a valid cache path for the given name, ensuring the directory exists.\"\"\"\n    path = os.path.join('cache', name)\n    os.makedirs(path, exist_ok=True)\n    return path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.324129Z","iopub.execute_input":"2025-03-10T21:24:41.324355Z","iopub.status.idle":"2025-03-10T21:24:41.338124Z","shell.execute_reply.started":"2025-03-10T21:24:41.324334Z","shell.execute_reply":"2025-03-10T21:24:41.337238Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Hardware Accelaration\ntry:\n    # 1. Detect TPU hardware\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()  # Auto-detects TPU\n    \n    # 2. Initialize TPU system\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    \n    # 3. Create distributed training strategy\n    strategy = tf.distribute.TPUStrategy(tpu)\n    \n    print(f\"TPU detected: {tpu.cluster_spec().as_dict()['worker']}\")\n    print(f\"Number of TPU cores: {strategy.num_replicas_in_sync}\")\n\nexcept ValueError:\n    # Fallback to GPU/CPU\n    strategy = tf.distribute.get_strategy()\n    print(f\"Using {strategy.__class__.__name__} (CPU/GPU)\")\n\nprint(f\"Number of replicas: {strategy.num_replicas_in_sync}\")\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.339074Z","iopub.execute_input":"2025-03-10T21:24:41.339339Z","iopub.status.idle":"2025-03-10T21:24:41.347715Z","shell.execute_reply.started":"2025-03-10T21:24:41.339316Z","shell.execute_reply":"2025-03-10T21:24:41.346890Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set Batch size if we have a TPU \n# (NOTE BATCH_SIZE must be divisble by 8 for 8 core TPU)\nBATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.348243Z","iopub.execute_input":"2025-03-10T21:24:41.348441Z","iopub.status.idle":"2025-03-10T21:24:41.357984Z","shell.execute_reply.started":"2025-03-10T21:24:41.348421Z","shell.execute_reply":"2025-03-10T21:24:41.357049Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load and prepare data\n# data_path = 'data'\n\n# for kaggle nb\ndata_path = KaggleDatasets().get_gcs_path(\"gan-getting-started\")\n\n\nMONET_JPG= tf.io.gfile.glob(str(data_path + '/monet_jpg/*.jpg'))\nPHOTO_JPG = tf.io.gfile.glob(str(data_path + '/photo_jpg/*.jpg'))\nMONET_TFREC = tf.io.gfile.glob(str(data_path + '/monet_tfrec/*.tfrec'))\nPHOTO_TFREC = tf.io.gfile.glob(str(data_path + '/photo_tfrec/*.tfrec'))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.358772Z","iopub.execute_input":"2025-03-10T21:24:41.358997Z","iopub.status.idle":"2025-03-10T21:24:41.845520Z","shell.execute_reply.started":"2025-03-10T21:24:41.358975Z","shell.execute_reply":"2025-03-10T21:24:41.844224Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def decode_image(image):\n    \"\"\"\n    Decode JPEG: Converts raw bytes to a uint8 tensor.\n\n    Normalize: Scales pixel values from [0, 255] to [-1, 1] (standard for GANs).\n\n    Reshape: Forces images to 256x256x3 (CycleGAN’s default input size).\n    \"\"\"\n    image = tf.image.decode_jpeg(image, channels=3)      # Decode JPEG bytes\n    image = (tf.cast(image, tf.float32) / 127.5) - 1     # Normalize to [-1, 1]\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])         # Resize to 256x256\n    return image\n\ndef read_tfrecord(example):\n    \"\"\"\n    Parses a single TFRecord example to extract the image field (stored as bytes).\n\n    Passes the bytes to decode_image for preprocessing.\n    \"\"\"\n    tfrecord_format = {\"image\": tf.io.FixedLenFeature([], tf.string)}\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    return decode_image(example['image'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.846458Z","iopub.execute_input":"2025-03-10T21:24:41.846701Z","iopub.status.idle":"2025-03-10T21:24:41.852850Z","shell.execute_reply.started":"2025-03-10T21:24:41.846676Z","shell.execute_reply":"2025-03-10T21:24:41.851624Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_dataset(filenames, is_tfrec=True, image_size=IMAGE_SIZE):\n    \"\"\"\n    Load images from either TFRecords or directory of JPG/PNG files.\n    \n    Args:\n        filenames: List of file paths (TFRecords) or directory path (for images)\n        is_tfrec: Boolean flag indicating if input is TFRecords\n        image_size: Target size for resizing images\n    \"\"\"\n    if is_tfrec:\n        # Handle TFRecord files\n        dataset = tf.data.TFRecordDataset(filenames)\n        dataset = dataset.map(\n            read_tfrecord, \n            num_parallel_calls=tf.data.AUTOTUNE\n        )\n    else:\n        # Handle JPG/PNG images from directory\n        def process_image(file_path):\n            img = tf.io.read_file(file_path)\n            img = tf.image.decode_jpeg(img, channels=3)\n            img = tf.image.resize(img, image_size)\n            img = (tf.cast(img, tf.float32) / 127.5) - 1  # Normalize to [-1, 1]\n            return img\n        \n        dataset = tf.data.Dataset.list_files(filenames + \"/*.jpg\")\n        dataset = dataset.map(\n            process_image, \n            num_parallel_calls=tf.data.AUTOTUNE\n        )\n    \n    return dataset\n\n# Load datasets\nmonet_ds = load_dataset(MONET_TFREC).batch(1)\nphoto_ds = load_dataset(PHOTO_TFREC).batch(1)\n\n# Batches scaled by strategy.num_replicas_in_sync (for TPU/GPU parallelism). prefetch(32) overlaps data loading with training to avoid bottlenecks\nfast_photo_ds = load_dataset(PHOTO_TFREC).batch(32*strategy.num_replicas_in_sync).prefetch(32)\n\n# Subset (take(1024)) for Fréchet Inception Distance (FID) calculation. Larger batches (32 * replicas) for efficient evaluation.\nfid_photo_ds = load_dataset(PHOTO_TFREC).take(1024).batch(32*strategy.num_replicas_in_sync).prefetch(32)\nfid_monet_ds = load_dataset(MONET_TFREC).batch(32*strategy.num_replicas_in_sync).prefetch(32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:41.853872Z","iopub.execute_input":"2025-03-10T21:24:41.854113Z","iopub.status.idle":"2025-03-10T21:24:46.040069Z","shell.execute_reply.started":"2025-03-10T21:24:41.854088Z","shell.execute_reply":"2025-03-10T21:24:46.039130Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Take one sample from each dataset\nmonet_sample = next(iter(monet_ds.shuffle(10)))\nphoto_sample = next(iter(photo_ds.shuffle(10)))\n\n# Convert from [-1, 1] range to [0, 1] for visualization\nmonet_image = monet_sample[0].numpy() * 0.5 + 0.5  # [0] accesses first image in batch\nphoto_image = photo_sample[0].numpy() * 0.5 + 0.5\n\n# Create subplots\nplt.figure(figsize=(10, 5))\n\n# Plot Monet painting\nplt.subplot(1, 2, 1)\nplt.imshow(monet_image)\nplt.title(\"Monet Style Example\")\nplt.axis('off')\n\n# Plot Photo\nplt.subplot(1, 2, 2)\nplt.imshow(photo_image)\nplt.title(\"Real Photo Example\")\nplt.axis('off')\n\nplt.tight_layout()\nplt.show()\n\nprint(\"Monet image shape:\", monet_image.shape)\nprint(\"Photo image shape:\", photo_image.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:46.040959Z","iopub.execute_input":"2025-03-10T21:24:46.041181Z","iopub.status.idle":"2025-03-10T21:24:46.697997Z","shell.execute_reply.started":"2025-03-10T21:24:46.041157Z","shell.execute_reply":"2025-03-10T21:24:46.697186Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_gan_dataset(monet_files, photo_files, repeat=True, shuffle=True, batch_size=1):\n    # Re-load raw datasets fully\n    monet_ds = load_dataset(monet_files)\n    photo_ds = load_dataset(photo_files)\n\n    # Shuffle\n    if shuffle:\n        monet_ds = monet_ds.shuffle(2048)\n        photo_ds = photo_ds.shuffle(2048)\n\n    # Batch and optimize\n    monet_ds = monet_ds.batch(batch_size, drop_remainder=True)\n    photo_ds = photo_ds.batch(batch_size, drop_remainder=True)\n\n    # Use cache\n    # Use the helper function for cache paths\n    monet_ds = monet_ds.cache(get_cache_path('monet'))\n    photo_ds = photo_ds.cache(get_cache_path('photo'))\n   \n    # Repeat indefinitely for epochs\n    if repeat:\n        monet_ds = monet_ds.repeat()\n        photo_ds = photo_ds.repeat()\n\n    monet_ds = monet_ds.prefetch(AUTOTUNE)\n    photo_ds = photo_ds.prefetch(AUTOTUNE)\n\n    # Pair Monet and Photo batches\n    gan_ds = tf.data.Dataset.zip((monet_ds, photo_ds))\n    return gan_ds","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:46.698829Z","iopub.execute_input":"2025-03-10T21:24:46.699044Z","iopub.status.idle":"2025-03-10T21:24:46.706046Z","shell.execute_reply.started":"2025-03-10T21:24:46.699021Z","shell.execute_reply":"2025-03-10T21:24:46.704709Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"training_dataset = get_gan_dataset(\n    MONET_TFREC, # path to Monet paintings\n    PHOTO_TFREC, # path to real photos (domain B)\n    repeat=True, # loop dataset for multi-epoch training\n    shuffle=True, # shuffle data\n    batch_size=BATCH_SIZE\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:46.706855Z","iopub.execute_input":"2025-03-10T21:24:46.707083Z","iopub.status.idle":"2025-03-10T21:24:46.780703Z","shell.execute_reply.started":"2025-03-10T21:24:46.707058Z","shell.execute_reply":"2025-03-10T21:24:46.779760Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_inception_feature_extractor():\n    \"\"\"\n    Creates a modified InceptionV3 model for FID feature extraction.\n    Returns:\n        tf.keras.Model: Feature extractor using InceptionV3's \"mixed9\" layer outputs.\n    \"\"\"\n    # Load pre-trained InceptionV3 without classification head\n    inception = tf.keras.applications.InceptionV3(\n        include_top=False,  # Remove final classification layer\n        pooling=\"avg\",  # Add global average pooling after last conv layer\n        input_shape=(256, 256, 3)  # Match CycleGAN's image size\n    )\n    \n    # Extract intermediate features from \"mixed9\" layer\n    # Why \"mixed9\"? It captures high-level features before final pooling\n    mix9 = inception.get_layer(\"mixed9\").output  # Shape: (None, 8, 8, 2048)\n    \n    # Additional pooling to reduce spatial dimensions\n    features = layers.GlobalAveragePooling2D()(mix9)  # Shape: (None, 2048)\n    \n    # Build final feature extraction model\n    return tf.keras.Model(inputs=inception.input, outputs=features)\n\ndef calculate_activation_statistics(dataset, fid_model):\n    \"\"\"\n    Computes mean and covariance matrix of feature vectors from a dataset.\n    \n    Args:\n        dataset (tf.data.Dataset): Batched dataset of images (shape: [None, 256, 256, 3])\n        fid_model (tf.keras.Model): Feature extractor model\n    \n    Returns:\n        tuple: (mu, sigma) - Mean vector and covariance matrix\n    \"\"\"\n    # Initialize lists to collect activations\n    all_activations = []\n    \n    # Process dataset in batches\n    for batch in dataset:\n        # Extract features for current batch\n        act = fid_model(batch)\n        all_activations.append(act)\n    \n    # Concatenate all activations\n    act_matrix = tf.concat(all_activations, axis=0)\n    \n    # Calculate mean and covariance\n    mu = tf.reduce_mean(act_matrix, axis=0)\n    sigma = tfp.stats.covariance(act_matrix, sample_axis=0)  # Requires tensorflow-probability\n    \n    return mu, sigma\n\n# -------------------- Execution -------------------- \nwith strategy.scope(): # TPU/GPU integration\n    \n    # 1. Initialize feature extractor\n    inception_model = build_inception_feature_extractor()\n    inception_model.trainable = False\n    \n    # 2. Precompute real image statistics\n    # Ensure fid_monet_ds is properly batched\n    fid_monet_ds = load_dataset(MONET_TFREC).batch(32).prefetch(AUTOTUNE)\n    \n    # Calculate statistics\n    myFID_mu_real, myFID_sigma_real = calculate_activation_statistics(\n        fid_monet_ds,  # receives proper image tensors\n        inception_model\n    )\n    \n    # 3. Initialize FID tracking list\n    fids = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:46.781781Z","iopub.execute_input":"2025-03-10T21:24:46.782162Z","iopub.status.idle":"2025-03-10T21:24:59.426591Z","shell.execute_reply.started":"2025-03-10T21:24:46.782135Z","shell.execute_reply":"2025-03-10T21:24:59.425227Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with strategy.scope():\n    \n    def calculate_frechet_distance(mu1, sigma1, mu2, sigma2):\n        \"\"\"Computes the Fréchet Distance between two multivariate Gaussians.\"\"\"\n        # Squared L2 norm of mean difference\n        diff = mu1 - mu2\n        mean_term = tf.reduce_sum(diff**2)  # ||μ₁ - μ₂||²\n\n        # Covariance term: Tr(Σ₁ + Σ₂ - 2√(Σ₁Σ₂))\n        cov_product = tf.matmul(sigma1, sigma2)\n        cov_product = tf.cast(cov_product, tf.complex64)  # For sqrtm\n        covmean = tf.linalg.sqrtm(cov_product)\n        covmean = tf.math.real(covmean)  # Cast back to float32\n        covmean = tf.cast(covmean, tf.float32)\n\n        # Avoid NaN gradients by ensuring covmean is finite\n        covmean = tf.where(\n            tf.math.is_nan(covmean), \n            tf.zeros_like(covmean), \n            covmean\n        )\n\n        # Compute trace terms\n        tr_covmean = tf.linalg.trace(covmean)\n        trace_term = (\n            tf.linalg.trace(sigma1) + \n            tf.linalg.trace(sigma2) - \n            2 * tr_covmean\n        )\n\n        fid = mean_term + trace_term\n        return fid\n\n    def compute_fid(generator, inception_model, real_mu, real_sigma, dataset):\n    \n        \"\"\"Computes FID between generated and real images.\"\"\"\n        # Define a tf.function to generate images\n        @tf.function\n        def generate_images(images):\n            return generator(images, training=False)\n        \n        # Collect all generated activations\n        all_activations = []\n        for batch in dataset:\n            # Generate images\n            generated_images = generate_images(batch)\n            # Extract features\n            activations = inception_model(generated_images)\n            all_activations.append(activations)\n        \n        # Concatenate activations\n        gen_activations = tf.concat(all_activations, axis=0)\n        \n        # Compute generated statistics\n        gen_mu = tf.reduce_mean(gen_activations, axis=0)\n        gen_sigma = tfp.stats.covariance(gen_activations)\n        \n        # Calculate FID\n        fid_value = calculate_frechet_distance(\n            gen_mu, gen_sigma, \n            real_mu, real_sigma\n        )\n        return fid_value","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.427618Z","iopub.execute_input":"2025-03-10T21:24:59.427866Z","iopub.status.idle":"2025-03-10T21:24:59.436084Z","shell.execute_reply.started":"2025-03-10T21:24:59.427842Z","shell.execute_reply":"2025-03-10T21:24:59.435336Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def down_sample(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(GroupNormalization(groups=-1))\n    result.add(layers.LeakyReLU())\n    return result\n\ndef up_sample(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, padding='same',\n                                      kernel_initializer=initializer, use_bias=False))\n    result.add(GroupNormalization(groups=-1))\n    if apply_dropout:\n        result.add(layers.Dropout(0.5))\n    result.add(layers.ReLU())\n    return result","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.437055Z","iopub.execute_input":"2025-03-10T21:24:59.437402Z","iopub.status.idle":"2025-03-10T21:24:59.451091Z","shell.execute_reply.started":"2025-03-10T21:24:59.437379Z","shell.execute_reply":"2025-03-10T21:24:59.450321Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def Generator():\n    '''\n    1. Downsampling Path (Encoder)\n    Purpose: Compresses the input image into a low-dimensional bottleneck.\n\n    Layers:\n\n    8 down_sample blocks with increasing filters (64 → 512).\n\n    Each block reduces spatial resolution by half (stride=2).\n\n    First block omits instance normalization to preserve low-level details.\n\n    Output: 1x1x512 bottleneck tensor.\n\n    2. Upsampling Path (Decoder)\n    Purpose: Reconstructs the image from the bottleneck to the target domain.\n\n    Layers:\n\n    7 up_sample blocks with decreasing filters (512 → 64).\n\n    Transposed convolutions (stride=2) double spatial resolution.\n\n    Dropout (50%) in early layers to prevent overfitting.\n\n    Skip Connections: Concatenate features from the encoder to the decoder (U-Net structure).\n\n    3. Final Output Layer\n    Conv2DTranspose:\n\n    Output channels: 3 (RGB).\n\n    tanh activation: Normalizes outputs to [-1, 1], matching input normalization.\n    '''\n    inputs = layers.Input(shape=[256, 256, 3])\n    \n    # Downsampling path\n    down_stack = [\n        down_sample(64, 4, apply_instancenorm=False),  # 256x256 → 128x128\n        down_sample(128, 4),                            # 128x128 → 64x64\n        down_sample(256, 4),                            # 64x64 → 32x32\n        down_sample(512, 4),                            # 32x32 → 16x16\n        down_sample(512, 4),                            # 16x16 → 8x8\n        down_sample(512, 4),                            # 8x8 → 4x4\n        down_sample(512, 4),                            # 4x4 → 2x2\n        down_sample(512, 4),                            # 2x2 → 1x1 (bottleneck)\n    ]\n\n    # Upsampling path\n    up_stack = [\n        up_sample(512, 4, apply_dropout=True),  # 1x1 → 2x2\n        up_sample(512, 4, apply_dropout=True),   # 2x2 → 4x4\n        up_sample(512, 4, apply_dropout=True),   # 4x4 → 8x8\n        up_sample(512, 4),                       # 8x8 → 16x16\n        up_sample(256, 4),                       # 16x16 → 32x32\n        up_sample(128, 4),                       # 32x32 → 64x64\n        up_sample(64, 4),                        # 64x64 → 128x128\n    ]\n\n    # Final output layer\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(\n        3, 4, strides=2, padding='same',\n        kernel_initializer=initializer, activation='tanh'\n    )  # 128x128 → 256x256\n\n    x = inputs\n    skips = []\n    \n    # Downsampling\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n    \n    skips = reversed(skips[:-1])  # Omit the bottleneck layer\n    \n    # Upsampling with skip connections\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = layers.Concatenate()([x, skip])\n    \n    x = last(x)\n    return keras.Model(inputs=inputs, outputs=x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.451963Z","iopub.execute_input":"2025-03-10T21:24:59.452166Z","iopub.status.idle":"2025-03-10T21:24:59.462817Z","shell.execute_reply.started":"2025-03-10T21:24:59.452145Z","shell.execute_reply":"2025-03-10T21:24:59.461692Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Input (256x256x3)\n\n│\n\n├─Downsample 64 → 128x128x64\n\n├─Downsample 128 → 64x64x128\n\n├─Downsample 256 → 32x32x256\n\n├─Downsample 512 → 16x16x512\n\n├─Downsample 512 → 8x8x512\n\n├─Downsample 512 → 4x4x512\n\n├─Downsample 512 → 2x2x512\n\n└─Downsample 512 → 1x1x512 (bottleneck)\n\n│\n\n├─Upsample 512 → 2x2x512 (with skip from 2x2x512)\n\n├─Upsample 512 → 4x4x512 (with skip from 4x4x512)\n\n├─Upsample 512 → 8x8x512 (with skip from 8x8x512)\n\n├─Upsample 512 → 16x16x512 (with skip from 16x16x512)\n\n├─Upsample 256 → 32x32x256 (with skip from 32x32x256)\n\n├─Upsample 128 → 64x64x128 (with skip from 64x64x128)\n\n├─Upsample 64 → 128x128x64 (with skip from 128x128x64)\n\n│\n\n└─Output (256x256x3)","metadata":{}},{"cell_type":"code","source":"def Discriminator():\n    '''\n    Discriminator Architecture (Base)\n    Purpose: Feature extractor for real/fake classification\n\n    Structure:\n\n    3 Downsample Blocks: Reduce resolution while increasing filters (64→256)\n\n    Zero Padding: Expands spatial dimensions for subsequent convolutions\n\n    Final Conv Layer: 512 filters with stride 1 (no resolution change)\n\n    Output Shape: 33x33x512 feature maps\n\n    2. DHead (Decision Head)\n    Purpose: Final classification layer for adversarial loss\n\n    Structure:\n\n    Input: 33x33x512 features from base discriminator\n\n    Conv2D(1): Reduces to 30x30x1 \"patch\" predictions\n\n    No Activation: Raw logits for different loss functions\n\n    3. Design Choices\n    Separate Heads: Allows using different loss functions:\n\n    dHead1: Binary cross-entropy (BCE)\n\n    dHead2: Hinge loss\n\n    Instance Normalization: Stabilizes training by normalizing features per-image\n\n    Zero Padding: Preserves spatial dimensions after convolutions\n\n    LeakyReLU: (α=0.2) prevents dead neurons in discriminator\n    '''\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    \n    inp = layers.Input(shape=[256, 256, 3], name='input_image')\n    x = inp\n    \n    # Downsampling Path\n    down1 = down_sample(64, 4, False)(x)       # 256x256 → 128x128\n    down2 = down_sample(128, 4)(down1)         # 128x128 → 64x64\n    down3 = down_sample(256, 4)(down2)         # 64x64 → 32x32\n    \n    # Final Layers\n    zero_pad1 = layers.ZeroPadding2D()(down3)  # 32x32 → 34x34\n    conv = layers.Conv2D(512, 4, strides=1, \n                        kernel_initializer=initializer, \n                        use_bias=False)(zero_pad1)  # 34x34 → 31x31\n    norm1 = GroupNormalization(groups=-1, gamma_initializer=gamma_init)(conv)\n    leaky_relu = layers.LeakyReLU()(norm1)\n    zero_pad2 = layers.ZeroPadding2D()(leaky_relu)  # 31x31 → 33x33\n    \n    return keras.Model(inputs=inp, outputs=zero_pad2)\n\ndef DHead():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    \n    inp = layers.Input(shape=[33, 33, 512], name='input_image')\n    x = inp\n    last = layers.Conv2D(1, 4, strides=1, \n                        kernel_initializer=initializer)(x)  # 33x33 → 30x30\n    return keras.Model(inputs=inp, outputs=last)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.463849Z","iopub.execute_input":"2025-03-10T21:24:59.464093Z","iopub.status.idle":"2025-03-10T21:24:59.477659Z","shell.execute_reply.started":"2025-03-10T21:24:59.464069Z","shell.execute_reply":"2025-03-10T21:24:59.476720Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Input Image (256x256x3)\n\n│\n\n├─Downsample 64 → 128x128x64\n\n├─Downsample 128 → 64x64x128\n\n├─Downsample 256 → 32x32x256\n\n│\n\n├─ZeroPad2D → 34x34x256\n\n├─Conv2D 512 → 31x31x512\n\n├─InstanceNorm + LeakyReLU\n\n└─ZeroPad2D → 33x33x512 (Base Discriminator Output)\n\n│\n\n├─DHead:\n\n  ├─Conv2D 1 → 30x30x1 (Patch Predictions)","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def DiffAugment(x, policy='', channels_first=False):\n        if policy:\n            if channels_first:\n                x = tf.transpose(x, [0, 2, 3, 1])\n            for p in policy.split(','):\n                for f in AUGMENT_FNS[p]:\n                    x = f(x)\n            if channels_first:\n                x = tf.transpose(x, [0, 3, 1, 2])\n        return x\n\n\n    def rand_brightness(x):\n        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) - 0.5\n        x = x + magnitude\n        return x\n\n\n    def rand_saturation(x):\n        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) * 2\n        x_mean = tf.reduce_sum(x, axis=3, keepdims=True) * 0.3333333333333333333\n        x = (x - x_mean) * magnitude + x_mean\n        return x\n\n\n    def rand_contrast(x):\n        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) + 0.5\n        x_mean = tf.reduce_sum(x, axis=[1, 2, 3], keepdims=True) * 5.086e-6\n        x = (x - x_mean) * magnitude + x_mean\n        return x\n\n    def rand_translation(x, ratio=0.125):\n        batch_size = tf.shape(x)[0]\n        image_size = tf.shape(x)[1:3]\n        shift = tf.cast(tf.cast(image_size, tf.float32) * ratio + 0.5, tf.int32)\n        translation_x = tf.random.uniform([batch_size, 1], -shift[0], shift[0] + 1, dtype=tf.int32)\n        translation_y = tf.random.uniform([batch_size, 1], -shift[1], shift[1] + 1, dtype=tf.int32)\n        grid_x = tf.clip_by_value(tf.expand_dims(tf.range(image_size[0], dtype=tf.int32), 0) + translation_x + 1, 0, image_size[0] + 1)\n        grid_y = tf.clip_by_value(tf.expand_dims(tf.range(image_size[1], dtype=tf.int32), 0) + translation_y + 1, 0, image_size[1] + 1)\n        x = tf.gather_nd(tf.pad(x, [[0, 0], [1, 1], [0, 0], [0, 0]]), tf.expand_dims(grid_x, -1), batch_dims=1)\n        x = tf.transpose(tf.gather_nd(tf.pad(tf.transpose(x, [0, 2, 1, 3]), [[0, 0], [1, 1], [0, 0], [0, 0]]), tf.expand_dims(grid_y, -1), batch_dims=1), [0, 2, 1, 3])\n        return x\n\n\n    def rand_cutout(x, ratio=0.5):\n        batch_size = tf.shape(x)[0]\n        image_size = tf.shape(x)[1:3]\n        cutout_size = tf.cast(tf.cast(image_size, tf.float32) * ratio + 0.5, tf.int32)\n        offset_x = tf.random.uniform([tf.shape(x)[0], 1, 1], maxval=image_size[0] + (1 - cutout_size[0] % 2), dtype=tf.int32)\n        offset_y = tf.random.uniform([tf.shape(x)[0], 1, 1], maxval=image_size[1] + (1 - cutout_size[1] % 2), dtype=tf.int32)\n        grid_batch, grid_x, grid_y = tf.meshgrid(tf.range(batch_size, dtype=tf.int32), tf.range(cutout_size[0], dtype=tf.int32), tf.range(cutout_size[1], dtype=tf.int32), indexing='ij')\n        cutout_grid = tf.stack([grid_batch, grid_x + offset_x - cutout_size[0] // 2, grid_y + offset_y - cutout_size[1] // 2], axis=-1)\n        mask_shape = tf.stack([batch_size, image_size[0], image_size[1]])\n        cutout_grid = tf.maximum(cutout_grid, 0)\n        cutout_grid = tf.minimum(cutout_grid, tf.reshape(mask_shape - 1, [1, 1, 1, 3]))\n        mask = tf.maximum(1 - tf.scatter_nd(cutout_grid, tf.ones([batch_size, cutout_size[0], cutout_size[1]], dtype=tf.float32), mask_shape), 0)\n        x = x * tf.expand_dims(mask, axis=3)\n        return x\n\n\n    AUGMENT_FNS = {\n        'color': [rand_brightness, rand_saturation, rand_contrast],\n        'translation': [rand_translation],\n        'cutout': [rand_cutout],\n}\n    def aug_fn(image):\n        return DiffAugment(image,\"color,translation,cutout\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.478559Z","iopub.execute_input":"2025-03-10T21:24:59.478770Z","iopub.status.idle":"2025-03-10T21:24:59.498947Z","shell.execute_reply.started":"2025-03-10T21:24:59.478748Z","shell.execute_reply":"2025-03-10T21:24:59.498322Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CycleGan(keras.Model):\n    def __init__(self, m_gen, p_gen, m_disc, p_disc, dhead1=None, dhead2=None, lambda_cycle=3, lambda_id=3):\n        super().__init__()\n        self.m_gen = m_gen  # Monet generator (photos → paintings)\n        self.p_gen = p_gen  # Photo generator (paintings → photos)\n        self.m_disc = m_disc  # Monet discriminator (base)\n        self.p_disc = p_disc  # Photo discriminator (full)\n        self.dhead1 = dhead1  # First discriminator head\n        self.dhead2 = dhead2  # Second discriminator head (can be None)\n        self.lambda_cycle = lambda_cycle  # Cycle consistency weight\n        self.lambda_id = lambda_id  # Identity loss weight\n\n    def compile(self, m_gen_opt, p_gen_opt, m_disc_opt, p_disc_opt, \n               gen_loss_fn1, gen_loss_fn2, disc_loss_fn1, disc_loss_fn2,\n               cycle_loss_fn, identity_loss_fn, aug_fn):\n        super().compile()\n        # Optimizers\n        self.m_gen_opt = m_gen_opt\n        self.p_gen_opt = p_gen_opt\n        self.m_disc_opt = m_disc_opt\n        self.p_disc_opt = p_disc_opt\n        \n        # Loss Functions\n        self.gen_loss_fn1 = gen_loss_fn1  # e.g., BCE\n        self.gen_loss_fn2 = gen_loss_fn2  # e.g., Hinge\n        self.disc_loss_fn1 = disc_loss_fn1\n        self.disc_loss_fn2 = disc_loss_fn2\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n        \n        # Augmentation\n        self.aug_fn = aug_fn  # DiffAugment policy\n\n    def augment_batch(self, real_images, fake_images):\n        \"\"\"Apply data augmentation to both real and fake images.\"\"\"\n        # Concatenate images for batched augmentation (more efficient)\n        combined = tf.concat([real_images, fake_images], axis=0)\n        \n        # Apply augmentation\n        augmented = self.aug_fn(combined)\n        \n        # Split back into real and fake\n        batch_size = tf.shape(real_images)[0]\n        aug_real = augmented[:batch_size]\n        aug_fake = augmented[batch_size:]\n        \n        return aug_real, aug_fake\n\n    def train_step(self, batch_data):\n        real_monet, real_photo = batch_data\n        batch_size = tf.shape(real_monet)[0]\n        \n        with tf.GradientTape(persistent=True) as tape:\n            # Forward cycle: photo → monet → photo\n            fake_monet = self.m_gen(real_photo, training=True)\n            cycled_photo = self.p_gen(fake_monet, training=True)\n            \n            # Backward cycle: monet → photo → monet\n            fake_photo = self.p_gen(real_monet, training=True)\n            cycled_monet = self.m_gen(fake_photo, training=True)\n            \n            # Identity mapping\n            same_monet = self.m_gen(real_monet, training=True)\n            same_photo = self.p_gen(real_photo, training=True)\n            \n            # Apply augmentation if provided\n            aug_real_monet, aug_fake_monet = self.augment_batch(real_monet, fake_monet)\n            aug_real_photo, aug_fake_photo = self.augment_batch(real_photo, fake_photo)\n            \n            # Monet discriminator\n            disc_real_monet_features = self.m_disc(aug_real_monet, training=True)\n            disc_fake_monet_features = self.m_disc(aug_fake_monet, training=True)\n            \n            # Initialize loss values\n            monet_gen_loss = 0\n            monet_disc_loss = 0\n            monet_gen_loss2 = 0\n            monet_disc_loss2 = 0\n            \n            # Use discriminator head if available\n            if self.dhead1 is not None:\n                disc_real_monet = self.dhead1(disc_real_monet_features, training=True)\n                disc_fake_monet = self.dhead1(disc_fake_monet_features, training=True)\n                monet_gen_loss = self.gen_loss_fn1(disc_fake_monet)\n                monet_disc_loss = self.disc_loss_fn1(disc_real_monet, disc_fake_monet)\n                \n                # Second head (optional)\n                if self.dhead2 is not None:\n                    disc_real_monet2 = self.dhead2(disc_real_monet_features, training=True)\n                    disc_fake_monet2 = self.dhead2(disc_fake_monet_features, training=True)\n                    monet_gen_loss2 = self.gen_loss_fn2(disc_fake_monet2)\n                    monet_disc_loss2 = self.disc_loss_fn2(disc_real_monet2, disc_fake_monet2)\n            else:\n                # Use features directly (patch discriminator)\n                disc_real_monet = disc_real_monet_features\n                disc_fake_monet = disc_fake_monet_features\n                monet_gen_loss = self.gen_loss_fn1(disc_real_monet)\n                monet_disc_loss = self.disc_loss_fn1(disc_real_monet, disc_fake_monet)\n            \n            # Photo discriminator\n            disc_real_photo = self.p_disc(aug_real_photo, training=True)\n            disc_fake_photo = self.p_disc(aug_fake_photo, training=True)\n            photo_gen_loss = self.gen_loss_fn1(disc_fake_photo)\n            photo_disc_loss = self.disc_loss_fn1(disc_real_photo, disc_fake_photo)\n            \n            # Cycle consistency loss\n            cycle_loss = (\n                self.cycle_loss_fn(real_monet, cycled_monet) + \n                self.cycle_loss_fn(real_photo, cycled_photo)\n            ) * self.lambda_cycle\n            \n            # Identity loss\n            identity_loss = (\n                self.identity_loss_fn(real_monet, same_monet) + \n                self.identity_loss_fn(real_photo, same_photo)\n            ) * self.lambda_id\n            \n            # Total losses\n            total_monet_gen_loss = monet_gen_loss + monet_gen_loss2 + cycle_loss + identity_loss\n            total_photo_gen_loss = photo_gen_loss + cycle_loss + identity_loss\n            total_monet_disc_loss = monet_disc_loss + monet_disc_loss2\n        \n        # Calculate gradients\n        monet_gen_grads = tape.gradient(total_monet_gen_loss, self.m_gen.trainable_variables)\n        photo_gen_grads = tape.gradient(total_photo_gen_loss, self.p_gen.trainable_variables)\n        \n        # Apply generator gradients directly\n        self.m_gen_opt.apply_gradients(zip(monet_gen_grads, self.m_gen.trainable_variables))\n        self.p_gen_opt.apply_gradients(zip(photo_gen_grads, self.p_gen.trainable_variables))\n        \n        # Handle discriminator (complete networks at once instead of separating heads)\n        # This avoids unnecessary computation in graph mode\n        all_m_disc_vars = self.m_disc.trainable_variables\n        if self.dhead1 is not None:\n            all_m_disc_vars += self.dhead1.trainable_variables\n        if self.dhead2 is not None:\n            all_m_disc_vars += self.dhead2.trainable_variables\n        \n        monet_disc_grads = tape.gradient(total_monet_disc_loss, all_m_disc_vars)\n        photo_disc_grads = tape.gradient(photo_disc_loss, self.p_disc.trainable_variables)\n        \n        self.m_disc_opt.apply_gradients(zip(monet_disc_grads, all_m_disc_vars))\n        self.p_disc_opt.apply_gradients(zip(photo_disc_grads, self.p_disc.trainable_variables))\n        \n        # Return metrics\n        return {\n            \"monet_gen_loss\": monet_gen_loss,\n            \"photo_gen_loss\": photo_gen_loss,\n            \"monet_disc_loss\": monet_disc_loss,\n            \"photo_disc_loss\": photo_disc_loss,\n            \"cycle_loss\": cycle_loss,\n            \"identity_loss\": identity_loss\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.499971Z","iopub.execute_input":"2025-03-10T21:24:59.500170Z","iopub.status.idle":"2025-03-10T21:24:59.524813Z","shell.execute_reply.started":"2025-03-10T21:24:59.500150Z","shell.execute_reply":"2025-03-10T21:24:59.523849Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with strategy.scope():\n    # ========================\n    # 1. Initialize Models\n    # ========================\n    monet_generator = Generator()  # Photos → Monet\n    photo_generator = Generator()  # Monet → Photos\n    \n    monet_discriminator = Discriminator()\n    photo_discriminator = Discriminator()\n    dhead1 = DHead()  # For BCE loss\n    dhead2 = DHead()  # For Hinge loss\n\n    # ========================\n    # 2. Define Optimizers\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    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    # ========================\n    # 3. Define Loss Functions\n    # ========================\n    def generator_loss1(generated):\n        return tf.reduce_mean(-generated)  # Hinge loss`\n\n    def generator_loss2(generated):\n        return tf.keras.losses.BinaryCrossentropy(from_logits=True)(\n            tf.ones_like(generated), generated)\n\n    def discriminator_loss1(real, generated):\n        real_loss = tf.reduce_mean(tf.minimum(0., -1. + real))\n        fake_loss = tf.reduce_mean(tf.minimum(0., -1. - generated))\n        return -tf.reduce_mean(real_loss + fake_loss)\n\n    def discriminator_loss2(real, generated):\n        real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)(\n            tf.ones_like(real), real)\n        fake_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)(\n            tf.zeros_like(generated), generated)\n        return 0.5 * (real_loss + fake_loss)\n    \n    def cycle_loss_fn(real, cycled):\n        return tf.reduce_mean(tf.abs(real - cycled))\n    \n    def identity_loss_fn(real, same):\n        return tf.reduce_mean(tf.abs(real - same))\n\n    # ========================\n    # 4. Compile CycleGAN\n    # ========================\n    cycle_gan = CycleGan(\n        monet_generator,\n        photo_generator,\n        monet_discriminator,\n        photo_discriminator,\n        dhead1,\n        dhead2,\n        lambda_cycle=10,  # Stronger cycle consistency\n        lambda_id=0.5     # Weaker identity loss\n    )\n\n    cycle_gan.compile(\n        m_gen_opt=monet_generator_optimizer,\n        p_gen_opt=photo_generator_optimizer,\n        m_disc_opt=monet_discriminator_optimizer,\n        p_disc_opt=photo_discriminator_optimizer,\n        gen_loss_fn1=generator_loss1,\n        gen_loss_fn2=generator_loss2,\n        disc_loss_fn1=discriminator_loss1,\n        disc_loss_fn2=discriminator_loss2,\n        cycle_loss_fn=cycle_loss_fn,\n        identity_loss_fn=identity_loss_fn,\n        aug_fn=aug_fn\n    )\n\n# ========================\n# 4.5 Monitoring\n# ========================\nclass CycleGANMonitor(Callback):\n    def __init__(self, sample_photo, sample_monet, monet_generator, photo_generator, epoch_interval=5):\n        \"\"\"\n        Args:\n            sample_photo: Batch of sample photos (normalized to [-1, 1])\n            sample_monet: Batch of sample Monet paintings (normalized to [-1, 1])\n            monet_generator: Generator that converts photos to Monet style\n            photo_generator: Generator that converts Monet to photos\n            epoch_interval: How often to generate samples (in epochs)\n        \"\"\"\n        self.sample_photo = sample_photo\n        self.sample_monet = sample_monet\n        self.monet_generator = monet_generator\n        self.photo_generator = photo_generator\n        self.epoch_interval = epoch_interval\n\n    def _denormalize(self, image):\n        \"\"\"Convert from [-1, 1] range to [0, 1] for visualization\"\"\"\n        return (image * 0.5) + 0.5\n    @tf.function\n    def generate_predictions(self, images, generator):\n        return generator(images, training=False)\n\n    def _plot_predictions(self, epoch=None):\n        # Generate predictions with tf.function\n        monet_pred = self.generate_predictions(self.sample_photo, self.monet_generator)\n        photo_pred = self.generate_predictions(self.sample_monet, self.photo_generator)\n\n        plt.figure(figsize=(18, 8))\n        \n        # Photo → Monet translations\n        for i in range(min(3, len(self.sample_photo))):  # Show first 3 samples\n            plt.subplot(2, 3, i+1)\n            plt.imshow(self._denormalize(self.sample_photo[i]))\n            plt.title(f\"Input Photo {i+1}\")\n            plt.axis(\"off\")\n            \n            plt.subplot(2, 3, i+4)\n            plt.imshow(self._denormalize(monet_pred[i]))\n            plt.title(f\"Generated Monet {i+1}\" + (f\" (Epoch {epoch})\" if epoch else \"\"))\n            plt.axis(\"off\")\n\n        plt.tight_layout()\n        plt.show()\n\n    def on_epoch_end(self, epoch, logs=None):\n        \"\"\"Generate samples at specified intervals\"\"\"\n        if (epoch+1) % self.epoch_interval == 0:  # +1 to avoid epoch 0\n            self._plot_predictions(epoch+1)\n\n    def on_train_end(self, logs=None):\n        \"\"\"Final visualization after training\"\"\"\n        self._plot_predictions()\n\n# Prepare sample images\nsample_photo = next(iter(photo_ds.take(1)))  # Get 1 batch of photos\nsample_monet = next(iter(monet_ds.take(1)))  # Get 1 batch of Monet paintings\n\n# Create callback\nviz_callback = CycleGANMonitor(\n    sample_photo=sample_photo,\n    sample_monet=sample_monet,\n    monet_generator=monet_generator,\n    photo_generator=photo_generator,\n    epoch_interval=5  # Generate samples every 5 epochs\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:24:59.525868Z","iopub.execute_input":"2025-03-10T21:24:59.526108Z","iopub.status.idle":"2025-03-10T21:25:00.263298Z","shell.execute_reply.started":"2025-03-10T21:24:59.526088Z","shell.execute_reply":"2025-03-10T21:25:00.261674Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ========================\n# 5. Training Loop\n# ========================\nfids = []  # Track FID scores during training\n\nfor epoch in range(1, EPOCHS+1):\n    print(f\"Epoch {epoch}/{EPOCHS}\")\n    \n    # Train for one epoch\n    history = cycle_gan.fit(\n        training_dataset,\n        epochs=1,\n        steps_per_epoch=STEPS_PER_EPOCH,\n        callbacks=[viz_callback],  # Use the visualization callback\n        verbose=1\n    )\n    \n    # Periodic Evaluation\n    if epoch % FID_INTERVAL == 0:\n        try:\n            # Calculate FID\n            fid_score = compute_fid(\n                monet_generator,\n                inception_model,\n                myFID_mu_real,\n                myFID_sigma_real,\n                fid_photo_ds.take(64)  # Use subset for faster evaluation\n            )\n            fids.append(fid_score.numpy())\n            print(f\"FID after epoch {epoch}: {fid_score:.2f}\")\n            \n            # Save checkpoint\n            checkpoint_dir = f\"./cyclegan_checkpoints/epoch_{epoch}\"\n            checkpoint = tf.train.Checkpoint(\n                monet_generator=monet_generator,\n                photo_generator=photo_generator,\n                monet_generator_optimizer=monet_generator_optimizer,\n                photo_generator_optimizer=photo_generator_optimizer\n            )\n            checkpoint.save(checkpoint_dir)\n            print(f\"Checkpoint saved at epoch {epoch}\")\n        except Exception as e:\n            print(f\"Error during FID calculation: {e}\")\n            continue\n\n# Learning rate decay after halfway\nif epoch == EPOCHS // 2:\n    print(\"Reducing learning rate by factor of 10\")\n    # Reduce learning rates\n    monet_generator_optimizer.learning_rate = monet_generator_optimizer.learning_rate * 0.1\n    photo_generator_optimizer.learning_rate = photo_generator_optimizer.learning_rate * 0.1\n    monet_discriminator_optimizer.learning_rate = monet_discriminator_optimizer.learning_rate * 0.1\n    photo_discriminator_optimizer.learning_rate = photo_discriminator_optimizer.learning_rate * 0.1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:25:00.264110Z","iopub.execute_input":"2025-03-10T21:25:00.264384Z","iopub.status.idle":"2025-03-10T21:25:37.723940Z","shell.execute_reply.started":"2025-03-10T21:25:00.264358Z","shell.execute_reply":"2025-03-10T21:25:37.722843Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ========================\n# 6. Final Save\n# ========================\nprint(\"Saving final models...\")\nmonet_generator.save(\"monet_generator_final.h5\")\nphoto_generator.save(\"photo_generator_final.h5\")\n\n# Plot FID progression\nif len(fids) > 0:\n    plt.figure(figsize=(10, 5))\n    plt.plot(np.arange(0, EPOCHS, FID_INTERVAL)[0:len(fids)], fids)\n    plt.title(\"FID Score During Training\")\n    plt.xlabel(\"Epochs\")\n    plt.ylabel(\"FID\")\n    plt.savefig(\"fid_progress.png\")\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:25:37.724802Z","iopub.status.idle":"2025-03-10T21:25:37.725410Z","shell.execute_reply.started":"2025-03-10T21:25:37.725000Z","shell.execute_reply":"2025-03-10T21:25:37.725047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ds_iter = iter(photo_ds)\nfor n_sample in range(3):\n        example_sample = next(ds_iter)\n        generated_sample = monet_generator(example_sample)\n        \n        f = plt.figure(figsize=(32, 32))\n        \n        plt.subplot(121)\n        plt.title('Input image')\n        plt.imshow(example_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        \n        plt.subplot(122)\n        plt.title('Generated image')\n        plt.imshow(generated_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:25:37.726024Z","iopub.status.idle":"2025-03-10T21:25:37.726584Z","shell.execute_reply.started":"2025-03-10T21:25:37.726189Z","shell.execute_reply":"2025-03-10T21:25:37.726225Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create a temporary directory if it doesn't exist\nos.makedirs(\"../tmp\", exist_ok=True)\nzip_path = \"../tmp/images.zip\"\n\n# Create a function to generate and process images in batches\ndef generate_monet_images(photo_dataset, generator_model, batch_size=32):\n    \"\"\"\n    Generates Monet-style images from a photo dataset and saves them directly to a zip file.\n    \n    Args:\n        photo_dataset: TensorFlow dataset of photos\n        generator_model: Trained Monet generator model\n        batch_size: Number of images to process at once\n    \"\"\"\n    start_time = time.time()\n    \n    # Check if dataset is already batched\n    # We won't re-batch it if it's already batched\n    dataset = photo_dataset\n    \n    # Create zip file\n    with zipfile.ZipFile(zip_path, 'w', compression=zipfile.ZIP_DEFLATED) as zipf:\n        image_count = 0\n        \n        # Add a simple progress bar\n        for batch in tqdm(dataset):\n            # Handle potential batching issues by checking the tensor rank\n            tensor_rank = len(batch.shape)\n            \n            # If tensor has too many dimensions, we need to reshape\n            if tensor_rank > 4:  # Should be [batch, height, width, channels]\n                print(f\"Warning: Input tensor has shape {batch.shape}, reshaping...\")\n                # Assuming the first two dimensions are batch dimensions\n                batch_size = batch.shape[0] * batch.shape[1]\n                batch = tf.reshape(batch, [batch_size, *batch.shape[2:]])\n                print(f\"Reshaped to {batch.shape}\")\n            \n            # Generate Monet-style images\n            generated_images = generator_model(batch, training=False)\n            \n            # Convert from [-1, 1] to [0, 255] uint8\n            generated_images = ((generated_images.numpy() * 0.5 + 0.5) * 255).astype(np.uint8)\n            \n            # Process each image in the batch\n            for img in generated_images:\n                # Convert to PIL Image\n                pil_img = Image.fromarray(img)\n                \n                # Create in-memory buffer\n                img_buffer = io.BytesIO()\n                \n                # Save as JPEG to buffer\n                pil_img.save(img_buffer, format='JPEG', quality=95)\n                \n                # Reset buffer position\n                img_buffer.seek(0)\n                \n                # Add to zip file with sequential naming\n                img_name = f\"{image_count + 1}.jpg\"\n                zipf.writestr(img_name, img_buffer.getvalue())\n                \n                image_count += 1\n                \n                # Optional: Add status update every 1000 images\n                if image_count % 1000 == 0:\n                    elapsed = time.time() - start_time\n                    print(f\"Generated {image_count} images in {elapsed:.2f} seconds\")\n    \n    elapsed = time.time() - start_time\n    print(f\"Generation complete! Created {image_count} images in {elapsed:.2f} seconds\")\n    print(f\"Zip file saved to: {zip_path}\")\n    \n    # Optional: Calculate zip file size\n    zip_size_mb = os.path.getsize(zip_path) / (1024 * 1024)\n    print(f\"Zip file size: {zip_size_mb:.2f} MB\")\n    \n    return image_count\n\n# Copy zip to final Kaggle output location\ndef finalize_submission():\n    \"\"\"Copies the zip file to the Kaggle output location\"\"\"\n    import shutil\n    output_path = \"images.zip\"\n    shutil.copy(zip_path, output_path)\n    print(f\"Submission file copied to: {output_path}\")\n    \n    # Verify contents\n    with zipfile.ZipFile(output_path, 'r') as zipf:\n        file_count = len(zipf.namelist())\n        print(f\"Verified zip contains {file_count} images\")\n\n# Prepare a proper dataset for generation\n# CRITICAL: Don't use load_dataset that already batches the data\n# We need to handle the batching carefully here\ndef prepare_photo_dataset(photo_files, batch_size=32):\n    # Load the TFRecord dataset without prior batching\n    dataset = tf.data.TFRecordDataset(photo_files)\n    \n    # Parse the records\n    def read_tfrecord(example):\n        tfrecord_format = {\"image\": tf.io.FixedLenFeature([], tf.string)}\n        example = tf.io.parse_single_example(example, tfrecord_format)\n        image = tf.image.decode_jpeg(example['image'], channels=3)\n        image = (tf.cast(image, tf.float32) / 127.5) - 1\n        image = tf.reshape(image, [256, 256, 3])\n        return image\n    \n    # Map the parsing function and batch\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=tf.data.AUTOTUNE)\n    dataset = dataset.batch(batch_size)\n    dataset = dataset.prefetch(tf.data.AUTOTUNE)\n    \n    return dataset\n\n# 1. Prepare the dataset properly\nphoto_ds_for_generation = prepare_photo_dataset(PHOTO_TFREC)\n\n# 2. Generate images and create zip\ntotal_images = generate_monet_images(photo_ds_for_generation, monet_generator)\n\n# 3. Finalize submission\nfinalize_submission()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T21:25:37.728012Z","iopub.status.idle":"2025-03-10T21:25:37.728630Z","shell.execute_reply.started":"2025-03-10T21:25:37.728177Z","shell.execute_reply":"2025-03-10T21:25:37.728222Z"}},"outputs":[],"execution_count":null}]}