{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":117682,"databundleVersionId":15062069,"sourceType":"competition"},{"sourceId":291858356,"sourceType":"kernelVersion"},{"sourceId":291899286,"sourceType":"kernelVersion"},{"sourceId":292168025,"sourceType":"kernelVersion"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n#for dirname, _, filenames in os.walk('/kaggle/input'):\n#    for filename in filenames:\n#        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:39:58.585025Z","iopub.execute_input":"2026-01-15T07:39:58.585420Z","iopub.status.idle":"2026-01-15T07:39:58.589898Z","shell.execute_reply.started":"2026-01-15T07:39:58.585388Z","shell.execute_reply":"2026-01-15T07:39:58.589134Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"All credits for the model to : https://www.tensorflow.org/tutorials/generative/pix2pix\nExcept as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. \n\nAuthors : Copyright 2019 The TensorFlow Authors.\n\nCode used here is Licensed under the Apache License, Version 2.0 (the \"License\"); \nhttps://github.com/tensorflow/docs/blob/master/site/en/tutorials/generative/pix2pix.ipynb\nChanges to the code: \n* dataset\n* jitter function\n* resolution changed to 320x320","metadata":{}},{"cell_type":"markdown","source":"a conditional generative adversarial network (cGAN) called pix2pix that learns a mapping from input images to output images","metadata":{}},{"cell_type":"markdown","source":"Image-to-Image Translation with Conditional Adversarial Networks\n","metadata":{}},{"cell_type":"markdown","source":"Conclusion first:\nIt makes no sense to waste GPU. As seen in the link below: the model learns labels, or more precisely label - section pairs. The mask dataset (2 value) was replaced in this experiment with the background color\nthe mask label is linked to the section... so ...\nhttps://www.kaggle.com/code/crischir/pix2pix-1c-filter","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"# path = '/kaggle/input/vesuvius-split-dataset/pix2pix_dataset/train'\n# # Count only items that are files\n# file_count = len([f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))])\n\n# print(f\"Total files: {file_count}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:39:58.591089Z","iopub.execute_input":"2026-01-15T07:39:58.591367Z","iopub.status.idle":"2026-01-15T07:39:58.607916Z","shell.execute_reply.started":"2026-01-15T07:39:58.591330Z","shell.execute_reply":"2026-01-15T07:39:58.607270Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Total files: 50841","metadata":{}},{"cell_type":"code","source":"pip install \"protobuf<4.21.0\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:39:58.608619Z","iopub.execute_input":"2026-01-15T07:39:58.608795Z","iopub.status.idle":"2026-01-15T07:40:01.734259Z","shell.execute_reply.started":"2026-01-15T07:39:58.608778Z","shell.execute_reply":"2026-01-15T07:40:01.733515Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tensorflow as tf\n\nimport os\nimport pathlib\nimport time\nimport datetime\n\nfrom matplotlib import pyplot as plt\nfrom IPython import display","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.736340Z","iopub.execute_input":"2026-01-15T07:40:01.736678Z","iopub.status.idle":"2026-01-15T07:40:01.740853Z","shell.execute_reply.started":"2026-01-15T07:40:01.736648Z","shell.execute_reply":"2026-01-15T07:40:01.740125Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"One image sample","metadata":{}},{"cell_type":"code","source":"from IPython.display import Image\nImage(filename='/kaggle/input/dataset-image-vesuvius/combined_out/1004283650_p100.png') ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.741789Z","iopub.execute_input":"2026-01-15T07:40:01.742501Z","iopub.status.idle":"2026-01-15T07:40:01.764998Z","shell.execute_reply.started":"2026-01-15T07:40:01.742479Z","shell.execute_reply":"2026-01-15T07:40:01.764343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from PIL import Image\n\n# Load the image\nimg = Image.open('/kaggle/input/dataset-image-vesuvius/combined_out/1004283650_p100.png')\n\n# Get dimensions\nwidth, height = img.size\n\nprint(f\"Width: {width}px\")\nprint(f\"Height: {height}px\")\nprint(f\"Format: {img.format}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.765946Z","iopub.execute_input":"2026-01-15T07:40:01.766141Z","iopub.status.idle":"2026-01-15T07:40:01.808582Z","shell.execute_reply.started":"2026-01-15T07:40:01.766122Z","shell.execute_reply":"2026-01-15T07:40:01.807894Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Each original image is of size 320 x 640 containing two 320 x 320 images:","metadata":{}},{"cell_type":"code","source":"sample_image = tf.io.read_file(str('/kaggle/input/dataset-image-vesuvius/combined_out/1004283650_p100.png'))\nsample_image = tf.io.decode_jpeg(sample_image)\nprint(sample_image.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.809388Z","iopub.execute_input":"2026-01-15T07:40:01.809602Z","iopub.status.idle":"2026-01-15T07:40:01.829467Z","shell.execute_reply.started":"2026-01-15T07:40:01.809583Z","shell.execute_reply":"2026-01-15T07:40:01.828693Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"One issue is 3 channel with same info ","metadata":{}},{"cell_type":"code","source":"plt.figure()\nplt.imshow(sample_image)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.830319Z","iopub.execute_input":"2026-01-15T07:40:01.830621Z","iopub.status.idle":"2026-01-15T07:40:01.977767Z","shell.execute_reply.started":"2026-01-15T07:40:01.830598Z","shell.execute_reply":"2026-01-15T07:40:01.977146Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"You need to separate real scan images images from the label images—all of which will be of size 320 x 320.\n\nDefine a function that loads image files and outputs two image tensors:","metadata":{}},{"cell_type":"code","source":"# def load(image_file):\n#   # Read and decode an image file to a uint8 tensor\n#   image = tf.io.read_file(image_file)\n#   image = tf.io.decode_jpeg(image)\n\n#   # Split each image tensor into two tensors:\n#   # - one with a real building facade image\n#   # - one with an architecture label image \n#   w = tf.shape(image)[1]\n#   w = w // 2\n#   input_image = image[:, w:, :]\n#   real_image = image[:, :w, :]\n\n#   # Convert both images to float32 tensors\n#   input_image = tf.cast(input_image, tf.float32)\n#   real_image = tf.cast(real_image, tf.float32)\n\n#   return input_image, real_image\n\ndef load(image_file):\n    # Read and decode\n    image = tf.io.read_file(image_file)\n    image = tf.io.decode_jpeg(image)\n\n    # Split the image in half\n    w = tf.shape(image)[1]\n    w = w // 2\n    input_image = image[:, w:, :]\n    real_image = image[:, :w, :]\n    \n    # Get current dimensions\n    # Note: we check only one image since they are usually identical\n    curr_h = tf.shape(input_image)[0]\n    curr_w = tf.shape(input_image)[1]\n\n    # Only resize if dimensions don't match 320x320\n    if curr_h != 320 or curr_w != 320:\n        input_image = tf.image.resize(input_image, [320, 320],\n                                    method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n        real_image = tf.image.resize(real_image, [320, 320],\n                                   method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n\n    # Convert to float32\n    input_image = tf.cast(input_image, tf.float32)\n    real_image = tf.cast(real_image, tf.float32)\n\n    return input_image, real_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.980047Z","iopub.execute_input":"2026-01-15T07:40:01.980299Z","iopub.status.idle":"2026-01-15T07:40:01.986938Z","shell.execute_reply.started":"2026-01-15T07:40:01.980277Z","shell.execute_reply":"2026-01-15T07:40:01.986376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"inp, re = load(str('/kaggle/input/dataset-image-vesuvius/combined_out/1004283650_p111.png'))\n# Casting to int for matplotlib to display the images\nplt.figure()\nplt.imshow(inp / 255.0)\nplt.figure()\nplt.imshow(re / 255.0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:01.987935Z","iopub.execute_input":"2026-01-15T07:40:01.988272Z","iopub.status.idle":"2026-01-15T07:40:02.273679Z","shell.execute_reply.started":"2026-01-15T07:40:01.988249Z","shell.execute_reply":"2026-01-15T07:40:02.273110Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Preprocess\nAs described in the pix2pix paper, you need to apply random jittering and mirroring to preprocess the training set.\nWe will skip the jitter part.\n\n\nOrig=ginal paper say that :\n\nDefine several functions that:\n\n* Resize each 256 x 256 image to a larger height and width—286 x 286.\n* Randomly crop it back to 256 x 256.\n* Randomly flip the image horizontally i.e., left to right (random mirroring).\n* Normalize the images to the [-1, 1] range.","metadata":{}},{"cell_type":"code","source":"path = '/kaggle/input/dataset-image-vesuvius/combined_out/'\n# listdir returns a list of all files and folders\nfiles = os.listdir(path)\nprint(f\"Total number of files: {len(files)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.274507Z","iopub.execute_input":"2026-01-15T07:40:02.274805Z","iopub.status.idle":"2026-01-15T07:40:02.772019Z","shell.execute_reply.started":"2026-01-15T07:40:02.274782Z","shell.execute_reply":"2026-01-15T07:40:02.771246Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# The training set consist of a larger number of images then the original dataset\n# BUFFER_SIZE = 400\n#BUFFER_SIZE = len(files)\nBUFFER_SIZE = 1000\n# The batch size of 1 produced better results for the U-Net in the original pix2pix experiment\nBATCH_SIZE = 1\n# Each image is 256x256 in size\nIMG_WIDTH = 320\nIMG_HEIGHT = 320","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.773084Z","iopub.execute_input":"2026-01-15T07:40:02.773473Z","iopub.status.idle":"2026-01-15T07:40:02.777069Z","shell.execute_reply.started":"2026-01-15T07:40:02.773450Z","shell.execute_reply":"2026-01-15T07:40:02.776403Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Image resize","metadata":{}},{"cell_type":"code","source":"def resize(input_image, real_image, height, width):\n  input_image = tf.image.resize(input_image, [height, width],\n                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n  real_image = tf.image.resize(real_image, [height, width],\n                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n\n  return input_image, real_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.777950Z","iopub.execute_input":"2026-01-15T07:40:02.778363Z","iopub.status.idle":"2026-01-15T07:40:02.795903Z","shell.execute_reply.started":"2026-01-15T07:40:02.778330Z","shell.execute_reply":"2026-01-15T07:40:02.795145Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Crop image","metadata":{}},{"cell_type":"code","source":"def random_crop(input_image, real_image):\n  stacked_image = tf.stack([input_image, real_image], axis=0)\n  cropped_image = tf.image.random_crop(\n      stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])\n\n  return cropped_image[0], cropped_image[1]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.796752Z","iopub.execute_input":"2026-01-15T07:40:02.797015Z","iopub.status.idle":"2026-01-15T07:40:02.809079Z","shell.execute_reply.started":"2026-01-15T07:40:02.796990Z","shell.execute_reply":"2026-01-15T07:40:02.808447Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Normalize","metadata":{}},{"cell_type":"code","source":"# Normalizing the images to [-1, 1]\ndef normalize(input_image, real_image):\n  input_image = (input_image / 127.5) - 1\n  real_image = (real_image / 127.5) - 1\n\n  return input_image, real_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.809886Z","iopub.execute_input":"2026-01-15T07:40:02.810161Z","iopub.status.idle":"2026-01-15T07:40:02.823162Z","shell.execute_reply.started":"2026-01-15T07:40:02.810124Z","shell.execute_reply":"2026-01-15T07:40:02.822447Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Augumentation","metadata":{}},{"cell_type":"code","source":"@tf.function()\ndef random_jitter(input_image, real_image):\n  #No resizing for now  \n  # # Resizing to 286x286\n  # input_image, real_image = resize(input_image, real_image, 300, 300)\n\n  # # Random cropping back to 256x256\n  # input_image, real_image = random_crop(input_image, real_image)\n\n  if tf.random.uniform(()) > 0.5:\n    # Random mirroring\n    input_image = tf.image.flip_left_right(input_image)\n    real_image = tf.image.flip_left_right(real_image)\n\n# # 2. Random vertical flipping (Up-Down)\n#   if tf.random.uniform(()) > 0.3:\n#     input_image = tf.image.flip_up_down(input_image)\n#     real_image = tf.image.flip_up_down(real_image)\n\n  return input_image, real_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.824089Z","iopub.execute_input":"2026-01-15T07:40:02.824409Z","iopub.status.idle":"2026-01-15T07:40:02.837877Z","shell.execute_reply.started":"2026-01-15T07:40:02.824374Z","shell.execute_reply":"2026-01-15T07:40:02.837182Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Inspect jitter ","metadata":{}},{"cell_type":"markdown","source":"Input","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(6, 6))\nfor i in range(4):\n  rj_inp, rj_re = random_jitter(inp, re)\n  plt.subplot(2, 2, i + 1)\n  plt.imshow(rj_inp / 255.0)\n  plt.axis('off')\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:02.838732Z","iopub.execute_input":"2026-01-15T07:40:02.838982Z","iopub.status.idle":"2026-01-15T07:40:03.034812Z","shell.execute_reply.started":"2026-01-15T07:40:02.838961Z","shell.execute_reply":"2026-01-15T07:40:03.034084Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Output","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(6, 6))\nfor i in range(4):\n  rj_inp, rj_re = random_jitter(inp, re)\n  plt.subplot(2, 2, i + 1)\n  plt.imshow(rj_re / 255.0)\n  plt.axis('off')\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:03.035757Z","iopub.execute_input":"2026-01-15T07:40:03.036017Z","iopub.status.idle":"2026-01-15T07:40:03.137368Z","shell.execute_reply.started":"2026-01-15T07:40:03.035995Z","shell.execute_reply":"2026-01-15T07:40:03.136776Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Load images for train","metadata":{}},{"cell_type":"code","source":"def load_image_train(image_file):\n  input_image, real_image = load(image_file)\n  input_image, real_image = random_jitter(input_image, real_image)\n  input_image, real_image = normalize(input_image, real_image)\n\n  return input_image, real_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:03.138140Z","iopub.execute_input":"2026-01-15T07:40:03.138442Z","iopub.status.idle":"2026-01-15T07:40:03.142354Z","shell.execute_reply.started":"2026-01-15T07:40:03.138417Z","shell.execute_reply":"2026-01-15T07:40:03.141725Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_image_test(image_file):\n  input_image, real_image = load(image_file)\n  input_image, real_image = resize(input_image, real_image,\n                                   IMG_HEIGHT, IMG_WIDTH)\n  input_image, real_image = normalize(input_image, real_image)\n\n  return input_image, real_image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:03.143231Z","iopub.execute_input":"2026-01-15T07:40:03.143568Z","iopub.status.idle":"2026-01-15T07:40:03.156418Z","shell.execute_reply.started":"2026-01-15T07:40:03.143534Z","shell.execute_reply":"2026-01-15T07:40:03.155705Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Build an input pipeline with tf.data","metadata":{}},{"cell_type":"markdown","source":"Shuffle files","metadata":{}},{"cell_type":"code","source":"##comment if you consider shuffle is needed\nprint('shuffle dataset path')\nimport random\n# 1. Get the list of all file paths\nall_files = tf.io.gfile.glob('/kaggle/input/dataset-image-vesuvius/combined_out/*.png')\n\n# 2. SHUFFLE THE FILENAMES (This uses almost zero RAM)\n# This ensures that even with a small buffer, the model sees images \n# from totally different parts of the dataset every epoch.\nrandom.shuffle(all_files)\ntrain_dataset = tf.data.Dataset.from_tensor_slices(all_files)\n# 4. Map the load function\ntrain_dataset = train_dataset.map(load, num_parallel_calls=tf.data.AUTOTUNE)\n\n# 5. Use a small buffer for \"local\" shuffling\n# Since the filenames were already shuffled, a buffer of 100 is now \n# just as effective as a buffer of 10,000.\ntrain_dataset = train_dataset.shuffle(buffer_size=BUFFER_SIZE)\ntrain_dataset = train_dataset.batch(BATCH_SIZE)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"No shuffle\")\n##uncomment if shuffle is not needed\n#train_dataset = tf.data.Dataset.list_files(str('/kaggle/input/vesuvius-split-dataset/pix2pix_dataset/train/*.png'))\n# train_dataset = train_dataset.map(load_image_train,\n#                                   num_parallel_calls=tf.data.AUTOTUNE)\n# train_dataset = train_dataset.shuffle(BUFFER_SIZE)\n# train_dataset = train_dataset.batch(BATCH_SIZE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:03.157275Z","iopub.execute_input":"2026-01-15T07:40:03.157563Z","iopub.status.idle":"2026-01-15T07:40:23.554715Z","shell.execute_reply.started":"2026-01-15T07:40:03.157532Z","shell.execute_reply":"2026-01-15T07:40:23.553945Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"try:\n    # Attempt to load using the glob pattern\n    test_dataset = tf.data.Dataset.list_files(str('/kaggle/input/vesuvius-split-dataset/pix2pix_dataset/test/*.png'))\nexcept tf.errors.InvalidArgumentError:\n    # Corrected syntax: removed the extra ')' and added the glob pattern\n    test_dataset = tf.data.Dataset.list_files('/kaggle/input/vesuvius-split-dataset/pix2pix_dataset/test/*.png')\n\ntest_dataset = test_dataset.map(load_image_test)\ntest_dataset = test_dataset.batch(BATCH_SIZE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:23.555731Z","iopub.execute_input":"2026-01-15T07:40:23.556013Z","iopub.status.idle":"2026-01-15T07:40:26.477431Z","shell.execute_reply.started":"2026-01-15T07:40:23.555982Z","shell.execute_reply":"2026-01-15T07:40:26.476664Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# try:\n#   test_dataset = tf.data.Dataset.list_files(str('/kaggle/input/vesuvius-split-dataset/pix2pix_dataset/test/*.png'))\n# except tf.errors.InvalidArgumentError:\n#   test_dataset = tf.data.Dataset.list_files('/kaggle/input/vesuvius-split-dataset/pix2pix_dataset/test/'))\n# test_dataset = test_dataset.map(load_image_test)\n# test_dataset = test_dataset.batch(BATCH_SIZE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.478418Z","iopub.execute_input":"2026-01-15T07:40:26.478699Z","iopub.status.idle":"2026-01-15T07:40:26.482032Z","shell.execute_reply.started":"2026-01-15T07:40:26.478671Z","shell.execute_reply":"2026-01-15T07:40:26.481471Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Build the generator\n\nThe generator of your pix2pix cGAN is a modified U-Net. A U-Net consists of an encoder (downsampler) and decoder (upsampler). (You can find out more about it in the Image segmentation tutorial and on the U-Net project website.)\n\nEach block in the encoder is: Convolution -> Batch normalization -> Leaky ReLU\nEach block in the decoder is: Transposed convolution -> Batch normalization -> Dropout (applied to the first 3 blocks) -> ReLU\nThere are skip connections between the encoder and decoder (as in the U-Net).\nDefine the downsampler (encoder):\n","metadata":{}},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.485131Z","iopub.execute_input":"2026-01-15T07:40:26.485444Z","iopub.status.idle":"2026-01-15T07:40:26.498594Z","shell.execute_reply.started":"2026-01-15T07:40:26.485410Z","shell.execute_reply":"2026-01-15T07:40:26.497911Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def downsample(filters, size, apply_batchnorm=True):\n  initializer = tf.random_normal_initializer(0., 0.02)\n\n  result = tf.keras.Sequential()\n  result.add(\n      tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n\n  if apply_batchnorm:\n    result.add(tf.keras.layers.BatchNormalization())\n\n  result.add(tf.keras.layers.LeakyReLU())\n\n  return result","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.499259Z","iopub.execute_input":"2026-01-15T07:40:26.499823Z","iopub.status.idle":"2026-01-15T07:40:26.513171Z","shell.execute_reply.started":"2026-01-15T07:40:26.499800Z","shell.execute_reply":"2026-01-15T07:40:26.512676Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"down_model = downsample(3, 4)\ndown_result = down_model(tf.expand_dims(inp, 0))\nprint (down_result.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.514017Z","iopub.execute_input":"2026-01-15T07:40:26.514385Z","iopub.status.idle":"2026-01-15T07:40:26.551833Z","shell.execute_reply.started":"2026-01-15T07:40:26.514353Z","shell.execute_reply":"2026-01-15T07:40:26.551159Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False):\n  initializer = tf.random_normal_initializer(0., 0.02)\n\n  result = tf.keras.Sequential()\n  result.add(\n    tf.keras.layers.Conv2DTranspose(filters, size, strides=2,\n                                    padding='same',\n                                    kernel_initializer=initializer,\n                                    use_bias=False))\n\n  result.add(tf.keras.layers.BatchNormalization())\n\n  if apply_dropout:\n      result.add(tf.keras.layers.Dropout(0.5))\n\n  result.add(tf.keras.layers.ReLU())\n\n  return result","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.552698Z","iopub.execute_input":"2026-01-15T07:40:26.552915Z","iopub.status.idle":"2026-01-15T07:40:26.557626Z","shell.execute_reply.started":"2026-01-15T07:40:26.552895Z","shell.execute_reply":"2026-01-15T07:40:26.556946Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"up_model = upsample(3, 4)\nup_result = up_model(down_result)\nprint (up_result.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.558500Z","iopub.execute_input":"2026-01-15T07:40:26.558954Z","iopub.status.idle":"2026-01-15T07:40:26.588032Z","shell.execute_reply.started":"2026-01-15T07:40:26.558933Z","shell.execute_reply":"2026-01-15T07:40:26.587299Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Adapted for 320 .... to be checked out ","metadata":{}},{"cell_type":"code","source":"def Generator():\n  # Change input shape to 320x320\n  # inputs = tf.keras.layers.Input(shape=[320, 320, 3])\n  inputs = tf.keras.layers.Input(shape=[IMG_HEIGHT, IMG_WIDTH, 3])\n  down_stack = [\n    downsample(64, 4, apply_batchnorm=False),  # (bs, 160, 160, 64)\n    downsample(128, 4),  # (bs, 80, 80, 128)\n    downsample(256, 4),  # (bs, 40, 40, 256)\n    downsample(512, 4),  # (bs, 20, 20, 512)\n    downsample(512, 4),  # (bs, 10, 10, 512)\n    downsample(512, 4),  # (bs, 5, 5, 512)\n    # We stop here or add one more. If we add another, it becomes 3x3 or 2x2 \n    # depending on padding. Let's keep 6 levels for 320 to avoid tiny feature maps.\n  ]\n\n  up_stack = [\n    upsample(512, 4, apply_dropout=True),  # (bs, 10, 10, 1024)\n    upsample(512, 4, apply_dropout=True),  # (bs, 20, 20, 1024)\n    upsample(512, 4, apply_dropout=True),  # (bs, 40, 40, 1024)\n    upsample(256, 4),  # (bs, 80, 80, 512)\n    upsample(128, 4),  # (bs, 160, 160, 256)\n  ]\n\n  initializer = tf.random_normal_initializer(0., 0.02)\n  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n                                         strides=2,\n                                         padding='same',\n                                         kernel_initializer=initializer,\n                                         activation='tanh')  # (bs, 320, 320, 3)\n\n  x = inputs\n\n  # Downsampling through the model\n  skips = []\n  for down in down_stack:\n    x = down(x)\n    skips.append(x)\n\n  skips = reversed(skips[:-1])\n\n  # Upsampling and establishing the skip connections\n  for up, skip in zip(up_stack, skips):\n    x = up(x)\n    x = tf.keras.layers.Concatenate()([x, skip])\n\n  x = last(x)\n\n  return tf.keras.Model(inputs=inputs, outputs=x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.588995Z","iopub.execute_input":"2026-01-15T07:40:26.589298Z","iopub.status.idle":"2026-01-15T07:40:26.595824Z","shell.execute_reply.started":"2026-01-15T07:40:26.589265Z","shell.execute_reply":"2026-01-15T07:40:26.595118Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"generator = Generator()\ntf.keras.utils.plot_model(generator, show_shapes=True, dpi=64)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.596658Z","iopub.execute_input":"2026-01-15T07:40:26.596956Z","iopub.status.idle":"2026-01-15T07:40:26.923909Z","shell.execute_reply.started":"2026-01-15T07:40:26.596913Z","shell.execute_reply":"2026-01-15T07:40:26.923200Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Test the generator:","metadata":{}},{"cell_type":"code","source":"gen_output = generator(inp[tf.newaxis, ...], training=False)\nplt.imshow(gen_output[0, ...])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:26.924905Z","iopub.execute_input":"2026-01-15T07:40:26.925221Z","iopub.status.idle":"2026-01-15T07:40:27.145360Z","shell.execute_reply.started":"2026-01-15T07:40:26.925187Z","shell.execute_reply":"2026-01-15T07:40:27.144612Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Define the generator loss\n\nGANs learn a loss that adapts to the data, while cGANs learn a structured loss that penalizes a possible structure that differs from the network output and the target image, as described in the pix2pix paper.\n\nThe generator loss is a sigmoid cross-entropy loss of the generated images and an array of ones.\nThe pix2pix paper also mentions the L1 loss, which is a MAE (mean absolute error) between the generated image and the target image.\nThis allows the generated image to become structurally similar to the target image.\nThe formula to calculate the total generator loss is gan_loss + LAMBDA * l1_loss, where LAMBDA = 100. This value was decided by the authors of the paper.\n\n","metadata":{}},{"cell_type":"code","source":"LAMBDA = 100","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.146235Z","iopub.execute_input":"2026-01-15T07:40:27.146552Z","iopub.status.idle":"2026-01-15T07:40:27.150052Z","shell.execute_reply.started":"2026-01-15T07:40:27.146520Z","shell.execute_reply":"2026-01-15T07:40:27.149371Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.151044Z","iopub.execute_input":"2026-01-15T07:40:27.151513Z","iopub.status.idle":"2026-01-15T07:40:27.166044Z","shell.execute_reply.started":"2026-01-15T07:40:27.151491Z","shell.execute_reply":"2026-01-15T07:40:27.165362Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generator_loss(disc_generated_output, gen_output, target):\n  gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output)\n\n  # Mean absolute error\n  l1_loss = tf.reduce_mean(tf.abs(target - gen_output))\n\n  total_gen_loss = gan_loss + (LAMBDA * l1_loss)\n\n  return total_gen_loss, gan_loss, l1_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.166958Z","iopub.execute_input":"2026-01-15T07:40:27.167213Z","iopub.status.idle":"2026-01-15T07:40:27.182412Z","shell.execute_reply.started":"2026-01-15T07:40:27.167193Z","shell.execute_reply":"2026-01-15T07:40:27.181665Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Build the discriminator\n\nThe discriminator in the pix2pix cGAN is a convolutional PatchGAN classifier—it tries to classify if each image patch is real or not real, as described in the pix2pix paper.\n\nEach block in the discriminator is: Convolution -> Batch normalization -> Leaky ReLU.\nThe shape of the output after the last layer is (batch_size, 30, 30, 1).\nEach 30 x 30 image patch of the output classifies a 70 x 70 portion of the input image.\nThe discriminator receives 2 inputs:\nThe input image and the target image, which it should classify as real.\nThe input image and the generated image (the output of the generator), which it should classify as fake.\nUse tf.concat([inp, tar], axis=-1) to concatenate these 2 inputs together.\nLet's define the discriminator:\n\n\n","metadata":{}},{"cell_type":"code","source":"def Discriminator():\n  initializer = tf.random_normal_initializer(0., 0.02)\n\n  # 1. Update input shapes to 320x320\n  # inp = tf.keras.layers.Input(shape=[320, 320, 3], name='input_image')\n  # tar = tf.keras.layers.Input(shape=[320, 320, 3], name='target_image')\n  inp = tf.keras.layers.Input(shape=[IMG_HEIGHT, IMG_WIDTH, 3], name='input_image')\n  tar = tf.keras.layers.Input(shape=[IMG_HEIGHT, IMG_WIDTH, 3], name='target_image')\n  x = tf.keras.layers.concatenate([inp, tar])  # (batch_size, 320, 320, 6)\n\n  # 2. Downsampling steps\n  down1 = downsample(64, 4, False)(x)  # (batch_size, 160, 160, 64)\n  down2 = downsample(128, 4)(down1)    # (batch_size, 80, 80, 128)\n  down3 = downsample(256, 4)(down2)    # (batch_size, 40, 40, 256)\n\n  # 3. Final Patch Construction\n  # Adding zero padding to maintain spatial control before the stride-1 convolution\n  zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)  # (batch_size, 42, 42, 256)\n  \n  conv = tf.keras.layers.Conv2D(512, 4, strides=1,\n                                kernel_initializer=initializer,\n                                use_bias=False)(zero_pad1)  # (batch_size, 39, 39, 512)\n\n  batchnorm1 = tf.keras.layers.BatchNormalization()(conv)\n  leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)\n\n  zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  # (batch_size, 41, 41, 512)\n\n  # Final 1-channel prediction map\n  last = tf.keras.layers.Conv2D(1, 4, strides=1,\n                                kernel_initializer=initializer)(zero_pad2)  # (batch_size, 38, 38, 1)\n\n  return tf.keras.Model(inputs=[inp, tar], outputs=last)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.183048Z","iopub.execute_input":"2026-01-15T07:40:27.183264Z","iopub.status.idle":"2026-01-15T07:40:27.195893Z","shell.execute_reply.started":"2026-01-15T07:40:27.183244Z","shell.execute_reply":"2026-01-15T07:40:27.195181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"discriminator = Discriminator()\ntf.keras.utils.plot_model(discriminator, show_shapes=True, dpi=64)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.196800Z","iopub.execute_input":"2026-01-15T07:40:27.197180Z","iopub.status.idle":"2026-01-15T07:40:27.337966Z","shell.execute_reply.started":"2026-01-15T07:40:27.197144Z","shell.execute_reply":"2026-01-15T07:40:27.337214Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Test discriminator","metadata":{}},{"cell_type":"code","source":"disc_out = discriminator([inp[tf.newaxis, ...], gen_output], training=False)\nplt.imshow(disc_out[0, ..., -1], vmin=-20, vmax=20, cmap='RdBu_r')\nplt.colorbar()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.338889Z","iopub.execute_input":"2026-01-15T07:40:27.339172Z","iopub.status.idle":"2026-01-15T07:40:27.538970Z","shell.execute_reply.started":"2026-01-15T07:40:27.339144Z","shell.execute_reply":"2026-01-15T07:40:27.538182Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Define the discriminator loss\n\nThe discriminator_loss function takes 2 inputs: real images and generated images.\nreal_loss is a sigmoid cross-entropy loss of the real images and an array of ones(since these are the real images).\ngenerated_loss is a sigmoid cross-entropy loss of the generated images and an array of zeros (since these are the fake images).\nThe total_loss is the sum of real_loss and generated_loss.\n\n","metadata":{}},{"cell_type":"code","source":"def discriminator_loss(disc_real_output, disc_generated_output):\n  real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output)\n\n  generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output)\n\n  total_disc_loss = real_loss + generated_loss\n\n  return total_disc_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.539938Z","iopub.execute_input":"2026-01-15T07:40:27.540188Z","iopub.status.idle":"2026-01-15T07:40:27.544178Z","shell.execute_reply.started":"2026-01-15T07:40:27.540165Z","shell.execute_reply":"2026-01-15T07:40:27.543592Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Define the optimizers and a checkpoint-saver","metadata":{}},{"cell_type":"code","source":"generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\ndiscriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.545207Z","iopub.execute_input":"2026-01-15T07:40:27.545497Z","iopub.status.idle":"2026-01-15T07:40:27.563787Z","shell.execute_reply.started":"2026-01-15T07:40:27.545467Z","shell.execute_reply":"2026-01-15T07:40:27.563229Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"checkpoint_dir = './training_checkpoints'\ncheckpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\ncheckpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n                                 discriminator_optimizer=discriminator_optimizer,\n                                 generator=generator,\n                                 discriminator=discriminator)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.564578Z","iopub.execute_input":"2026-01-15T07:40:27.564868Z","iopub.status.idle":"2026-01-15T07:40:27.571753Z","shell.execute_reply.started":"2026-01-15T07:40:27.564847Z","shell.execute_reply":"2026-01-15T07:40:27.571230Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Generate images\n\nWrite a function to plot some images during training.\n\nPass images from the test set to the generator.\nThe generator will then translate the input image into the output.\nThe last step is to plot the predictions and voila!\n","metadata":{}},{"cell_type":"code","source":"def generate_images(model, test_input, tar):\n  prediction = model(test_input, training=True)\n  plt.figure(figsize=(15, 15))\n\n  display_list = [test_input[0], tar[0], prediction[0]]\n  title = ['Input Image', 'Ground Truth', 'Predicted Image']\n\n  for i in range(3):\n    plt.subplot(1, 3, i+1)\n    plt.title(title[i])\n    # Getting the pixel values in the [0, 1] range to plot.\n    plt.imshow(display_list[i] * 0.5 + 0.5)\n    plt.axis('off')\n  plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.572703Z","iopub.execute_input":"2026-01-15T07:40:27.573041Z","iopub.status.idle":"2026-01-15T07:40:27.586988Z","shell.execute_reply.started":"2026-01-15T07:40:27.573008Z","shell.execute_reply":"2026-01-15T07:40:27.586487Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Test the function:","metadata":{}},{"cell_type":"code","source":"for example_input, example_target in test_dataset.take(1):\n  generate_images(generator, example_input, example_target)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.587877Z","iopub.execute_input":"2026-01-15T07:40:27.588123Z","iopub.status.idle":"2026-01-15T07:40:27.956222Z","shell.execute_reply.started":"2026-01-15T07:40:27.588101Z","shell.execute_reply":"2026-01-15T07:40:27.955649Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for example_input, example_target in test_dataset.take(2):\n  generate_images(generator, example_input, example_target)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:27.957059Z","iopub.execute_input":"2026-01-15T07:40:27.957275Z","iopub.status.idle":"2026-01-15T07:40:28.641577Z","shell.execute_reply.started":"2026-01-15T07:40:27.957255Z","shell.execute_reply":"2026-01-15T07:40:28.640830Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Training\n\nFor each example input generates an output.\nThe discriminator receives the input_image and the generated image as the first input. The second input is the input_image and the target_image.\nNext, calculate the generator and the discriminator loss.\nThen, calculate the gradients of loss with respect to both the generator and the discriminator variables(inputs) and apply those to the optimizer.\nFinally, log the losses to TensorBoard.","metadata":{}},{"cell_type":"code","source":"log_dir=\"logs/\"\n\nsummary_writer = tf.summary.create_file_writer(\n  log_dir + \"fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:28.642562Z","iopub.execute_input":"2026-01-15T07:40:28.643055Z","iopub.status.idle":"2026-01-15T07:40:28.648338Z","shell.execute_reply.started":"2026-01-15T07:40:28.643031Z","shell.execute_reply":"2026-01-15T07:40:28.647731Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@tf.function\ndef train_step(input_image, target, step):\n  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n    gen_output = generator(input_image, training=True)\n\n    disc_real_output = discriminator([input_image, target], training=True)\n    disc_generated_output = discriminator([input_image, gen_output], training=True)\n\n    gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target)\n    disc_loss = discriminator_loss(disc_real_output, disc_generated_output)\n\n  generator_gradients = gen_tape.gradient(gen_total_loss,\n                                          generator.trainable_variables)\n  discriminator_gradients = disc_tape.gradient(disc_loss,\n                                               discriminator.trainable_variables)\n\n  generator_optimizer.apply_gradients(zip(generator_gradients,\n                                          generator.trainable_variables))\n  discriminator_optimizer.apply_gradients(zip(discriminator_gradients,\n                                              discriminator.trainable_variables))\n\n  with summary_writer.as_default():\n    tf.summary.scalar('gen_total_loss', gen_total_loss, step=step//1000)\n    tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=step//1000)\n    tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=step//1000)\n    tf.summary.scalar('disc_loss', disc_loss, step=step//1000)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:28.649189Z","iopub.execute_input":"2026-01-15T07:40:28.649462Z","iopub.status.idle":"2026-01-15T07:40:28.660766Z","shell.execute_reply.started":"2026-01-15T07:40:28.649440Z","shell.execute_reply":"2026-01-15T07:40:28.660092Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def fit(train_ds, test_ds, steps):\n  example_input, example_target = next(iter(test_ds.take(1)))\n  start = time.time()\n\n  for step, (input_image, target) in train_ds.repeat().take(steps).enumerate():\n    if (step) % 1000 == 0:\n      display.clear_output(wait=True)\n\n      if step != 0:\n        print(f'Time taken for 1000 steps: {time.time()-start:.2f} sec\\n')\n\n      start = time.time()\n\n      generate_images(generator, example_input, example_target)\n      print(f\"Step: {step//1000}k\")\n\n    train_step(input_image, target, step)\n\n    # Training step\n    if (step+1) % 10 == 0:\n      print('.', end='', flush=True)\n\n\n    # Save (checkpoint) the model every 10k steps\n    if (step + 1) % 10000 == 0:\n      checkpoint.save(file_prefix=checkpoint_prefix)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:28.661738Z","iopub.execute_input":"2026-01-15T07:40:28.662290Z","iopub.status.idle":"2026-01-15T07:40:28.679757Z","shell.execute_reply.started":"2026-01-15T07:40:28.662268Z","shell.execute_reply":"2026-01-15T07:40:28.679095Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The actual training loop. Since this tutorial can run of more than one dataset, and the datasets vary greatly in size the training loop is setup to work in steps instead of epochs.\n\nIterates over the number of steps.\nEvery 10 steps print a dot (.).\nEvery 1k steps: clear the display and run generate_images to show the progress.\nEvery 5k steps: save a checkpoint.","metadata":{}},{"cell_type":"code","source":"# fit(train_dataset, test_dataset, steps=40000)\nfit(train_dataset, test_dataset, steps=200000)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T07:40:28.680572Z","iopub.execute_input":"2026-01-15T07:40:28.681214Z","iopub.status.idle":"2026-01-15T09:15:19.738147Z","shell.execute_reply.started":"2026-01-15T07:40:28.681190Z","shell.execute_reply":"2026-01-15T09:15:19.737567Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %load_ext tensorboard\n# %tensorboard --logdir {log_dir}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T09:16:17.801573Z","iopub.execute_input":"2026-01-15T09:16:17.802185Z","iopub.status.idle":"2026-01-15T09:16:24.844601Z","shell.execute_reply.started":"2026-01-15T09:16:17.802153Z","shell.execute_reply":"2026-01-15T09:16:24.843974Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#ls {checkpoint_dir}\n# Restoring the latest checkpoint in checkpoint_dir\n#checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T09:15:19.739187Z","iopub.execute_input":"2026-01-15T09:15:19.739433Z","iopub.status.idle":"2026-01-15T09:15:19.742813Z","shell.execute_reply.started":"2026-01-15T09:15:19.739410Z","shell.execute_reply":"2026-01-15T09:15:19.742292Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run the trained model on a few examples from the test set\nfor inp, tar in test_dataset.take(15):\n  generate_images(generator, inp, tar)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T09:15:19.743696Z","iopub.execute_input":"2026-01-15T09:15:19.744290Z","iopub.status.idle":"2026-01-15T09:15:22.210617Z","shell.execute_reply.started":"2026-01-15T09:15:19.744257Z","shell.execute_reply":"2026-01-15T09:15:22.209903Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run the trained model on a few examples from the test set\nfor inp, tar in test_dataset.take(10):\n  generate_images(generator, inp, tar)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T09:16:27.333732Z","iopub.execute_input":"2026-01-15T09:16:27.334380Z","iopub.status.idle":"2026-01-15T09:16:30.502738Z","shell.execute_reply.started":"2026-01-15T09:16:27.334351Z","shell.execute_reply":"2026-01-15T09:16:30.502147Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# https://www.kaggle.com/code/crischir/vesuvius-2d-filter \n\nimport tensorflow as tf\nimport numpy as np\nfrom PIL import Image\nimport os\n\n# --- 1. Replicate the Architecture (Must match the .ipynb exactly) ---\ndef downsample(filters, size, apply_batchnorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',\n                                      kernel_initializer=initializer, use_bias=False))\n    if apply_batchnorm:\n        result.add(tf.keras.layers.BatchNormalization())\n    result.add(tf.keras.layers.LeakyReLU())\n    return result\n\ndef upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same',\n                                               kernel_initializer=initializer, use_bias=False))\n    result.add(tf.keras.layers.BatchNormalization())\n    if apply_dropout:\n        result.add(tf.keras.layers.Dropout(0.5))\n    result.add(tf.keras.layers.ReLU())\n    return result\n\ndef Generator():\n    inputs = tf.keras.layers.Input(shape=[256, 256, 3]) # Update to 320, 320 if using your adapted version\n    down_stack = [\n        downsample(64, 4, apply_batchnorm=False),\n        downsample(128, 4),\n        downsample(256, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n    ]\n    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 = tf.keras.layers.Conv2DTranspose(3, 4, strides=2, padding='same',\n                                           kernel_initializer=initializer, activation='tanh')\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 = tf.keras.layers.Concatenate()([x, skip])\n    x = last(x)\n    return tf.keras.Model(inputs=inputs, outputs=x)\n\n# --- 2. Load and Restore Checkpoint ---\ngenerator = Generator()\n# Checkpoint must include the keys used in training (even if just doing inference)\ncheckpoint = tf.train.Checkpoint(generator=generator) \n\n# Path to the directory containing 'ckpt-1.index', etc.\ncheckpoint_dir = './training_checkpoints' \ncheckpoint.restore(tf.train.latest_checkpoint(checkpoint_dir)).expect_partial()\n\n# --- 3. Preprocess and Predict ---\ndef predict_single_image(image_path):\n    # Load and resize\n    img = Image.open(image_path).convert('RGB')\n    img = img.resize((256, 256)) # Ensure this matches the Generator input shape\n    img_array = np.array(img)\n    \n    # Normalize to [-1, 1]\n    img_array = (img_array / 127.5) - 1\n    img_tensor = tf.expand_dims(tf.cast(img_array, tf.float32), 0)\n\n    # Run inference\n    prediction = generator(img_tensor, training=False)\n    \n    # Convert back to [0, 255]\n    prediction = (prediction[0] * 0.5 + 0.5) * 255.0\n    return prediction.numpy().astype(np.uint8)\n\n# Usage\nresult = predict_single_image('your_input_image.jpg')\nImage.fromarray(result).save('prediction_output.png')","metadata":{"_kg_hide-output":true,"_kg_hide-input":true}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}