GAN_train_2021-11-25.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "387c0065",
   "metadata": {
    "id": "387c0065"
   },
   "source": [
    "# Train an anti-aliased U-net based GAN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a29f76",
   "metadata": {
    "id": "34a29f76"
   },
   "source": [
    "## Imports and GPU identification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57ce6cac",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 5853,
     "status": "ok",
     "timestamp": 1637408482853,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "57ce6cac",
    "outputId": "af73f4f8-13c2-4f81-dfd3-e8512249d33b"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import os\n",
    "from PIL import Image\n",
    "from scipy import special\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from tensorflow.python.client import device_lib\n",
    "from datetime import datetime\n",
    "import time\n",
    "from IPython.display import clear_output\n",
    "from keras.applications.inception_v3 import InceptionV3\n",
    "from keras.applications.inception_v3 import preprocess_input\n",
    "from scipy.linalg import sqrtm\n",
    "import glob\n",
    "\n",
    "\n",
    "#%rm -rf ./logs/\n",
    "\n",
    "print(device_lib.list_local_devices())\n",
    "print(tf.config.experimental.get_memory_info(\"/device:GPU:0\")['current'])"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "id": "99069d54"
   },
   "source": [
    "## Test Keras convolution\n",
    "Find out what the Keras convolution automatic zero paddings are for padding = \"SAME\", depending on the input size and kernel size. For skip connections it is important that the data stays centered spatially. These tests need not be run for the rest of the notebook to work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8514f9ea",
   "metadata": {
    "id": "8514f9ea"
   },
   "outputs": [],
   "source": [
    "if False:\n",
    "  print(\"Downsampling, even size input, even size kernel:\")\n",
    "  print(\"Input:\")\n",
    "  input_ = tf.ones((1, 8, 8, 1), dtype=tf.float32)\n",
    "  print(input_[0, :, :, 0].numpy())\n",
    "  print(\"Kernel:\")\n",
    "  kernel = tf.ones((4, 4, 1, 1), dtype=tf.float32)\n",
    "  print(kernel[:, :, 0, 0].numpy())\n",
    "  print(\"Output of strides = 2, padding = 'SAME' convolution:\")\n",
    "  conv = tf.nn.conv2d(input = input_, filters = kernel, strides = 2, padding = 'SAME')\n",
    "  print(conv[0, :, :, 0].numpy())\n",
    "  print(\"The output is properly centered.\")\n",
    "  print(\"Output size = input size / 2\")\n",
    "\n",
    "  print(\"\\nDownsampling, odd size input, even size kernel:\")\n",
    "  print(\"Input:\")\n",
    "  input_ = tf.ones((1, 7, 7, 1), dtype=tf.float32)\n",
    "  print(input_[0, :, :, 0].numpy())\n",
    "  print(\"Kernel:\")\n",
    "  kernel = tf.ones((4, 4, 1, 1), dtype=tf.float32)\n",
    "  print(kernel[:, :, 0, 0].numpy())\n",
    "  print(\"Output of strides = 2, padding = 'SAME' convolution:\")\n",
    "  conv = tf.nn.conv2d(input = input_, filters = kernel, strides = 2, padding = 'SAME')\n",
    "  print(conv[0, :, :, 0].numpy())\n",
    "  print(\"Invalid output centering! Don't do this.\")\n",
    "\n",
    "  print(\"\\nDownsampling, even size input, odd size kernel:\")\n",
    "  print(\"Input:\")\n",
    "  input_ = tf.ones((1, 8, 8, 1), dtype=tf.float32)\n",
    "  print(input_[0, :, :, 0].numpy())\n",
    "  print(\"Kernel:\")\n",
    "  kernel = tf.ones((3, 3, 1, 1), dtype=tf.float32)\n",
    "  print(kernel[:, :, 0, 0].numpy())\n",
    "  print(\"Output of strides = 2, padding = 'SAME' convolution:\")\n",
    "  conv = tf.nn.conv2d(input = input_, filters = kernel, strides = 1, padding = 'SAME')\n",
    "  print(conv[0, :, :, 0].numpy())\n",
    "  print(\"Invalid output centering! Don't do this.\")\n",
    "\n",
    "  print(\"\\nDownsampling, odd size input, odd size kernel:\")\n",
    "  print(\"Input:\")\n",
    "  input_ = tf.ones((1, 9, 9, 1), dtype=tf.float32)\n",
    "  print(input_[0, :, :, 0].numpy())\n",
    "  print(\"Kernel:\")\n",
    "  kernel = tf.ones((5, 5, 1, 1), dtype=tf.float32)\n",
    "  print(kernel[:, :, 0, 0].numpy())\n",
    "  print(\"Output of strides = 2, padding = 'SAME' convolution:\")\n",
    "  conv = tf.nn.conv2d(input = input_, filters = kernel, strides = 1, padding = 'SAME')\n",
    "  print(conv[0, :, :, 0].numpy())\n",
    "  print(\"The output is properly centered.\")\n",
    "  print(\"Output size = (input size + 1) / 2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70e8acd7",
   "metadata": {
    "id": "70e8acd7"
   },
   "outputs": [],
   "source": [
    "if False:\n",
    "  print(\"Upsampling, any size input, even size kernel:\")\n",
    "  print(\"Input:\")\n",
    "  input_ = tf.ones((1, 4, 4, 1), dtype=tf.float32)\n",
    "  print(input_[0, :, :, 0].numpy())\n",
    "  tile = tf.constant([[0, 0], [0, 4]], dtype=tf.float32)\n",
    "  input_shape = input_.get_shape()\n",
    "  tiles = tf.reshape(tf.tile(tile, [input_shape[-3], input_shape[-2]]), (1, input_shape[-3]*2, input_shape[-2]*2, 1))\n",
    "  print(\"Diluted:\")\n",
    "  diluted = tf.keras.layers.UpSampling2D(size = (2, 2), interpolation='nearest')(input_)*tiles\n",
    "  print(diluted[0, :, :, 0].numpy());\n",
    "  print(\"Kernel:\")\n",
    "  kernel = tf.ones((4, 4, 1, 1), dtype=tf.float32)\n",
    "  print(kernel[:, :, 0, 0].numpy())\n",
    "  print(\"Output of strides = 1, padding = 'SAME' convolution:\")\n",
    "  conv = tf.nn.conv2d(input = diluted, filters = kernel, strides = 1, padding = 'SAME')\n",
    "  print(conv[0, :, :, 0].numpy())\n",
    "  print(\"The output is properly centered.\")\n",
    "  print(\"Output size = input size * 2\")\n",
    "\n",
    "  print(\"\\nUpsampling, any size input, odd size kernel:\")\n",
    "  print(\"Input:\")\n",
    "  input_ = tf.ones((1, 4, 4, 1), dtype=tf.float32)\n",
    "  print(input_[0, :, :, 0].numpy())\n",
    "  tile = tf.constant([[0, 0], [0, 4]], dtype=tf.float32)\n",
    "  input_shape = input_.get_shape()\n",
    "  tiles = tf.reshape(tf.tile(tile, [input_shape[-3], input_shape[-2]])[1:,1:], (1, input_shape[-3]*2 - 1, input_shape[-2]*2 - 1, 1))\n",
    "  print(\"Diluted:\")\n",
    "  diluted = tf.keras.layers.UpSampling2D(size = (2, 2), interpolation='nearest')(input_)[:,1:,1:,:]*tiles\n",
    "  print(diluted[0, :, :, 0].numpy());\n",
    "  print(\"Kernel:\")\n",
    "  kernel = tf.ones((3, 3, 1, 1), dtype=tf.float32)\n",
    "  print(kernel[:, :, 0, 0].numpy())\n",
    "  print(\"Output of strides = 1, padding = 'SAME' convolution:\")\n",
    "  conv = tf.nn.conv2d(input = diluted, filters = kernel, strides = 1, padding = 'SAME')\n",
    "  print(conv[0, :, :, 0].numpy())\n",
    "  print(\"The output is properly centered.\")\n",
    "  print(\"Output size = input size * 2 - 1\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f53594bd",
   "metadata": {
    "id": "f53594bd"
   },
   "source": [
    "## Circularly symmetric low-pass or anti-aliasing filter\n",
    "Functions for generating a filter kernel with a parameterized cutoff frequency and kernel size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e996b06",
   "metadata": {
    "id": "8e996b06"
   },
   "outputs": [],
   "source": [
    "def circular_lowpass_kernel(omega_c, N):  # omega = cutoff frequency in radians (pi is max), N = horizontal size of the kernel, also its vertical size.\n",
    "  with np.errstate(divide='ignore',invalid='ignore'):    \n",
    "      kernel = np.fromfunction(lambda x, y: omega_c*special.j1(omega_c*np.sqrt((x - (N - 1)/2)**2 + (y - (N - 1)/2)**2))/(2*np.pi*np.sqrt((x - (N - 1)/2)**2 + (y - (N - 1)/2)**2)), [N, N])\n",
    "  if N % 2:\n",
    "      kernel[(N - 1)//2, (N - 1)//2] = omega_c**2/(4*np.pi)\n",
    "  return kernel\n",
    "\n",
    "def rotated_cosine_window(N):  # N = horizontal size of the targeted kernel, also its vertical size, must be odd.\n",
    "  return np.fromfunction(lambda y, x: np.maximum(np.cos(np.pi/2*np.sqrt(((x - (N - 1)/2)/((N - 1)/2 + 1))**2 + ((y - (N - 1)/2)/((N - 1)/2 + 1))**2)), 0), [N, N])\n",
    "\n",
    "def windowed_circular_lowpass_kernel(cutoff, kernel_size):\n",
    "  return (circular_lowpass_kernel(cutoff, kernel_size)*rotated_cosine_window(kernel_size)).astype(np.float32);\n",
    "\n",
    "def test_kernel(kernel):\n",
    "  kernel_size = kernel.shape[0] #height, should be same as width\n",
    "\n",
    "  print(\"Kernel:\")\n",
    "  plt.imshow(kernel, vmin=-0.25, vmax=0.25, cmap='bwr', extent=[-kernel_size/2, kernel_size/2, -kernel_size/2, kernel_size/2])\n",
    "  plt.colorbar()\n",
    "  plt.show()\n",
    "\n",
    "  print(\"Frequency response:\")\n",
    "  freq_resp = np.abs(np.fft.fftshift(np.fft.fft2(np.roll(np.pad(kernel, 101, mode = 'constant', constant_values = 0), shift = (-kernel_size//2-99, -kernel_size//2-99), axis = (0, 1)))))\n",
    "  plt.imshow(freq_resp, vmin=-1.2, vmax=1.2, cmap='bwr', extent=[-np.pi, np.pi, -np.pi, np.pi])\n",
    "  plt.colorbar()\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28c89670",
   "metadata": {
    "id": "28c89670"
   },
   "source": [
    "## Load training target images and low-pass filter them\n",
    "\n",
    "Note: Do NOT use the filtered PNG images as network input, because they have been quantized. They can be saved for visual inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccebb777",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1030,
     "status": "error",
     "timestamp": 1637410885981,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "ccebb777",
    "outputId": "86844914-76c9-4aa3-d0b9-67d432e131fe"
   },
   "outputs": [],
   "source": [
    "# Configuration\n",
    "oversampling = 2  # Oversampling factor, 2 is a good choice\n",
    "aa_even_kernel_size = 8 # Size of even sized antialiasing filter\n",
    "aa_odd_kernel_size = aa_even_kernel_size + 1 # Size of odd sized antialiasing filter\n",
    "approx_batch_size = 4\n",
    "train_image_cropped_size = 512\n",
    "photos_dir = \"source_grass_images_subset\"\n",
    "save_filtered_photos_dir = None\n",
    "#save_filtered_photos_dir = \"filtered_grass_images_subset\" # save folder, or None to disable saving\n",
    "\n",
    "print(\"Even size kernel:\")\n",
    "test_kernel(windowed_circular_lowpass_kernel(np.pi/oversampling, aa_even_kernel_size))\n",
    "print(\"Odd size kernel:\")\n",
    "test_kernel(windowed_circular_lowpass_kernel(np.pi/oversampling, aa_odd_kernel_size))\n",
    "\n",
    "def channeled_kernel(kernel_prototype, num_channels):\n",
    "  kernel_expanded = tf.expand_dims(kernel_prototype, 2)\n",
    "  kernel_repeated = tf.repeat(kernel_expanded, num_channels, axis=2)\n",
    "  return tf.expand_dims(kernel_repeated, -1)\n",
    "\n",
    "def load_and_filter_image(filepath, channeled_kernel):\n",
    "  image = tf.io.decode_png(tf.io.read_file(filepath), channels=3, dtype=tf.dtypes.uint8)\n",
    "  image = tf.image.convert_image_dtype(image, dtype=tf.float32) * 2 - 1\n",
    "  image = tf.nn.depthwise_conv2d([image], channeled_kernel, [1, 1, 1, 1], \"VALID\")[0]\n",
    "  if save_filtered_photos_dir is not None:\n",
    "    basename = tf.py_function(lambda string_tensor: os.path.basename(string_tensor.numpy()), [filepath], tf.string)\n",
    "    tf.io.write_file(save_filtered_photos_dir + \"/\" + basename, tf.io.encode_png(tf.image.convert_image_dtype((image + 1)/2, dtype=tf.dtypes.uint8, saturate=True)))\n",
    "  return image\n",
    "\n",
    "train_dataset = tf.data.Dataset.list_files(\"test.png\")\n",
    "train_dataset = train_dataset.map(lambda filepath: load_and_filter_image(filepath, channeled_kernel(windowed_circular_lowpass_kernel(np.pi/oversampling, aa_even_kernel_size), 3)), num_parallel_calls=tf.data.experimental.AUTOTUNE, deterministic=False)\n",
    "\n",
    "train_dataset = train_dataset.cache() # Load photos now instead of doing it again for each copy made in the next line\n",
    "train_dataset = train_dataset.repeat(tf.math.floordiv(approx_batch_size, train_dataset.cardinality()))\n",
    "train_dataset = train_dataset.cache()\n",
    "\n",
    "batch_size = train_dataset.cardinality() # batch size\n",
    "print(\"Batch size = \" + str(batch_size.numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eba26b3d",
   "metadata": {
    "id": "eba26b3d"
   },
   "source": [
    "## Isotropic lowpass filtering layers\n",
    "These layers apply a circularly symmmetrical lowpass filter. They preserve spatial centering of the data to facilitate skip connections and center cropping in the network.\n",
    "\n",
    "`IsotropicLowpassFilterLayer` accepts `strides = 1` to preserve the feature map size and `strides = 2` to downsample it to half size. With `strides = 1`, the kernel size must be odd. With downsampling, only even input and kernel sizes are supported, resulting in even-sized output. `cutoff` is the cutoff frequency (before downsampling).\n",
    "\n",
    "`IsotropicLowpassFilteredUpsampleLayer` doubles the spatial size of the feature map. `cutoff` is the cutoff frequency after upsampling. `kernel_size` must be even."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a334bf1",
   "metadata": {
    "id": "7a334bf1"
   },
   "outputs": [],
   "source": [
    "class IsotropicLowpassFilterLayer(tf.keras.layers.Layer):\n",
    "  def __init__(self, cutoff, kernel_size, strides = 1, **kwargs):\n",
    "    super(IsotropicLowpassFilterLayer, self).__init__(**kwargs)\n",
    "    self.kernel_size = kernel_size\n",
    "    self.strides = strides        \n",
    "    assert self.strides == 1 or self.strides == 2, \"Accepted values for strides are 1 and 2. Tuples or arrays are not accepted.\"\n",
    "    if self.strides == 1:\n",
    "      assert self.kernel_size % 2 == 1, \"kernel_size must be odd when strides = 1\"\n",
    "    if self.strides == 2:\n",
    "      assert self.kernel_size % 2 == 0, \"kernel_size must be even when strides = 2\"\n",
    "\n",
    "    self.kernel_prototype = windowed_circular_lowpass_kernel(cutoff, self.kernel_size)\n",
    "\n",
    "  def build(self, input_shape):\n",
    "    self.filtering = tf.keras.layers.DepthwiseConv2D(self.kernel_size, trainable=False, use_bias=False, padding='same', strides=self.strides, weights=[channeled_kernel(self.kernel_prototype, input_shape[-1])])\n",
    "\n",
    "  def call(self, inputs):\n",
    "    return self.filtering(inputs)\n",
    "\n",
    "def get_upsample_multiplication_tensor(input_shape):\n",
    "  tile = tf.constant([[0, 0], [0, 4]], dtype=tf.float32)\n",
    "  return tf.reshape(tf.tile(tile, [input_shape[-3], input_shape[-2]]), (1, input_shape[-3]*2, input_shape[-2]*2, 1))\n",
    "\n",
    "class IsotropicLowpassFilteredUpsampleLayer(tf.keras.layers.Layer):\n",
    "  def __init__(self, cutoff, kernel_size, **kwargs):\n",
    "    super(IsotropicLowpassFilteredUpsampleLayer, self).__init__(**kwargs)\n",
    "    self.kernel_size = kernel_size\n",
    "    assert self.kernel_size % 2 == 0, \"kernel_size must be even\"\n",
    "    self.upsampling = tf.keras.layers.UpSampling2D(size = (2, 2), interpolation='nearest')\n",
    "    self.multiply = tf.keras.layers.Multiply()\n",
    "    #self.zeropadding = tf.keras.layers.ZeroPadding2D(padding = ((0, 1), (0, 1)))  # not needed for input even -> output even\n",
    "    self.kernel_prototype = windowed_circular_lowpass_kernel(cutoff, kernel_size)\n",
    "    self.kernel_size = kernel_size\n",
    "\n",
    "  def build(self, input_shape):\n",
    "    self.multiplication_tensor = get_upsample_multiplication_tensor(input_shape)\n",
    "    self.filtering = tf.keras.layers.DepthwiseConv2D(self.kernel_size, trainable=False, use_bias=False, padding='same', weights=[channeled_kernel(self.kernel_prototype, input_shape[-1])])\n",
    "\n",
    "  def call(self, inputs):\n",
    "    x = self.upsampling(inputs)\n",
    "    x = self.multiply([x, self.multiplication_tensor])\n",
    "    #x = self.zeropadding(x)  # not needed for input even -> output even\n",
    "    x = self.filtering(x)\n",
    "    return x    \n",
    "\n",
    "if False:      \n",
    "  print(\"Test input:\")\n",
    "  input_shapes = (aa_odd_kernel_size + 2, aa_odd_kernel_size + 2, 3)\n",
    "  eye = tf.repeat(tf.expand_dims(tf.repeat(tf.expand_dims(tf.eye(input_shapes[0], input_shapes[1]), -1), repeats = input_shapes[2], axis = -1), 0), repeats = batch_size, axis = 0)\n",
    "  plt.imshow(eye[0])\n",
    "  plt.show()\n",
    "\n",
    "  def TestModel(input_shapes):\n",
    "    inputs = tf.keras.Input(input_shapes)\n",
    "    outputs = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size)(inputs)\n",
    "    model = tf.keras.Model(inputs=inputs, outputs=outputs)    \n",
    "    return model\n",
    "\n",
    "  print(\"Isotropic lowpass filtered:\")\n",
    "  model = TestModel(input_shapes)\n",
    "  output = model(eye)\n",
    "  plt.imshow(output[0])\n",
    "  plt.show()\n",
    "\n",
    "  def TestModel(input_shapes):\n",
    "    inputs = tf.keras.Input(input_shapes)\n",
    "    outputs = IsotropicLowpassFilteredUpsampleLayer(np.pi/oversampling, aa_even_kernel_size)(inputs)\n",
    "    model = tf.keras.Model(inputs=inputs, outputs=outputs)    \n",
    "    return model\n",
    "\n",
    "  print(\"Isotropic lowpass filtered upsampled:\")\n",
    "  model = TestModel(input_shapes)\n",
    "  output = model(eye)\n",
    "  plt.imshow(output[0])\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22f73bdd",
   "metadata": {
    "id": "22f73bdd"
   },
   "source": [
    "## RandomCropper\n",
    "Crop images randomly. This should be done in the beginning of each epoch by calling RandomCropper.recrop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4127d6fb",
   "metadata": {
    "id": "4127d6fb"
   },
   "outputs": [],
   "source": [
    "class RandomCropper:\n",
    "  def __init__(self, dataset, first_seed = None):\n",
    "    if first_seed is None:\n",
    "      first_seed = 1\n",
    "      #first_seed = int(str(int((datetime.now() - datetime(1970, 1, 1)).total_seconds()*1_000_000))[::-1][:16])\n",
    "    self.dataset_seeded = dataset.enumerate(first_seed)\n",
    "\n",
    "  def recrop(self, size):\n",
    "    cropped_dataset = self.dataset_seeded.map(lambda seed, image: tf.image.stateless_random_crop(image, (size, size, 3), [seed, 0]), num_parallel_calls=tf.data.experimental.AUTOTUNE, deterministic=False)\n",
    "    self.dataset_seeded = self.dataset_seeded.map(lambda seed, image: (seed + self.dataset_seeded.cardinality(), image), num_parallel_calls=tf.data.experimental.AUTOTUNE, deterministic=False)\n",
    "    return cropped_dataset\n",
    "\n",
    "if False:\n",
    "  train_dataset_cropper = RandomCropper(train_dataset)\n",
    "  print(\"First recrop, first item:\")\n",
    "  iterator = iter(train_dataset_cropper.recrop(train_image_cropped_size))\n",
    "  plt.imshow((next(iterator) + 1)/2)\n",
    "  plt.show()\n",
    "  print(\"Second recrop, first item:\")\n",
    "  iterator = iter(train_dataset_cropper.recrop(train_image_cropped_size))\n",
    "  plt.imshow((next(iterator) + 1)/2)\n",
    "  plt.show()\n",
    "  print(\"Second recrop, second item:\")\n",
    "  plt.imshow((next(iterator) + 1)/2)\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb309349",
   "metadata": {
    "id": "cb309349"
   },
   "source": [
    "## Generator model\n",
    "Define the generator network, parameterized by input image size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c547c65",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 9838,
     "status": "ok",
     "timestamp": 1637393243658,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "8c547c65",
    "outputId": "ad2dc5cd-4fc8-4fa5-b612-7dd9aa039c0f"
   },
   "outputs": [],
   "source": [
    "def Generator(input_size, down_nonlinearity, up_nonlinearity, initializer, learned_kernel_size, num_sublayers, num_resamplings, max_features, use_skip_connections):      \n",
    "  inputs = tf.keras.Input(shape = (input_size, input_size, 1), name = \"input\")\n",
    "\n",
    "  x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size, name = \"input_antialias\")(inputs)\n",
    "\n",
    "  skip_connection_stack = []\n",
    "  skip_connection_stack.append(x)\n",
    "\n",
    "  n = 64\n",
    "  for j in range(num_resamplings):\n",
    "    for i in range(num_sublayers):\n",
    "      x = tf.keras.layers.Conv2D(max_features if n > max_features else n, learned_kernel_size, kernel_initializer=initializer, use_bias = False, padding=\"same\", name = f\"down_{j}_subblock_{i}_learned_conv\")(x)\n",
    "      if j > 0:\n",
    "        x = tf.keras.layers.BatchNormalization(name = f\"down_{j}_subblock_{i}_batchnorm\")(x) \n",
    "      x = tf.keras.layers.Activation(down_nonlinearity, name = f\"down_{j}_subblock_{i}_nonlinearity\")(x)\n",
    "      x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size if i < num_sublayers - 1 else aa_even_kernel_size, strides = 1 if i < num_sublayers - 1 else 2, name = f\"down_{j}_subblock_{i}_antialias\" if i < num_sublayers - 1 else f\"down_{j}_subblock_{i}_antialiased_downsample\")(x)\n",
    "    x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size, name = f\"down_{j}_antialias\")(x)\n",
    "    if j < num_resamplings - 1:\n",
    "      skip_connection_stack.append(x)\n",
    "    n *= 2\n",
    "\n",
    "  for j in range(num_resamplings):\n",
    "    n //= 2\n",
    "    x = IsotropicLowpassFilteredUpsampleLayer(np.pi/oversampling, aa_even_kernel_size, name = f\"up_{j}_antialiased_upsample\")(x)\n",
    "    if use_skip_connections:\n",
    "      skip_connection = skip_connection_stack.pop()\n",
    "      x = tf.keras.layers.Concatenate()([x, skip_connection])\n",
    "    #x = tf.keras.layers.Multiply(x[..., 0:x.get_shape()[-1]/4])(x.get_shape()[-1]/4\n",
    "    for i in range(num_sublayers):\n",
    "      x = tf.keras.layers.Conv2D(max_features if n > max_features else n, learned_kernel_size, kernel_initializer=initializer, use_bias = False, padding=\"same\", name = f\"up_{j}_subblock_{i}_learned_conv\")(x)\n",
    "      if j < num_resamplings - 1:\n",
    "        x = tf.keras.layers.BatchNormalization(name = f\"up_{j}_subblock_{i}_batchnorm\")(x)\n",
    "        #x = tf.keras.layers.Dropout(0.5)(x)\n",
    "      x = tf.keras.layers.Activation(up_nonlinearity, name = f\"up_{j}_subblock_{i}_nonlinearity\")(x)\n",
    "      x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size, name = f\"up_{j}_subblock_{i}_antialias\")(x)\n",
    "\n",
    "  x = tf.keras.layers.Conv2D(3, learned_kernel_size, kernel_initializer=initializer, use_bias=False, padding=\"same\", name = \"to_rgb_learned_conv\")(x)\n",
    "  x = tf.keras.layers.Activation(\"tanh\", name = \"output_nonlinearity\")(x)  \n",
    "  x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size, name = \"output_antialias\")(x)\n",
    "  #x = tf.keras.layers.CenterCrop(x.get_shape()[-3]//2, x.get_shape()[-2]//2, name = \"output_crop\")(x)\n",
    "\n",
    "  return  tf.keras.Model(inputs=inputs, outputs=x)\n",
    "\n",
    "generator_input_size = train_image_cropped_size\n",
    "\n",
    "generator = Generator(generator_input_size, tf.keras.layers.LeakyReLU(), tf.keras.layers.LeakyReLU(), \"glorot_uniform\", learned_kernel_size = 5, num_sublayers = 3, num_resamplings = 7, max_features = 256, use_skip_connections = True)\n",
    "generator.build((generator_input_size, generator_input_size, 1))\n",
    "\n",
    "print(\"Generator network:\")\n",
    "display(tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1OZoCCAPYU6b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 556
    },
    "executionInfo": {
     "elapsed": 18016,
     "status": "ok",
     "timestamp": 1637393290774,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "1OZoCCAPYU6b",
    "outputId": "8d16ff56-bb04-4c18-a85f-2f5181e84299"
   },
   "outputs": [],
   "source": [
    "# Generate, display\n",
    "z = tf.random.uniform((1, generator_input_size, generator_input_size, 1), minval=-1, maxval=1)\n",
    "print(\"Generator input:\")\n",
    "plt.imshow(np.repeat((z[0].numpy() + 1)/2, repeats = 3, axis = -1))\n",
    "plt.show()\n",
    "print(\"Generator output:\")\n",
    "x = generator(z, training = True)[0]\n",
    "plt.imshow((x.numpy() + 1)/2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a55b8c",
   "metadata": {
    "id": "c1a55b8c"
   },
   "source": [
    "## Generator loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3a25bfa",
   "metadata": {
    "id": "c3a25bfa"
   },
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
    "def generator_loss(disc_generated_output, gen_output):\n",
    "    gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output)\n",
    "    return gan_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1f5e861",
   "metadata": {
    "id": "d1f5e861"
   },
   "source": [
    "## Discriminator model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aa49568",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 8743,
     "status": "ok",
     "timestamp": 1637393334298,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "5aa49568",
    "outputId": "9be19d8c-a79a-46a8-afce-f6437302eff4"
   },
   "outputs": [],
   "source": [
    "def Discriminator(input_size, down_nonlinearity, initializer, learned_kernel_size, num_sublayers, num_resamplings, max_features):      \n",
    "    inputs = tf.keras.Input(shape = (input_size, input_size, 3), name = \"input\")\n",
    "  \n",
    "    x = inputs\n",
    "  \n",
    "    n = 64\n",
    "    for j in range(num_resamplings):\n",
    "        for i in range(num_sublayers):\n",
    "            x = tf.keras.layers.Conv2D(max_features if n > max_features else n, learned_kernel_size, kernel_initializer=initializer, use_bias=True, padding=\"same\", name = f\"down_{j}_subblock_{i}_learned_conv\")(x)\n",
    "            if j > 0:\n",
    "                x = tf.keras.layers.BatchNormalization(name = f\"down_{j}_subblock_{i}_batchnorm\")(x)             \n",
    "            x = tf.keras.layers.Activation(down_nonlinearity, name = f\"down_{j}_subblock_{i}_nonlinearity\")(x)\n",
    "            x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size if i < num_sublayers - 1 else aa_even_kernel_size, strides = 1 if i < num_sublayers - 1 else 2, name = f\"down_{j}_subblock_{i}_antialias\" if i < num_sublayers - 1 else f\"down_{j}_subblock_{i}_antialiased_downsample\")(x)\n",
    "        x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size, name = f\"down_{j}_antialias\")(x)\n",
    "        n *= 2\n",
    "\n",
    "    x = tf.keras.layers.Conv2D(1, learned_kernel_size, kernel_initializer=initializer, use_bias=True, padding=\"same\", name = \"to_output_learned_conv\")(x)\n",
    "    x = IsotropicLowpassFilterLayer(np.pi/oversampling, aa_odd_kernel_size, name = \"output_antialias\")(x)\n",
    "\n",
    "    return  tf.keras.Model(inputs=inputs, outputs=x)\n",
    "\n",
    "discriminator = Discriminator(train_image_cropped_size, tf.keras.layers.LeakyReLU(), \"glorot_uniform\", learned_kernel_size = 5, num_sublayers = 3, num_resamplings = 6, max_features = 256)\n",
    "discriminator.build((train_image_cropped_size, train_image_cropped_size, 1))\n",
    "\n",
    "print(\"Discriminator network:\")\n",
    "display(tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52zQsH9DDqFp",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 11429,
     "status": "ok",
     "timestamp": 1637393388006,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "52zQsH9DDqFp",
    "outputId": "c1333353-d886-4f0c-f88c-497ac6391a53"
   },
   "outputs": [],
   "source": [
    "z = tf.random.uniform((1, generator_input_size, generator_input_size, 1), minval=-1, maxval=1)\n",
    "\n",
    "print(\"Sample training target image:\")\n",
    "train_dataset_cropper = RandomCropper(train_dataset)\n",
    "iterator = iter(train_dataset_cropper.recrop(train_image_cropped_size))\n",
    "train_image = next(iterator)\n",
    "plt.imshow((train_image + 1)/2)\n",
    "plt.show()\n",
    "\n",
    "print(\"Discriminator output:\")\n",
    "disc_out = discriminator(train_image[tf.newaxis, ...], training=True)\n",
    "plt.imshow(disc_out[0, ..., -1], cmap='RdBu_r')\n",
    "plt.clim(-10,10)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "print(\"Sample generated image, before training:\")\n",
    "generated_image = generator(z, training = True)[0]\n",
    "plt.imshow((generated_image.numpy() + 1)/2)\n",
    "plt.show()\n",
    "\n",
    "print(\"Discriminator output:\")\n",
    "disc_out = discriminator(generated_image[tf.newaxis, ...], training=True)\n",
    "plt.imshow(disc_out[0, ..., -1], cmap='RdBu_r')\n",
    "plt.clim(-10,10)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a061b3c",
   "metadata": {
    "id": "6a061b3c"
   },
   "source": [
    "## Discriminator loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c0be09",
   "metadata": {
    "id": "10c0be09"
   },
   "outputs": [],
   "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",
    "  generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output)\n",
    "  total_disc_loss = real_loss + generated_loss\n",
    "  return total_disc_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9689d9",
   "metadata": {
    "id": "7d9689d9"
   },
   "source": [
    "## Optimizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80d9eb7a",
   "metadata": {
    "id": "80d9eb7a"
   },
   "outputs": [],
   "source": [
    "# tf.keras.optimizers.Adam(learning_rate=0.01, beta_1=0.5, beta_2=0.99, epsilon=1e-07)\n",
    "#generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
    "#discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
    "generator_optimizer = tf.keras.optimizers.SGD(1e-3) # start with 1e-3, then change it to 1e-4\n",
    "discriminator_optimizer = tf.keras.optimizers.SGD(1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7lsBlQDJyF3",
   "metadata": {
    "id": "d7lsBlQDJyF3"
   },
   "source": [
    "## Set up model checkpoint loading or saving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "_2T66-x4Ju2o",
   "metadata": {
    "id": "_2T66-x4Ju2o"
   },
   "outputs": [],
   "source": [
    "checkpoint_dir = './training_checkpoints_gan'\n",
    "checkpoint_prefix_latest = os.path.join(checkpoint_dir, \"ckpt_latest\")\n",
    "checkpoint_prefix_best = os.path.join(checkpoint_dir, \"ckpt_best\")\n",
    "checkpoint = tf.train.Checkpoint(generator=generator, discriminator=discriminator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "GQwxXXH-KNdT",
   "metadata": {
    "id": "GQwxXXH-KNdT"
   },
   "source": [
    "## Save checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "PVUwD7tQKQPN",
   "metadata": {
    "id": "PVUwD7tQKQPN"
   },
   "outputs": [],
   "source": [
    "#checkpoint.save(file_prefix=checkpoint_prefix)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f11bb4ca",
   "metadata": {
    "id": "f11bb4ca"
   },
   "source": [
    "## Train the GAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_red(images1, images2, channel):\n",
    "    # calculate activations\n",
    "    # calculate mean and covariance statistics\n",
    "    mu1, sigma1 = np.mean(images1[:, :, channel], axis=0), np.cov(images1[:, :, channel], rowvar=False)\n",
    "    mu2, sigma2 = np.mean(images2[:, :, channel], axis=0), np.cov(images2[:, :, channel], rowvar=False)\n",
    "    # calculate sum squared difference between means\n",
    "    ssdiff = np.sum((mu1 - mu2)**2.0)\n",
    "    # calculate sqrt of product between cov\n",
    "    covmean = sqrtm(sigma1.dot(sigma2))\n",
    "    # check and correct imaginary numbers from sqrt\n",
    "    if np.iscomplexobj(covmean):\n",
    "        covmean = covmean.real\n",
    "    # calculate score\n",
    "    red = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)\n",
    "    return red"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fid between images1 and images1\n",
    "red = calculate_red(train_image, train_image, 0)\n",
    "print('RED (same): %.3f' % red)\n",
    "# fid between images1 and images2\n",
    "red = calculate_red(train_image, generated_image, 0)\n",
    "print('RED (different): %.3f' % red)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_red(train_image, generated_image, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_fid_step1(model, repeats, generator):\n",
    "    \n",
    "    act1 = []\n",
    "    act2 = []\n",
    "    \n",
    "    train_dataset_cropper = RandomCropper(train_dataset)\n",
    "    cropped_batch = tf.stack(list(train_dataset_cropper.recrop(train_image_cropped_size)))\n",
    "    \n",
    "    for i in range(repeats):\n",
    "        # calculate activations\n",
    "        z = tf.random.uniform((batch_size, generator_input_size, generator_input_size, 1), minval=0, maxval=None)\n",
    "        generated_images = generator(z, training = True)\n",
    "        act1.append(model.predict(generated_images))\n",
    "        act2.append(model.predict(cropped_batch))\n",
    "    \n",
    "    act1 = np.array(act1)\n",
    "    act2 = np.array(act2)\n",
    "    \n",
    "    act1 = np.reshape(act1, (4 * repeats, 2048))\n",
    "    act2 = np.reshape(act2, (4 * repeats, 2048))\n",
    "    \n",
    "    return act1, act2\n",
    "\n",
    "def calculate_fid_step2(acts1, acts2):\n",
    "    # calculate mean and covariance statistics\n",
    "    mu1, sigma1 = acts1.mean(axis=0), np.cov(acts1, rowvar=False)\n",
    "    mu2, sigma2 = acts2.mean(axis=0), np.cov(acts2, rowvar=False)\n",
    "    # calculate sum squared difference between means\n",
    "    ssdiff = np.sum((mu1 - mu2)**2.0)\n",
    "    # calculate sqrt of product between cov\n",
    "    covmean = sqrtm(sigma1.dot(sigma2))\n",
    "    # check and correct imaginary numbers from sqrt\n",
    "    if np.iscomplexobj(covmean):\n",
    "        covmean = covmean.real\n",
    "    # calculate score\n",
    "    fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)\n",
    "    return fid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_ckpt(latest = True, number = 1):\n",
    "  if latest:\n",
    "    checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))\n",
    "  else:\n",
    "    checkpoint.restore('training_checkpoints_testlatest3/ckpt-' + str(number))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_ckpt()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "output_embedded_package_id": "1P7huFJ75YjFiJIKRCURlmwGt47SLa34D"
    },
    "executionInfo": {
     "elapsed": 6683748,
     "status": "error",
     "timestamp": 1637400129812,
     "user": {
      "displayName": "Olli Niemitalo",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GgS9Yyd7FywbLGk607s8oFV4Ic3akt_JuveOeZCeg=s64",
      "userId": "16659099300963047240"
     },
     "user_tz": -120
    },
    "id": "96aac46d",
    "outputId": "3dcecae5-5280-42b2-9b0d-9d897efbcec1"
   },
   "outputs": [],
   "source": [
    "epochs = 10000\n",
    "display_every = 5 # How many epochs to skip between data displays\n",
    "\n",
    "#checkpoint.save(file_prefix=checkpoint_prefix)\n",
    "#checkpoint.restore(file_prefix=checkpoint_prefix)\n",
    "\n",
    "def crop_output(thing):\n",
    "    cropper = tf.keras.layers.CenterCrop(thing.get_shape()[-3]//4, thing.get_shape()[-2]//4, name = \"discriminator_output_crop\")\n",
    "    return cropper\n",
    "\n",
    "@tf.function\n",
    "def train_step(cropped_batch, z, epoch):\n",
    "  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
    "    gen_output = generator(z, training=True)\n",
    "\n",
    "    disc_real_output = discriminator(cropped_batch, training=True)\n",
    "    disc_output_cropper = crop_output(disc_real_output)\n",
    "    disc_real_output = disc_output_cropper(disc_real_output)\n",
    "    disc_generated_output = discriminator(gen_output, training=True)\n",
    "    disc_generated_output = disc_output_cropper(disc_generated_output)\n",
    "\n",
    "    gen_gan_loss = generator_loss(disc_generated_output, gen_output)\n",
    "    disc_loss = discriminator_loss(disc_real_output, disc_generated_output)\n",
    "\n",
    "    generator_gradients = gen_tape.gradient(gen_gan_loss, generator.trainable_variables)\n",
    "    discriminator_gradients = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n",
    "    \n",
    "    generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))\n",
    "    discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables))\n",
    "\n",
    "  with summary_writer.as_default():\n",
    "    tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch)\n",
    "    tf.summary.scalar('disc_loss', disc_loss, step=epoch)\n",
    "\n",
    "def checkpoint_save(epoch, mode, inEpochs = 100):\n",
    "    if mode == \"every\":\n",
    "        if (epoch + 1) % inEpochs == 0:\n",
    "            checkpoint.save(file_prefix=checkpoint_prefix_latest)\n",
    "    elif mode == \"fid\":\n",
    "        checkpoint.save(file_prefix=checkpoint_prefix_best)\n",
    "    \n",
    "def plot_batches(subject, disc_out):\n",
    "    plt.subplots(batch_size.numpy()//2, 4, figsize = (28,7*(batch_size.numpy()//2)))\n",
    "    for i in range(batch_size):\n",
    "        plt.subplot(batch_size.numpy()//2, 4, 1 + i*2)\n",
    "        #cropped_target_image = tf.keras.layers.CenterCrop(cropped_batch.get_shape()[-3]//2, cropped_batch.get_shape()[-2]//2)(cropped_batch[0])\n",
    "        plt.imshow(np.clip((subject[i].numpy() + 1)/2, 0, 1))\n",
    "        plt.subplot(batch_size.numpy()//2, 4, 2 + i*2)\n",
    "        plt.imshow(disc_out[i][..., -1], cmap='RdBu_r')\n",
    "        plt.clim(-10,10)\n",
    "    plt.show()\n",
    "    \n",
    "def image_save(image, epoch):\n",
    "    name = str(epoch).zfill(4)\n",
    "    tf.keras.utils.save_img((\"pictures/\" + name + \".png\"), image[0])\n",
    "    \n",
    "def fit(epochs, dataset):\n",
    "  train_dataset_cropper = RandomCropper(train_dataset)\n",
    "  red_list = []\n",
    "  fid_list = []\n",
    "  const_z = tf.random.uniform((batch_size, generator_input_size, generator_input_size, 1), minval=0, maxval=None)\n",
    "    \n",
    "  for epoch in range(epochs + 1):\n",
    "    z = tf.random.uniform((batch_size, generator_input_size, generator_input_size, 1), minval=0, maxval=None)    \n",
    "    cropped_batch = tf.stack(list(train_dataset_cropper.recrop(train_image_cropped_size)))\n",
    "    # Train\n",
    "    start = time.time()\n",
    "    clear_output(wait=True)\n",
    "\n",
    "    if epoch % display_every == 0:\n",
    "      print(\"Target images and discriminator outputs (only 2x2 center is used in loss):\")\n",
    "      disc_out = discriminator(cropped_batch, training=True)\n",
    "      plot_batches(cropped_batch, disc_out)\n",
    "\n",
    "      print(\"Generated images and discriminator outputs (only 2x2 center is used in loss):\")\n",
    "      generated_images = generator(z, training = True)\n",
    "      disc_out = discriminator(generated_images, training=True)\n",
    "      #cropped_generated_image = tf.keras.layers.CenterCrop(generated_image.get_shape()[-3]//2, generated_image.get_shape()[-2]//2)(generated_image)\n",
    "      plot_batches(generated_images, disc_out)\n",
    "        \n",
    "      generated_images_4save = generator(const_z, training = True)\n",
    "      image_save(generated_images_4save, epoch)\n",
    "    \n",
    "    model = InceptionV3(include_top=False, pooling='avg')\n",
    "\n",
    "    act1, act2 = calculate_fid_step1(model, 16, generator)\n",
    "    latest_fid = calculate_fid_step2(act1, act2)\n",
    "    \n",
    "    fid_list.append(latest_fid)\n",
    "    \n",
    "    latest_red = np.mean([\n",
    "        (calculate_red(np.clip((cropped_batch[0].numpy() + 1)/2, 0, 1), np.clip((generated_images[0].numpy() + 1)/2, 0, 1), 0)),\n",
    "         (calculate_red(np.clip((cropped_batch[0].numpy() + 1)/2, 0, 1), np.clip((generated_images[0].numpy() + 1)/2, 0, 1), 1)),\n",
    "          (calculate_red(np.clip((cropped_batch[0].numpy() + 1)/2, 0, 1), np.clip((generated_images[0].numpy() + 1)/2, 0, 1), 2))\n",
    "    ])\n",
    "                \n",
    "    red_list.append(latest_red)\n",
    "    \n",
    "    print('Latest RED: %.3f' % latest_red)\n",
    "    plt.plot(range(len(red_list)), red_list)\n",
    "    plt.plot(range(len(fid_list)), fid_list)\n",
    "    plt.show()\n",
    "\n",
    "    train_step(cropped_batch, z, epoch)\n",
    "    \n",
    "    if epoch > 0 and latest_fid < minimal_fid * 0.99:\n",
    "        minimal_fid = latest_fid \n",
    "        checkpoint_save(epoch, mode = \"fid\")\n",
    "    elif epoch == 0:\n",
    "        minimal_fid = latest_fid\n",
    "    checkpoint_save(epoch, mode = \"every\")\n",
    "    \n",
    "    print('Latest FID: %.3f' % latest_fid)\n",
    "    print('Minimal FID: %.3f' % minimal_fid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "652215f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "log_dir = \"logs/\"\n",
    "im_dir = \"images/fit\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\") + \"/\"\n",
    "os.makedirs(im_dir)\n",
    "summary_writer = tf.summary.create_file_writer(\n",
    "  log_dir + \"fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60343dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext tensorboard\n",
    "#%tensorboard --logdir {log_dir} --host localhost --port 6213 #Windows\n",
    "%tensorboard --logdir {log_dir} --bind_all #Docker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296d4107",
   "metadata": {},
   "outputs": [],
   "source": [
    "fit(epochs, train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_files = os.listdir(\"pictures/\")\n",
    "sorted_files.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for picture in sorted_files[1:]:\n",
    "    display(Image.open(\"pictures/\" + picture))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "GAN_train_2021-11-25.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}