main.ipynb

{
 "metadata": {
  "kernelspec": {
   "language": "python",
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.12",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  }
 },
 "nbformat_minor": 4,
 "nbformat": 4,
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Team: Orel Ben Zaken 31 Omer Luxembourg 20"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Introduction and Setup\n",
    "\n",
    "https://github.com/omerlux/Something-of-a-Painter"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "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",
    "\n",
    "import numpy as np # linear algebra\n",
    "import 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",
    "\n",
    "import 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",
    "execution": {
     "iopub.status.busy": "2022-08-09T02:03:59.147541Z",
     "iopub.execute_input": "2022-08-09T02:03:59.148013Z",
     "iopub.status.idle": "2022-08-09T02:04:00.484543Z",
     "shell.execute_reply.started": "2022-08-09T02:03:59.147976Z",
     "shell.execute_reply": "2022-08-09T02:04:00.483210Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 30,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "import os, random, json, PIL, shutil, re\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from kaggle_datasets import KaggleDatasets\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras.layers as L\n",
    "import tensorflow_addons as tfa\n",
    "from tensorflow.keras import Model, losses, optimizers\n",
    "from tqdm import tqdm\n",
    "import cv2\n",
    "\n",
    "try:\n",
    "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
    "    print(f'Running on TPU {tpu.master()}')\n",
    "except ValueError:\n",
    "    tpu = None\n",
    "\n",
    "if tpu:\n",
    "    tf.config.experimental_connect_to_cluster(tpu)\n",
    "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
    "    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
    "else:\n",
    "    strategy = tf.distribute.get_strategy()\n",
    "\n",
    "\n",
    "REPLICAS = strategy.num_replicas_in_sync\n",
    "print(f'REPLICAS: {REPLICAS}')\n",
    "AUTO = tf.data.experimental.AUTOTUNE"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:00.486617Z",
     "iopub.execute_input": "2022-08-09T02:04:00.486980Z",
     "iopub.status.idle": "2022-08-09T02:04:00.497660Z",
     "shell.execute_reply.started": "2022-08-09T02:04:00.486947Z",
     "shell.execute_reply": "2022-08-09T02:04:00.496762Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 31,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Model parameters"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "HEIGHT = 256\n",
    "WIDTH = 256\n",
    "CHANNELS = 3\n",
    "OUTPUT_CHANNELS = 3\n",
    "EPOCHS = 8\n",
    "BATCH_SIZE = 16"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:00.498629Z",
     "iopub.execute_input": "2022-08-09T02:04:00.498986Z",
     "iopub.status.idle": "2022-08-09T02:04:00.513746Z",
     "shell.execute_reply.started": "2022-08-09T02:04:00.498953Z",
     "shell.execute_reply": "2022-08-09T02:04:00.512681Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 32,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load data\n",
    "We want to keep our photo dataset and our Monet dataset separate. First, load in the filenames of the TFRecords.\n",
    "\n",
    "All the images for the competition are already sized to 256x256. As these images are RGB images, set the channel to 3. Additionally, we need to scale the images to a [-1, 1] scale. Because we are building a generative model, we don't need the labels or the image id so we'll only return the image from the TFRecord."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "GCS_PATH = KaggleDatasets().get_gcs_path()\n",
    "\n",
    "MONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\n",
    "PHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\n",
    "\n",
    "def count_data_items(filenames):\n",
    "    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n",
    "    return np.sum(n)\n",
    "\n",
    "n_monet_samples = count_data_items(MONET_FILENAMES)\n",
    "n_photo_samples = count_data_items(PHOTO_FILENAMES)\n",
    "\n",
    "print(f'Monet TFRecord files: {len(MONET_FILENAMES)}')\n",
    "print(f'Monet image files: {n_monet_samples}')\n",
    "print(f'Photo TFRecord files: {len(PHOTO_FILENAMES)}')\n",
    "print(f'Photo image files: {n_photo_samples}')"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:00.516229Z",
     "iopub.execute_input": "2022-08-09T02:04:00.516646Z",
     "iopub.status.idle": "2022-08-09T02:04:01.063199Z",
     "shell.execute_reply.started": "2022-08-09T02:04:00.516608Z",
     "shell.execute_reply": "2022-08-09T02:04:01.061961Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 33,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "bar plot of the number of images per type (MONET or PHOTOs)"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_axes([0,0,1,1])\n",
    "image_names = ['MONET', 'PHOTOS']\n",
    "values = [n_monet_samples, n_photo_samples]\n",
    "ax.bar(image_names[0], values[0], color = 'r')\n",
    "ax.bar(image_names[1], values[1], color = 'b')\n",
    "plt.show()"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:01.064759Z",
     "iopub.execute_input": "2022-08-09T02:04:01.065107Z",
     "iopub.status.idle": "2022-08-09T02:04:01.198405Z",
     "shell.execute_reply.started": "2022-08-09T02:04:01.065075Z",
     "shell.execute_reply": "2022-08-09T02:04:01.197306Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 34,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Histograms - distribution of pixels at each channel\n",
    "*"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "BASE_FOLDER_PATH = \"../input/gan-getting-started/\"\n",
    "MONET_FOLDER_PATH = os.path.join(BASE_FOLDER_PATH, \"monet_jpg\")\n",
    "PHOTO_FOLDER_PATH = os.path.join(BASE_FOLDER_PATH, \"photo_jpg\")"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:01.200146Z",
     "iopub.execute_input": "2022-08-09T02:04:01.201074Z",
     "iopub.status.idle": "2022-08-09T02:04:01.207407Z",
     "shell.execute_reply.started": "2022-08-09T02:04:01.201009Z",
     "shell.execute_reply": "2022-08-09T02:04:01.206344Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 35,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "def image_type_read_all(path):\n",
    "    image_list = []\n",
    "    for i, image_name in enumerate(os.listdir(path)):\n",
    "        if i==300:\n",
    "            break\n",
    "        image = cv2.imread(os.path.join(path, image_name))\n",
    "        image_list.append(image)\n",
    "    \n",
    "    all_images = np.stack(image_list)\n",
    "    print(all_images.shape)\n",
    "    return all_images\n",
    "        \n",
    "\n",
    "print(f\"Monet images:\")\n",
    "monet_all_images = image_type_read_all(MONET_FOLDER_PATH)\n",
    "print(\"-\" * 10)\n",
    "print(f\"Photo images:\")\n",
    "photo_all_images = image_type_read_all(PHOTO_FOLDER_PATH)\n",
    "print(\"-\" * 10)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:01.209714Z",
     "iopub.execute_input": "2022-08-09T02:04:01.210447Z",
     "iopub.status.idle": "2022-08-09T02:04:02.643573Z",
     "shell.execute_reply.started": "2022-08-09T02:04:01.210374Z",
     "shell.execute_reply": "2022-08-09T02:04:02.642509Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 36,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "print(f\"Monet images:\")\n",
    "monet_all_images = monet_all_images.reshape((-1,3))\n",
    "print(monet_all_images.shape)\n",
    "print(\"-\" * 10)\n",
    "print(f\"Photo images:\")\n",
    "photo_all_images = photo_all_images.reshape((-1,3))\n",
    "print(photo_all_images.shape)\n",
    "print(\"-\" * 10)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:02.645237Z",
     "iopub.execute_input": "2022-08-09T02:04:02.645832Z",
     "iopub.status.idle": "2022-08-09T02:04:02.651879Z",
     "shell.execute_reply.started": "2022-08-09T02:04:02.645794Z",
     "shell.execute_reply": "2022-08-09T02:04:02.651132Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 37,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "x = monet_all_images / 127.5 - 1\n",
    "y = photo_all_images / 127.5 - 1\n",
    "\n",
    "bins = np.linspace(-2, 2, 20)\n",
    "\n",
    "fig, (axR, axG, axB) = plt.subplots(3)\n",
    "fig.suptitle('Histograms - distribution of pixels at each channel')\n",
    "axR.hist(x[:,0], bins, alpha=0.5, color='r', label='Monet')\n",
    "axR.hist(y[:,0], bins, alpha=0.5, color='b', label='PHOTOS')\n",
    "axR.legend(loc='upper right')\n",
    "axR.set_title('Red channel')\n",
    "axG.hist(x[:,1], bins, alpha=0.5, color='r', label='Monet')\n",
    "axG.hist(y[:,1], bins, alpha=0.5, color='b', label='PHOTOS')\n",
    "axG.legend(loc='upper right')\n",
    "axG.set_title('Green channel')\n",
    "axB.hist(x[:,2], bins, alpha=0.5, color='r', label='Monet')\n",
    "axB.hist(y[:,2], bins, alpha=0.5, color='b', label='PHOTOS')\n",
    "axB.legend(loc='upper right')\n",
    "axB.set_title('Blue channel')\n",
    "axB.set_xlabel('pixel value')\n",
    "\n",
    "# plt.hist(x, bins, alpha=0.5, label='Monet')\n",
    "# plt.hist(y, bins, alpha=0.5, label='PHOTOS')\n",
    "# plt.legend(loc='upper right')\n",
    "fig.tight_layout()\n",
    "fig.show()"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:02.652980Z",
     "iopub.execute_input": "2022-08-09T02:04:02.653951Z",
     "iopub.status.idle": "2022-08-09T02:04:11.123948Z",
     "shell.execute_reply.started": "2022-08-09T02:04:02.653916Z",
     "shell.execute_reply": "2022-08-09T02:04:11.122920Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 38,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Auxiliar functions\n",
    "Auxiliary functions for: read organization and augmentation the dataset"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def decode_image(image):\n",
    "    image = tf.image.decode_jpeg(image, channels=CHANNELS)\n",
    "    image = (tf.cast(image, tf.float32) / 127.5) - 1\n",
    "    image = tf.reshape(image, [HEIGHT, WIDTH, CHANNELS])\n",
    "    return image\n",
    "\n",
    "def read_tfrecord(example):\n",
    "    tfrecord_format = {\n",
    "        'image_name': tf.io.FixedLenFeature([], tf.string),\n",
    "        'image':      tf.io.FixedLenFeature([], tf.string),\n",
    "        'target':     tf.io.FixedLenFeature([], tf.string)\n",
    "    }\n",
    "    example = tf.io.parse_single_example(example, tfrecord_format)\n",
    "    image = decode_image(example['image'])\n",
    "    return image\n",
    "\n",
    "def load_dataset(filenames):\n",
    "    dataset = tf.data.TFRecordDataset(filenames)\n",
    "    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)\n",
    "    return dataset\n",
    "\n",
    "def get_gan_dataset(monet_files, photo_files, augment=None, repeat=True, shuffle=True, batch_size=1):\n",
    "\n",
    "    monet_ds = load_dataset(monet_files)\n",
    "    photo_ds = load_dataset(photo_files)\n",
    "\n",
    "#     if augment:\n",
    "        \n",
    "        \n",
    "    if repeat:\n",
    "        monet_ds = monet_ds.repeat()\n",
    "        photo_ds = photo_ds.repeat()\n",
    "    if shuffle:\n",
    "        monet_ds = monet_ds.shuffle(2048)#max(n_monet_samples, n_photo_samples))\n",
    "        photo_ds = photo_ds.shuffle(2048)#max(n_monet_samples, n_photo_samples))\n",
    "        \n",
    "    monet_ds = monet_ds.batch(batch_size, drop_remainder=True)\n",
    "    photo_ds = photo_ds.batch(batch_size, drop_remainder=True)\n",
    "    monet_ds = monet_ds.cache()\n",
    "    photo_ds = photo_ds.cache()\n",
    "    monet_ds = monet_ds.prefetch(AUTO)\n",
    "    photo_ds = photo_ds.prefetch(AUTO)\n",
    "    \n",
    "    gan_ds = tf.data.Dataset.zip((monet_ds, photo_ds))\n",
    "    \n",
    "    return gan_ds\n",
    "\n",
    "def display_samples(ds, row, col):\n",
    "    ds_iter = iter(ds)\n",
    "    plt.figure(figsize=(15, int(15*row/col)))\n",
    "    for j in range(row*col):\n",
    "        example_sample = next(ds_iter)\n",
    "        plt.subplot(row,col,j+1)\n",
    "        plt.axis('off')\n",
    "        plt.imshow(example_sample[0] * 0.5 + 0.5)\n",
    "    plt.show()"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:11.127313Z",
     "iopub.execute_input": "2022-08-09T02:04:11.127736Z",
     "iopub.status.idle": "2022-08-09T02:04:11.141922Z",
     "shell.execute_reply.started": "2022-08-09T02:04:11.127698Z",
     "shell.execute_reply": "2022-08-09T02:04:11.141030Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 39,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "def plot_simple_augmentation(ds, transform):\n",
    "    ds_iter = iter(ds)\n",
    "    example_sample = next(ds_iter)\n",
    "    img = example_sample[0] * 0.5 + 0.5\n",
    "    \n",
    "    plt.figure(figsize=(10, 5))\n",
    "    \n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.imshow(img)\n",
    "    plt.axis(\"off\")\n",
    "    \n",
    "    plt.subplot(1, 2, 2)\n",
    "    x = transform(image=img)\n",
    "    plt.imshow(x)\n",
    "    plt.axis(\"off\")\n",
    "\n",
    "    plt.show()\n",
    "    \n",
    "def display_augmentation_samples(ds, transform,num_images=1):\n",
    "    ds_iter = iter(ds)\n",
    "      \n",
    "    plt.figure(figsize=(5*num_images, 10))\n",
    "    \n",
    "    for j in range(num_images):\n",
    "        example_sample = next(ds_iter)\n",
    "        img = example_sample * 0.5 + 0.5\n",
    "        x = transform(img)\n",
    "        plt.subplot(2,num_images,j+1)\n",
    "        plt.axis('off')\n",
    "        plt.imshow(img[0])\n",
    "        plt.subplot(2,num_images,(j+num_images)+1)\n",
    "        plt.axis('off')\n",
    "        plt.imshow(x[0])\n",
    "    plt.show()"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:11.142981Z",
     "iopub.execute_input": "2022-08-09T02:04:11.143656Z",
     "iopub.status.idle": "2022-08-09T02:04:11.159380Z",
     "shell.execute_reply.started": "2022-08-09T02:04:11.143620Z",
     "shell.execute_reply": "2022-08-09T02:04:11.158120Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 40,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Augmentations"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Dataset Augmentations - preprocessing\n",
    "\n",
    "Data augmentation for GANs should be done very carefully, especially for tasks similar to style transfer, if we apply transformations that can change too much the style of the data (e.g. brightness, contrast, saturation) it can cause the generator to do not efficiently learn the base style, so in this case, we are using only spatial transformations like, flips, rotates and crops.\n"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def DataAugment(height=HEIGHT, width=WIDTH, channels=CHANNELS):\n",
    "    inputs = L.Input(shape=[height, width, channels])\n",
    "\n",
    "    flip = tf.keras.layers.experimental.preprocessing.RandomFlip('horizontal')\n",
    "    rotate = tf.keras.layers.experimental.preprocessing.RandomRotation(factor=(-.25, 0.25), fill_mode='reflect')\n",
    "    crop = tf.keras.layers.experimental.preprocessing.RandomCrop(int(height / 1.5), int(width / 1.5))\n",
    "    resize = tf.keras.layers.experimental.preprocessing.Resizing(height, width)\n",
    "\n",
    "    outputs = flip(inputs)\n",
    "    outputs = rotate(outputs)\n",
    "    outputs = crop(outputs)\n",
    "    outputs = resize(outputs)\n",
    "    return Model(inputs=inputs, outputs=outputs)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:11.161044Z",
     "iopub.execute_input": "2022-08-09T02:04:11.162090Z",
     "iopub.status.idle": "2022-08-09T02:04:11.175125Z",
     "shell.execute_reply.started": "2022-08-09T02:04:11.162052Z",
     "shell.execute_reply": "2022-08-09T02:04:11.174016Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 41,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Diff-Augmentation"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "with strategy.scope():\n",
    "# Differentiable Augmentation for Data-Efficient GAN Training\n",
    "# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han\n",
    "# https://arxiv.org/pdf/2006.10738\n",
    "\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",
    "        factor = 1  # 1    # between (-factor/2, factor/2) around 0\n",
    "        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) * factor - factor / 2          # - 0.5\n",
    "        x = x + magnitude\n",
    "        return x\n",
    "\n",
    "\n",
    "    def rand_saturation(x):\n",
    "        factor = 2  # 2    # between (1 - factor/2, 1 + factor/2) - around 1\n",
    "        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) * factor + (1 - factor / 2)    # * 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",
    "\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,\n",
    "                                  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,\n",
    "                                  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]]),\n",
    "                                      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),\n",
    "                                     dtype=tf.int32)\n",
    "        offset_y = tf.random.uniform([tf.shape(x)[0], 1, 1], maxval=image_size[1] + (1 - cutout_size[1] % 2),\n",
    "                                     dtype=tf.int32)\n",
    "        grid_batch, grid_x, grid_y = tf.meshgrid(tf.range(batch_size, dtype=tf.int32),\n",
    "                                                 tf.range(cutout_size[0], dtype=tf.int32),\n",
    "                                                 tf.range(cutout_size[1], dtype=tf.int32), indexing='ij')\n",
    "        cutout_grid = tf.stack(\n",
    "            [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(\n",
    "            1 - tf.scatter_nd(cutout_grid, tf.ones([batch_size, cutout_size[0], cutout_size[1]], dtype=tf.float32),\n",
    "                              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": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:11.177055Z",
     "iopub.execute_input": "2022-08-09T02:04:11.177793Z",
     "iopub.status.idle": "2022-08-09T02:04:11.205496Z",
     "shell.execute_reply.started": "2022-08-09T02:04:11.177756Z",
     "shell.execute_reply": "2022-08-09T02:04:11.204487Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 42,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "def data_augment(image):\n",
    "    \n",
    "    p_spatial = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n",
    "    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n",
    "#     p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n",
    "    \n",
    "    # 90º rotations\n",
    "    if p_rotate > .8:\n",
    "        image = tf.image.rot90(image, k=3) # rotate 270º\n",
    "    elif p_rotate > .6:\n",
    "        image = tf.image.rot90(image, k=2) # rotate 180º\n",
    "    elif p_rotate > .4:\n",
    "        image = tf.image.rot90(image, k=1) # rotate 90º\n",
    "        \n",
    "    # Flips\n",
    "    image = tf.image.random_flip_left_right(image)\n",
    "    image = tf.image.random_flip_up_down(image)\n",
    "    if p_spatial > .75:\n",
    "        image = tf.image.transpose(image)\n",
    "        \n",
    "            \n",
    "    # Train on crops\n",
    "    image = tf.image.random_crop(image, size=[HEIGHT, WIDTH, CHANNELS])\n",
    "        \n",
    "    \n",
    "    return image"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:11.206968Z",
     "iopub.execute_input": "2022-08-09T02:04:11.207362Z",
     "iopub.status.idle": "2022-08-09T02:04:11.222484Z",
     "shell.execute_reply.started": "2022-08-09T02:04:11.207327Z",
     "shell.execute_reply": "2022-08-09T02:04:11.221483Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 43,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Visualization\n",
    "\n",
    "Let's visualize a few Monet example."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "display_samples(load_dataset(MONET_FILENAMES).batch(1), 2, 5)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:11.225415Z",
     "iopub.execute_input": "2022-08-09T02:04:11.226216Z",
     "iopub.status.idle": "2022-08-09T02:04:12.574089Z",
     "shell.execute_reply.started": "2022-08-09T02:04:11.226177Z",
     "shell.execute_reply": "2022-08-09T02:04:12.573012Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 44,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's visualize a few photo example."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "display_samples(load_dataset(PHOTO_FILENAMES).batch(1), 2, 5)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:12.575714Z",
     "iopub.execute_input": "2022-08-09T02:04:12.576103Z",
     "iopub.status.idle": "2022-08-09T02:04:14.202807Z",
     "shell.execute_reply.started": "2022-08-09T02:04:12.576063Z",
     "shell.execute_reply": "2022-08-09T02:04:14.201654Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 45,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Visualize augmentation few Monet example."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Dataset Augmentations - preprocessing"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "display_augmentation_samples(load_dataset(MONET_FILENAMES).batch(1), DataAugment(), num_images=4)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:14.204391Z",
     "iopub.execute_input": "2022-08-09T02:04:14.204790Z",
     "iopub.status.idle": "2022-08-09T02:04:16.953470Z",
     "shell.execute_reply.started": "2022-08-09T02:04:14.204755Z",
     "shell.execute_reply": "2022-08-09T02:04:16.952081Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 46,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Diff-Augmentation"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "display_augmentation_samples(load_dataset(MONET_FILENAMES).batch(1), aug_fn, num_images=4)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:16.955682Z",
     "iopub.execute_input": "2022-08-09T02:04:16.956728Z",
     "iopub.status.idle": "2022-08-09T02:04:18.225093Z",
     "shell.execute_reply.started": "2022-08-09T02:04:16.956679Z",
     "shell.execute_reply": "2022-08-09T02:04:18.224048Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 47,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Build the model "
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's first define our downsample and upsample methods (we'll be using a UNET architecture for our generator).\n",
    "\n",
    "**The downsample**, as the name suggests, reduces the 2D dimensions, the width and height, of the image by the stride. The stride is the length of the step the filter takes. Since the stride is 2, the filter is applied to every other pixel, hence reducing the weight and height by 2.\n",
    "\n",
    "**Upsample** does the opposite of downsample and increases the dimensions of the of the image. Conv2DTranspose does basically the opposite of a Conv2D layer.\n",
    "\n",
    "We'll be using an instance normalization instead of batch normalization. As the instance normalization is not standard in the TensorFlow API, we'll use the layer from TensorFlow Add-ons."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def downsample(filters, size, apply_instancenorm=True, strides=2):\n",
    "    initializer = tf.random_normal_initializer(0., 0.02)\n",
    "    gamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n",
    "\n",
    "    result = tf.keras.Sequential()\n",
    "    result.add(L.Conv2D(filters, size, strides=strides, padding='same',\n",
    "                        kernel_initializer=initializer, use_bias=False))\n",
    "\n",
    "    if apply_instancenorm:\n",
    "        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n",
    "\n",
    "    result.add(L.LeakyReLU())\n",
    "\n",
    "    return result\n",
    "\n",
    "def upsample(filters, size, apply_dropout=False, strides=2):\n",
    "    initializer = tf.random_normal_initializer(0., 0.02)\n",
    "    gamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n",
    "\n",
    "    result = tf.keras.Sequential()\n",
    "    result.add(L.Conv2DTranspose(filters, size, strides=strides, padding='same',\n",
    "                                 kernel_initializer=initializer, use_bias=False))\n",
    "\n",
    "    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n",
    "\n",
    "    if apply_dropout:\n",
    "        result.add(L.Dropout(0.5))\n",
    "\n",
    "    result.add(L.ReLU())\n",
    "\n",
    "    return result"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.226885Z",
     "iopub.execute_input": "2022-08-09T02:04:18.227950Z",
     "iopub.status.idle": "2022-08-09T02:04:18.237570Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.227909Z",
     "shell.execute_reply": "2022-08-09T02:04:18.236766Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 48,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Generator\n",
    "\n",
    "The generator first downsamples the input image and then upsample while establishing long skip connections. Skip connections are a way to help bypass the vanishing gradient problem by concatenating the output of a layer to multiple layers instead of only one. Here we concatenate the output of the downsample layer to the upsample layer in a symmetrical fashion."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def Generator():\n",
    "    inputs = L.Input(shape=[HEIGHT, WIDTH, CHANNELS])\n",
    "\n",
    "    # bs = batch size\n",
    "    down_stack = [\n",
    "        downsample(64, 4, apply_instancenorm=False), # (bs, 128, 128, 64)\n",
    "        downsample(128, 4),                          # (bs, 64, 64, 128)\n",
    "        downsample(256, 4),                          # (bs, 32, 32, 256)\n",
    "        downsample(512, 4),                          # (bs, 16, 16, 512)\n",
    "        downsample(512, 4),                          # (bs, 8, 8, 512)\n",
    "        downsample(512, 4),                          # (bs, 4, 4, 512)\n",
    "        downsample(512, 4),                          # (bs, 2, 2, 512)\n",
    "        downsample(512, 4),                          # (bs, 1, 1, 512)\n",
    "    ]\n",
    "\n",
    "    up_stack = [\n",
    "        upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)\n",
    "        upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)\n",
    "        upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)\n",
    "        upsample(512, 4),                     # (bs, 16, 16, 1024)\n",
    "        upsample(256, 4),                     # (bs, 32, 32, 512)\n",
    "        upsample(128, 4),                     # (bs, 64, 64, 256)\n",
    "        upsample(64, 4),                      # (bs, 128, 128, 128)\n",
    "    ]\n",
    "\n",
    "    initializer = tf.random_normal_initializer(0., 0.02)\n",
    "    last = L.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n",
    "                             strides=2,\n",
    "                             padding='same',\n",
    "                             kernel_initializer=initializer,\n",
    "                             activation='tanh') # (bs, 256, 256, 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 = L.Concatenate()([x, skip])\n",
    "\n",
    "    x = last(x)\n",
    "\n",
    "    return Model(inputs=inputs, outputs=x)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.238732Z",
     "iopub.execute_input": "2022-08-09T02:04:18.239595Z",
     "iopub.status.idle": "2022-08-09T02:04:18.255740Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.239560Z",
     "shell.execute_reply": "2022-08-09T02:04:18.254508Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 49,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Discriminator\n",
    "\n",
    "The discriminator takes in the input image and classifies it as real or fake (generated). Instead of outputing a single node, the discriminator outputs a smaller 2D image with higher pixel values indicating a real classification and lower values indicating a fake classification."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def Discriminator():\n",
    "    initializer = tf.random_normal_initializer(0., 0.02)\n",
    "    gamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n",
    "\n",
    "    inp = L.Input(shape=[HEIGHT, WIDTH, CHANNELS], name='input_image')\n",
    "\n",
    "    x = inp\n",
    "\n",
    "    down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64)\n",
    "    down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)\n",
    "    down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)\n",
    "\n",
    "    zero_pad1 = L.ZeroPadding2D()(down3) # (bs, 34, 34, 256)\n",
    "    conv = L.Conv2D(512, 4, strides=1,\n",
    "                    kernel_initializer=initializer,\n",
    "                    use_bias=False)(zero_pad1) # (bs, 31, 31, 512)\n",
    "\n",
    "    norm1 = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(conv)\n",
    "\n",
    "    leaky_relu = L.LeakyReLU()(norm1)\n",
    "\n",
    "    zero_pad2 = L.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)\n",
    "\n",
    "    last = L.Conv2D(1, 4, strides=1,\n",
    "                    kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1)\n",
    "\n",
    "    return Model(inputs=inp, outputs=last)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.257167Z",
     "iopub.execute_input": "2022-08-09T02:04:18.257634Z",
     "iopub.status.idle": "2022-08-09T02:04:18.272563Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.257581Z",
     "shell.execute_reply": "2022-08-09T02:04:18.271186Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 50,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### The whole model (CycleGAN)"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We will subclass a tf.keras.Model so that we can run fit() later to train our model. During the training step, the model transforms a photo to a Monet painting and then back to a photo. The difference between the original photo and the twice-transformed photo is the cycle-consistency loss. We want the original photo and the twice-transformed photo to be similar to one another.\n",
    "\n",
    "The losses are defined in the next section."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "class CycleGan(Model):\n",
    "    def __init__(\n",
    "        self, dsaug_layer,\n",
    "        monet_generator,\n",
    "        photo_generator,\n",
    "        monet_discriminator,\n",
    "        photo_discriminator,\n",
    "        lambda_cycle=10,\n",
    "    ):\n",
    "        super(CycleGan, self).__init__()\n",
    "        self.dsaug_layer = dsaug_layer\n",
    "        self.m_gen = monet_generator\n",
    "        self.p_gen = photo_generator\n",
    "        self.m_disc = monet_discriminator\n",
    "        self.p_disc = photo_discriminator\n",
    "        self.lambda_cycle = lambda_cycle\n",
    "\n",
    "    def compile(\n",
    "        self,\n",
    "        m_gen_optimizer,\n",
    "        p_gen_optimizer,\n",
    "        m_disc_optimizer,\n",
    "        p_disc_optimizer,\n",
    "        gen_loss_fn,\n",
    "        disc_loss_fn,\n",
    "        cycle_loss_fn,\n",
    "        identity_loss_fn,\n",
    "        ds_augment,\n",
    "        diffaugment\n",
    "    ):\n",
    "        super(CycleGan, self).compile()\n",
    "        self.m_gen_optimizer = m_gen_optimizer\n",
    "        self.p_gen_optimizer = p_gen_optimizer\n",
    "        self.m_disc_optimizer = m_disc_optimizer\n",
    "        self.p_disc_optimizer = p_disc_optimizer\n",
    "        self.gen_loss_fn = gen_loss_fn\n",
    "        self.disc_loss_fn = disc_loss_fn\n",
    "        self.cycle_loss_fn = cycle_loss_fn\n",
    "        self.identity_loss_fn = identity_loss_fn\n",
    "        self.ds_augment = ds_augment\n",
    "        self.diffaugment =  diffaugment\n",
    "        \n",
    "    def train_step(self, batch_data):\n",
    "        real_monet, real_photo = batch_data\n",
    "        \n",
    "        if self.ds_augment:\n",
    "            real_monet = self.dsaug_layer(real_monet)\n",
    "            real_photo = self.dsaug_layer(real_photo)\n",
    "        \n",
    "        batch_size = tf.shape(real_monet)[0]\n",
    "        \n",
    "        with tf.GradientTape(persistent=True) as tape:\n",
    "            # photo to monet back to photo\n",
    "            fake_monet = self.m_gen(real_photo, training=True)\n",
    "            cycled_photo = self.p_gen(fake_monet, training=True)\n",
    "\n",
    "            # monet to photo back to monet\n",
    "            fake_photo = self.p_gen(real_monet, training=True)\n",
    "            cycled_monet = self.m_gen(fake_photo, training=True)\n",
    "\n",
    "            # generating itself\n",
    "            same_monet = self.m_gen(real_monet, training=True)\n",
    "            same_photo = self.p_gen(real_photo, training=True)\n",
    "\n",
    "           # \"Diffaugment\": Augmentation is done before the discriminator\n",
    "            if len(self.diffaugment) != 0:\n",
    "                both_monet = tf.concat([real_monet, fake_monet], axis=0)\n",
    "                aug_monet = DiffAugment(both_monet, self.diffaugment)\n",
    "                real_monet = aug_monet[:batch_size]     # AUGMENTED!\n",
    "                fake_monet = aug_monet[batch_size:]     # AUGMENTED!\n",
    "\n",
    "            # discriminator used to check, inputing real images\n",
    "            disc_real_monet = self.m_disc(real_monet, training=True)\n",
    "            disc_real_photo = self.p_disc(real_photo, training=True)\n",
    "\n",
    "            # discriminator used to check, inputing fake images\n",
    "            disc_fake_monet = self.m_disc(fake_monet, training=True)\n",
    "            disc_fake_photo = self.p_disc(fake_photo, training=True)\n",
    "\n",
    "            # evaluates generator loss\n",
    "            monet_gen_loss = self.gen_loss_fn(disc_fake_monet)\n",
    "            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)\n",
    "\n",
    "            # evaluates total cycle consistency loss\n",
    "            total_cycle_loss = self.cycle_loss_fn(real_monet, cycled_monet, self.lambda_cycle) + self.cycle_loss_fn(real_photo, cycled_photo, self.lambda_cycle)\n",
    "\n",
    "            # evaluates total generator loss\n",
    "            total_monet_gen_loss = monet_gen_loss + total_cycle_loss + self.identity_loss_fn(real_monet, same_monet, self.lambda_cycle)\n",
    "            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo, self.lambda_cycle)\n",
    "\n",
    "            # evaluates discriminator loss\n",
    "            monet_disc_loss = self.disc_loss_fn(disc_real_monet, disc_fake_monet)\n",
    "            photo_disc_loss = self.disc_loss_fn(disc_real_photo, disc_fake_photo)\n",
    "\n",
    "        # Calculate the gradients for generator and discriminator\n",
    "        monet_generator_gradients = tape.gradient(total_monet_gen_loss,\n",
    "                                                  self.m_gen.trainable_variables)\n",
    "        photo_generator_gradients = tape.gradient(total_photo_gen_loss,\n",
    "                                                  self.p_gen.trainable_variables)\n",
    "\n",
    "        monet_discriminator_gradients = tape.gradient(monet_disc_loss,\n",
    "                                                      self.m_disc.trainable_variables)\n",
    "        photo_discriminator_gradients = tape.gradient(photo_disc_loss,\n",
    "                                                      self.p_disc.trainable_variables)\n",
    "\n",
    "        # Apply the gradients to the optimizer\n",
    "        self.m_gen_optimizer.apply_gradients(zip(monet_generator_gradients,\n",
    "                                                 self.m_gen.trainable_variables))\n",
    "\n",
    "        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients,\n",
    "                                                 self.p_gen.trainable_variables))\n",
    "\n",
    "        self.m_disc_optimizer.apply_gradients(zip(monet_discriminator_gradients,\n",
    "                                                  self.m_disc.trainable_variables))\n",
    "\n",
    "        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients,\n",
    "                                                  self.p_disc.trainable_variables))\n",
    "        \n",
    "        return {\n",
    "            'monet_gen_loss': total_monet_gen_loss,\n",
    "            'photo_gen_loss': total_photo_gen_loss,\n",
    "            'monet_disc_loss': monet_disc_loss,\n",
    "            'photo_disc_loss': photo_disc_loss,\n",
    "            'total_cycle_loss': total_cycle_loss\n",
    "        }"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.274232Z",
     "iopub.execute_input": "2022-08-09T02:04:18.274876Z",
     "iopub.status.idle": "2022-08-09T02:04:18.298339Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.274837Z",
     "shell.execute_reply": "2022-08-09T02:04:18.297345Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 51,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Loss functions\n",
    "\n",
    "The discriminator loss function below compares real images to a matrix of 1s and fake images to a matrix of 0s. The perfect discriminator will output all 1s for real images and all 0s for fake images. The discriminator loss outputs the average of the real and generated loss."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "with strategy.scope():\n",
    "    # Discriminator loss {0: fake, 1: real} (The discriminator loss outputs the average of the real and generated loss)\n",
    "    def discriminator_loss(real, generated):\n",
    "        real_loss = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(real), real)\n",
    "\n",
    "        generated_loss = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.zeros_like(generated), generated)\n",
    "\n",
    "        total_disc_loss = real_loss + generated_loss\n",
    "\n",
    "        return total_disc_loss * 0.5"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.299721Z",
     "iopub.execute_input": "2022-08-09T02:04:18.300828Z",
     "iopub.status.idle": "2022-08-09T02:04:18.313895Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.300786Z",
     "shell.execute_reply": "2022-08-09T02:04:18.312798Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 52,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "The generator wants to fool the discriminator into thinking the generated image is real. The perfect generator will have the discriminator output only 1s. Thus, it compares the generated image to a matrix of 1s to find the loss."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "with strategy.scope():\n",
    "# Generator loss\n",
    "    def generator_loss(generated):\n",
    "        return losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(generated), generated)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.315657Z",
     "iopub.execute_input": "2022-08-09T02:04:18.316352Z",
     "iopub.status.idle": "2022-08-09T02:04:18.325855Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.316312Z",
     "shell.execute_reply": "2022-08-09T02:04:18.324798Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 53,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "We want our original photo and the twice transformed photo to be similar to one another. Thus, we can calculate the cycle consistency loss be finding the average of their difference."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Cycle consistency loss (measures if original photo and the twice transformed photo to be similar to one another)\n",
    "with strategy.scope():\n",
    "    def calc_cycle_loss(real_image, cycled_image, LAMBDA):\n",
    "        loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n",
    "\n",
    "        return LAMBDA * loss1"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.327153Z",
     "iopub.execute_input": "2022-08-09T02:04:18.328205Z",
     "iopub.status.idle": "2022-08-09T02:04:18.341313Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.328158Z",
     "shell.execute_reply": "2022-08-09T02:04:18.340505Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 54,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "The identity loss compares the image with its generator (i.e. photo with photo generator). If given a photo as input, we want it to generate the same image as the image was originally a photo. The identity loss compares the input with the output of the generator."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Identity loss (compares the image with its generator (i.e. photo with photo generator))\n",
    "with strategy.scope():\n",
    "    def identity_loss(real_image, same_image, LAMBDA):\n",
    "        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n",
    "        return LAMBDA * 0.5 * loss"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.342581Z",
     "iopub.execute_input": "2022-08-09T02:04:18.343506Z",
     "iopub.status.idle": "2022-08-09T02:04:18.352685Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.343456Z",
     "shell.execute_reply": "2022-08-09T02:04:18.351876Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 55,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Train the CycleGAN \n",
    "\n",
    "**In practice the training was done on an external GPU and the graphs printed were using the interface Wandb**\n",
    "\n",
    "**Over Sampling**(OS) - change the 'step_size' argument in the line with 'gan_model.fit' to: **steps_per_epoch=(max(n_monet_samples, n_photo_samples)//4**\n",
    "\n",
    "**Under Sampling**(US) - change the 'step_size' argument in the line with 'gan_model.fit' to: **steps_per_epoch=300**\n",
    "\n",
    "**Dataset Augmentation**(Aug) - change the 'ds_augment' argument in the line with 'gan_model.compile' to: **ds_augment=True**. This model can only be run on the cpu and gpu (not with the tpu).\n",
    "\n",
    "**Diff-Augmentation**(DiffAug) - change the 'diffaugment' argument in the line with 'gan_model.compile', fro empty string \"\" to: **diffaugment=\"color,translation,cutout\"**\n"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "with strategy.scope():\n",
    "    dsaug_layer = DataAugment()\n",
    "    monet_generator = Generator() # transforms photos to Monet-esque paintings\n",
    "    photo_generator = Generator() # transforms Monet paintings to be more like photos\n",
    "\n",
    "    monet_discriminator = Discriminator() # differentiates real Monet paintings and generated Monet paintings\n",
    "    photo_discriminator = Discriminator() # differentiates real photos and generated photos"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:18.353931Z",
     "iopub.execute_input": "2022-08-09T02:04:18.354626Z",
     "iopub.status.idle": "2022-08-09T02:04:22.894031Z",
     "shell.execute_reply.started": "2022-08-09T02:04:18.354571Z",
     "shell.execute_reply": "2022-08-09T02:04:22.892985Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 56,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's compile our model. Since we used tf.keras.Model to build our CycleGAN, we can just ude the fit function to train our model."
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "with strategy.scope():\n",
    "    # Create generators optimizers\n",
    "    monet_generator_optimizer = optimizers.Adam(2e-4, beta_1=0.5)\n",
    "    photo_generator_optimizer = optimizers.Adam(2e-4, beta_1=0.5)\n",
    "\n",
    "    # Create discriminators optimizers\n",
    "    monet_discriminator_optimizer = optimizers.Adam(2e-4, beta_1=0.5)\n",
    "    photo_discriminator_optimizer = optimizers.Adam(2e-4, beta_1=0.5)\n",
    "    "
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:22.897863Z",
     "iopub.execute_input": "2022-08-09T02:04:22.898238Z",
     "iopub.status.idle": "2022-08-09T02:04:22.905290Z",
     "shell.execute_reply.started": "2022-08-09T02:04:22.898203Z",
     "shell.execute_reply": "2022-08-09T02:04:22.903688Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 57,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "with strategy.scope():\n",
    "     # Create GAN\n",
    "    gan_model = CycleGan(dsaug_layer, monet_generator, photo_generator, \n",
    "                         monet_discriminator, photo_discriminator)\n",
    "\n",
    "    gan_model.compile(m_gen_optimizer=monet_generator_optimizer,\n",
    "                      p_gen_optimizer=photo_generator_optimizer,\n",
    "                      m_disc_optimizer=monet_discriminator_optimizer,\n",
    "                      p_disc_optimizer=photo_discriminator_optimizer,\n",
    "                      gen_loss_fn=generator_loss,\n",
    "                      disc_loss_fn=discriminator_loss,\n",
    "                      cycle_loss_fn=calc_cycle_loss,\n",
    "                      identity_loss_fn=identity_loss, ds_augment=None,\n",
    "                      diffaugment=\"\")"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:22.906807Z",
     "iopub.execute_input": "2022-08-09T02:04:22.907297Z",
     "iopub.status.idle": "2022-08-09T02:04:22.939901Z",
     "shell.execute_reply.started": "2022-08-09T02:04:22.907249Z",
     "shell.execute_reply": "2022-08-09T02:04:22.938909Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 58,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "history = gan_model.fit(get_gan_dataset(MONET_FILENAMES, PHOTO_FILENAMES, batch_size=BATCH_SIZE), \n",
    "                        steps_per_epoch=(max(n_monet_samples, n_photo_samples)//4),\n",
    "                        epochs=EPOCHS).history"
   ],
   "metadata": {
    "_kg_hide-output": true,
    "execution": {
     "iopub.status.busy": "2022-08-09T02:04:22.942875Z",
     "iopub.execute_input": "2022-08-09T02:04:22.943393Z",
     "iopub.status.idle": "2022-08-09T02:07:27.519730Z",
     "shell.execute_reply.started": "2022-08-09T02:04:22.943345Z",
     "shell.execute_reply": "2022-08-09T02:07:27.518217Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 59,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Prediction"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Prediction functions"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def display_generated_samples(ds, model, n_samples):\n",
    "    ds_iter = iter(ds)\n",
    "    for n_sample in range(n_samples):\n",
    "        example_sample = next(ds_iter)\n",
    "        generated_sample = model.predict(example_sample)\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()\n",
    "        \n",
    "def predict_and_save(input_ds, generator_model, output_path):\n",
    "    i = 1\n",
    "    for img in input_ds:\n",
    "        prediction = generator_model(img, training=False)[0].numpy() # make predition\n",
    "        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)   # re-scale\n",
    "        im = PIL.Image.fromarray(prediction)\n",
    "        im.save(f'{output_path}{str(i)}.jpg')\n",
    "        i += 1"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:07:47.496153Z",
     "iopub.execute_input": "2022-08-09T02:07:47.496655Z",
     "iopub.status.idle": "2022-08-09T02:07:47.507238Z",
     "shell.execute_reply.started": "2022-08-09T02:07:47.496612Z",
     "shell.execute_reply": "2022-08-09T02:07:47.506155Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 60,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Visualize predictions"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "display_generated_samples(load_dataset(PHOTO_FILENAMES).batch(1), monet_generator, 2)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:07:52.075967Z",
     "iopub.execute_input": "2022-08-09T02:07:52.077254Z",
     "iopub.status.idle": "2022-08-09T02:07:54.875107Z",
     "shell.execute_reply.started": "2022-08-09T02:07:52.077199Z",
     "shell.execute_reply": "2022-08-09T02:07:54.873403Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 61,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Create submission file"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# os.makedirs('../images/') # Create folder to save generated images\n",
    "! mkdir ../images3\n",
    "\n",
    "predict_and_save(load_dataset(PHOTO_FILENAMES).batch(1), monet_generator, '../images3/')"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:08:02.007115Z",
     "iopub.execute_input": "2022-08-09T02:08:02.007644Z",
     "iopub.status.idle": "2022-08-09T02:09:39.498148Z",
     "shell.execute_reply.started": "2022-08-09T02:08:02.007598Z",
     "shell.execute_reply": "2022-08-09T02:09:39.496258Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 62,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "shutil.make_archive('/kaggle/working/images3/', 'zip', '../images3')\n",
    "\n",
    "print(f\"Generated samples: {len([name for name in os.listdir('../images3/') if os.path.isfile(os.path.join('../images3/', name))])}\")"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-08-09T02:09:54.587192Z",
     "iopub.execute_input": "2022-08-09T02:09:54.587712Z",
     "iopub.status.idle": "2022-08-09T02:09:54.700624Z",
     "shell.execute_reply.started": "2022-08-09T02:09:54.587667Z",
     "shell.execute_reply": "2022-08-09T02:09:54.699524Z"
    },
    "trusted": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 63,
   "outputs": []
  }
 ]
}