Celebrity_faces_generation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "mount_file_id": "1SIs-BjREmIc2u4_60W_uMWtVBLU-BU7a",
      "authorship_tag": "ABX9TyOZYFQAPw0p3Tls+SYUg2V7",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/karthiknadar1204/generative_deep_learning/blob/main/Celebrity_faces_generation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kX5iLcLYx00o"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from google.colab import files\n",
        "\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow.keras.backend as K\n",
        "from tensorflow.keras import (\n",
        "    layers,\n",
        "    models,\n",
        "    callbacks,\n",
        "    utils,\n",
        "    metrics,\n",
        "    losses,\n",
        "    optimizers,\n",
        ")\n",
        "\n",
        "from scipy.stats import norm\n",
        "import pandas as pd\n",
        "import zipfile\n",
        "import os\n",
        "\n",
        "# from notebooks.utils import sample_batch, display\n",
        "\n",
        "# from vae_utils import get_vector_from_label, add_vector_to_images, morph_faces"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "IMAGE_SIZE = 32\n",
        "CHANNELS = 3\n",
        "BATCH_SIZE = 128\n",
        "NUM_FEATURES = 128\n",
        "Z_DIM = 200\n",
        "LEARNING_RATE = 0.0005\n",
        "EPOCHS = 10\n",
        "BETA = 2000\n",
        "LOAD_MODEL = False"
      ],
      "metadata": {
        "id": "Hbf8vpGWx3UO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "uploaded = files.upload()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 74
        },
        "id": "VPg1i65REvNl",
        "outputId": "fe66942d-4b40-4875-c2e8-f03cc249d7d1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-5b32b8c4-cb2a-47bf-a2a1-035b16c9d373\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-5b32b8c4-cb2a-47bf-a2a1-035b16c9d373\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script>// Copyright 2017 Google LLC\n",
              "//\n",
              "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
              "// you may not use this file except in compliance with the License.\n",
              "// You may obtain a copy of the License at\n",
              "//\n",
              "//      http://www.apache.org/licenses/LICENSE-2.0\n",
              "//\n",
              "// Unless required by applicable law or agreed to in writing, software\n",
              "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
              "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
              "// See the License for the specific language governing permissions and\n",
              "// limitations under the License.\n",
              "\n",
              "/**\n",
              " * @fileoverview Helpers for google.colab Python module.\n",
              " */\n",
              "(function(scope) {\n",
              "function span(text, styleAttributes = {}) {\n",
              "  const element = document.createElement('span');\n",
              "  element.textContent = text;\n",
              "  for (const key of Object.keys(styleAttributes)) {\n",
              "    element.style[key] = styleAttributes[key];\n",
              "  }\n",
              "  return element;\n",
              "}\n",
              "\n",
              "// Max number of bytes which will be uploaded at a time.\n",
              "const MAX_PAYLOAD_SIZE = 100 * 1024;\n",
              "\n",
              "function _uploadFiles(inputId, outputId) {\n",
              "  const steps = uploadFilesStep(inputId, outputId);\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  // Cache steps on the outputElement to make it available for the next call\n",
              "  // to uploadFilesContinue from Python.\n",
              "  outputElement.steps = steps;\n",
              "\n",
              "  return _uploadFilesContinue(outputId);\n",
              "}\n",
              "\n",
              "// This is roughly an async generator (not supported in the browser yet),\n",
              "// where there are multiple asynchronous steps and the Python side is going\n",
              "// to poll for completion of each step.\n",
              "// This uses a Promise to block the python side on completion of each step,\n",
              "// then passes the result of the previous step as the input to the next step.\n",
              "function _uploadFilesContinue(outputId) {\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  const steps = outputElement.steps;\n",
              "\n",
              "  const next = steps.next(outputElement.lastPromiseValue);\n",
              "  return Promise.resolve(next.value.promise).then((value) => {\n",
              "    // Cache the last promise value to make it available to the next\n",
              "    // step of the generator.\n",
              "    outputElement.lastPromiseValue = value;\n",
              "    return next.value.response;\n",
              "  });\n",
              "}\n",
              "\n",
              "/**\n",
              " * Generator function which is called between each async step of the upload\n",
              " * process.\n",
              " * @param {string} inputId Element ID of the input file picker element.\n",
              " * @param {string} outputId Element ID of the output display.\n",
              " * @return {!Iterable<!Object>} Iterable of next steps.\n",
              " */\n",
              "function* uploadFilesStep(inputId, outputId) {\n",
              "  const inputElement = document.getElementById(inputId);\n",
              "  inputElement.disabled = false;\n",
              "\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  outputElement.innerHTML = '';\n",
              "\n",
              "  const pickedPromise = new Promise((resolve) => {\n",
              "    inputElement.addEventListener('change', (e) => {\n",
              "      resolve(e.target.files);\n",
              "    });\n",
              "  });\n",
              "\n",
              "  const cancel = document.createElement('button');\n",
              "  inputElement.parentElement.appendChild(cancel);\n",
              "  cancel.textContent = 'Cancel upload';\n",
              "  const cancelPromise = new Promise((resolve) => {\n",
              "    cancel.onclick = () => {\n",
              "      resolve(null);\n",
              "    };\n",
              "  });\n",
              "\n",
              "  // Wait for the user to pick the files.\n",
              "  const files = yield {\n",
              "    promise: Promise.race([pickedPromise, cancelPromise]),\n",
              "    response: {\n",
              "      action: 'starting',\n",
              "    }\n",
              "  };\n",
              "\n",
              "  cancel.remove();\n",
              "\n",
              "  // Disable the input element since further picks are not allowed.\n",
              "  inputElement.disabled = true;\n",
              "\n",
              "  if (!files) {\n",
              "    return {\n",
              "      response: {\n",
              "        action: 'complete',\n",
              "      }\n",
              "    };\n",
              "  }\n",
              "\n",
              "  for (const file of files) {\n",
              "    const li = document.createElement('li');\n",
              "    li.append(span(file.name, {fontWeight: 'bold'}));\n",
              "    li.append(span(\n",
              "        `(${file.type || 'n/a'}) - ${file.size} bytes, ` +\n",
              "        `last modified: ${\n",
              "            file.lastModifiedDate ? file.lastModifiedDate.toLocaleDateString() :\n",
              "                                    'n/a'} - `));\n",
              "    const percent = span('0% done');\n",
              "    li.appendChild(percent);\n",
              "\n",
              "    outputElement.appendChild(li);\n",
              "\n",
              "    const fileDataPromise = new Promise((resolve) => {\n",
              "      const reader = new FileReader();\n",
              "      reader.onload = (e) => {\n",
              "        resolve(e.target.result);\n",
              "      };\n",
              "      reader.readAsArrayBuffer(file);\n",
              "    });\n",
              "    // Wait for the data to be ready.\n",
              "    let fileData = yield {\n",
              "      promise: fileDataPromise,\n",
              "      response: {\n",
              "        action: 'continue',\n",
              "      }\n",
              "    };\n",
              "\n",
              "    // Use a chunked sending to avoid message size limits. See b/62115660.\n",
              "    let position = 0;\n",
              "    do {\n",
              "      const length = Math.min(fileData.byteLength - position, MAX_PAYLOAD_SIZE);\n",
              "      const chunk = new Uint8Array(fileData, position, length);\n",
              "      position += length;\n",
              "\n",
              "      const base64 = btoa(String.fromCharCode.apply(null, chunk));\n",
              "      yield {\n",
              "        response: {\n",
              "          action: 'append',\n",
              "          file: file.name,\n",
              "          data: base64,\n",
              "        },\n",
              "      };\n",
              "\n",
              "      let percentDone = fileData.byteLength === 0 ?\n",
              "          100 :\n",
              "          Math.round((position / fileData.byteLength) * 100);\n",
              "      percent.textContent = `${percentDone}% done`;\n",
              "\n",
              "    } while (position < fileData.byteLength);\n",
              "  }\n",
              "\n",
              "  // All done.\n",
              "  yield {\n",
              "    response: {\n",
              "      action: 'complete',\n",
              "    }\n",
              "  };\n",
              "}\n",
              "\n",
              "scope.google = scope.google || {};\n",
              "scope.google.colab = scope.google.colab || {};\n",
              "scope.google.colab._files = {\n",
              "  _uploadFiles,\n",
              "  _uploadFilesContinue,\n",
              "};\n",
              "})(self);\n",
              "</script> "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saving img_align_celeba (3).zip to img_align_celeba (3).zip\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Path to the uploaded ZIP file\n",
        "zip_file_path = \"/content/img_align_celeba(3).zip\"\n",
        "extract_dir = \"/content/img_align_celeba\"\n"
      ],
      "metadata": {
        "id": "rsB8hlICAEjB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Ensure the file exists and is a valid ZIP file\n",
        "if zipfile.is_zipfile(zip_file_path):\n",
        "    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:\n",
        "        zip_ref.extractall(extract_dir)\n",
        "else:\n",
        "    print(\"The specified file is not a valid ZIP file.\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ETM05JDBAEte",
        "outputId": "04c5d0c7-ed99-4cae-932c-feb5e0677e8a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The specified file is not a valid ZIP file.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Ensure the images are in a directory that TensorFlow can read from\n",
        "image_dir = os.path.join(extract_dir, 'img_align_celeba')"
      ],
      "metadata": {
        "id": "JzV6WiRUAJC_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the data\n",
        "train_data = utils.image_dataset_from_directory(\n",
        "    \"/content/img_align_celeba.zip\",\n",
        "    labels=None,\n",
        "    color_mode=\"rgb\",\n",
        "    image_size=(IMAGE_SIZE, IMAGE_SIZE),\n",
        "    batch_size=BATCH_SIZE,\n",
        "    shuffle=True,\n",
        "    seed=42,\n",
        "    interpolation=\"bilinear\",\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 376
        },
        "id": "udhvOxXFx3bO",
        "outputId": "927ab84e-f12a-4630-e298-a71bc6dd53d7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found 0 files belonging to 1 classes.\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "No images found in directory /content/img_align_celeba.zip. Allowed formats: ('.bmp', '.gif', '.jpeg', '.jpg', '.png')",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-31-3f230cdbb42c>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Load the data\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m train_data = utils.image_dataset_from_directory(\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;34m\"/content/img_align_celeba.zip\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mlabels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mcolor_mode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"rgb\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/keras/src/utils/image_dataset.py\u001b[0m in \u001b[0;36mimage_dataset_from_directory\u001b[0;34m(directory, labels, label_mode, class_names, color_mode, batch_size, image_size, shuffle, seed, validation_split, subset, interpolation, follow_links, crop_to_aspect_ratio, **kwargs)\u001b[0m\n\u001b[1;32m    301\u001b[0m         )\n\u001b[1;32m    302\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mimage_paths\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 303\u001b[0;31m             raise ValueError(\n\u001b[0m\u001b[1;32m    304\u001b[0m                 \u001b[0;34mf\"No images found in directory {directory}. \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    305\u001b[0m                 \u001b[0;34mf\"Allowed formats: {ALLOWLIST_FORMATS}\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mValueError\u001b[0m: No images found in directory /content/img_align_celeba.zip. Allowed formats: ('.bmp', '.gif', '.jpeg', '.jpg', '.png')"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Preprocess the data\n",
        "def preprocess(img):\n",
        "    img = tf.cast(img, \"float32\") / 255.0\n",
        "    return img\n",
        "\n",
        "\n",
        "train = train_data.map(lambda x: preprocess(x))"
      ],
      "metadata": {
        "id": "QDmwKhehx3gx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_sample = sample_batch(train)"
      ],
      "metadata": {
        "id": "m5DhPkGNx3ml"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Show some faces from the training set\n",
        "display(train_sample, cmap=None)"
      ],
      "metadata": {
        "id": "T6aUjNdtx3sX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Sampling(layers.Layer):\n",
        "    def call(self, inputs):\n",
        "        z_mean, z_log_var = inputs\n",
        "        batch = tf.shape(z_mean)[0]\n",
        "        dim = tf.shape(z_mean)[1]\n",
        "        epsilon = K.random_normal(shape=(batch, dim))\n",
        "        return z_mean + tf.exp(0.5 * z_log_var) * epsilon"
      ],
      "metadata": {
        "id": "2w1N1NZIx3yI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Encoder\n",
        "encoder_input = layers.Input(\n",
        "    shape=(IMAGE_SIZE, IMAGE_SIZE, CHANNELS), name=\"encoder_input\"\n",
        ")\n",
        "x = layers.Conv2D(NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\")(\n",
        "    encoder_input\n",
        ")\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Conv2D(NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Conv2D(NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Conv2D(NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "shape_before_flattening = K.int_shape(x)[1:]  # the decoder will need this!\n",
        "\n",
        "x = layers.Flatten()(x)\n",
        "z_mean = layers.Dense(Z_DIM, name=\"z_mean\")(x)\n",
        "z_log_var = layers.Dense(Z_DIM, name=\"z_log_var\")(x)\n",
        "z = Sampling()([z_mean, z_log_var])\n",
        "\n",
        "encoder = models.Model(encoder_input, [z_mean, z_log_var, z], name=\"encoder\")\n",
        "encoder.summary()"
      ],
      "metadata": {
        "id": "6Qtn-NSQx33s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Decoder\n",
        "decoder_input = layers.Input(shape=(Z_DIM,), name=\"decoder_input\")\n",
        "x = layers.Dense(np.prod(shape_before_flattening))(decoder_input)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Reshape(shape_before_flattening)(x)\n",
        "x = layers.Conv2DTranspose(\n",
        "    NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\"\n",
        ")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Conv2DTranspose(\n",
        "    NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\"\n",
        ")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Conv2DTranspose(\n",
        "    NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\"\n",
        ")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "x = layers.Conv2DTranspose(\n",
        "    NUM_FEATURES, kernel_size=3, strides=2, padding=\"same\"\n",
        ")(x)\n",
        "x = layers.BatchNormalization()(x)\n",
        "x = layers.LeakyReLU()(x)\n",
        "decoder_output = layers.Conv2DTranspose(\n",
        "    CHANNELS, kernel_size=3, strides=1, activation=\"sigmoid\", padding=\"same\"\n",
        ")(x)\n",
        "decoder = models.Model(decoder_input, decoder_output)\n",
        "decoder.summary()"
      ],
      "metadata": {
        "id": "J9uUkTLMx39G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class VAE(models.Model):\n",
        "    def __init__(self, encoder, decoder, **kwargs):\n",
        "        super(VAE, self).__init__(**kwargs)\n",
        "        self.encoder = encoder\n",
        "        self.decoder = decoder\n",
        "        self.total_loss_tracker = metrics.Mean(name=\"total_loss\")\n",
        "        self.reconstruction_loss_tracker = metrics.Mean(\n",
        "            name=\"reconstruction_loss\"\n",
        "        )\n",
        "        self.kl_loss_tracker = metrics.Mean(name=\"kl_loss\")\n",
        "\n",
        "    @property\n",
        "    def metrics(self):\n",
        "        return [\n",
        "            self.total_loss_tracker,\n",
        "            self.reconstruction_loss_tracker,\n",
        "            self.kl_loss_tracker,\n",
        "        ]\n",
        "\n",
        "    def call(self, inputs):\n",
        "        \"\"\"Call the model on a particular input.\"\"\"\n",
        "        z_mean, z_log_var, z = encoder(inputs)\n",
        "        reconstruction = decoder(z)\n",
        "        return z_mean, z_log_var, reconstruction\n",
        "\n",
        "    def train_step(self, data):\n",
        "        \"\"\"Step run during training.\"\"\"\n",
        "        with tf.GradientTape() as tape:\n",
        "            z_mean, z_log_var, reconstruction = self(data, training=True)\n",
        "            reconstruction_loss = tf.reduce_mean(\n",
        "                BETA * losses.mean_squared_error(data, reconstruction)\n",
        "            )\n",
        "            kl_loss = tf.reduce_mean(\n",
        "                tf.reduce_sum(\n",
        "                    -0.5\n",
        "                    * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)),\n",
        "                    axis=1,\n",
        "                )\n",
        "            )\n",
        "            total_loss = reconstruction_loss + kl_loss\n",
        "\n",
        "        grads = tape.gradient(total_loss, self.trainable_weights)\n",
        "        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))\n",
        "\n",
        "        self.total_loss_tracker.update_state(total_loss)\n",
        "        self.reconstruction_loss_tracker.update_state(reconstruction_loss)\n",
        "        self.kl_loss_tracker.update_state(kl_loss)\n",
        "\n",
        "        return {\n",
        "            \"loss\": self.total_loss_tracker.result(),\n",
        "            \"reconstruction_loss\": self.reconstruction_loss_tracker.result(),\n",
        "            \"kl_loss\": self.kl_loss_tracker.result(),\n",
        "        }\n",
        "\n",
        "    def test_step(self, data):\n",
        "        \"\"\"Step run during validation.\"\"\"\n",
        "        if isinstance(data, tuple):\n",
        "            data = data[0]\n",
        "\n",
        "        z_mean, z_log_var, reconstruction = self(data)\n",
        "        reconstruction_loss = tf.reduce_mean(\n",
        "            BETA * losses.mean_squared_error(data, reconstruction)\n",
        "        )\n",
        "        kl_loss = tf.reduce_mean(\n",
        "            tf.reduce_sum(\n",
        "                -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)),\n",
        "                axis=1,\n",
        "            )\n",
        "        )\n",
        "        total_loss = reconstruction_loss + kl_loss\n",
        "\n",
        "        return {\n",
        "            \"loss\": total_loss,\n",
        "            \"reconstruction_loss\": reconstruction_loss,\n",
        "            \"kl_loss\": kl_loss,\n",
        "        }"
      ],
      "metadata": {
        "id": "7qdJreFOx4Cb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a variational autoencoder\n",
        "vae = VAE(encoder, decoder)"
      ],
      "metadata": {
        "id": "ndVC9mC2x4H2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compile the variational autoencoder\n",
        "optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)\n",
        "vae.compile(optimizer=optimizer)"
      ],
      "metadata": {
        "id": "rL_NeJNOx4Nf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a model save checkpoint\n",
        "model_checkpoint_callback = callbacks.ModelCheckpoint(\n",
        "    filepath=\"./checkpoint\",\n",
        "    save_weights_only=False,\n",
        "    save_freq=\"epoch\",\n",
        "    monitor=\"loss\",\n",
        "    mode=\"min\",\n",
        "    save_best_only=True,\n",
        "    verbose=0,\n",
        ")\n",
        "\n",
        "tensorboard_callback = callbacks.TensorBoard(log_dir=\"./logs\")\n",
        "\n",
        "\n",
        "class ImageGenerator(callbacks.Callback):\n",
        "    def __init__(self, num_img, latent_dim):\n",
        "        self.num_img = num_img\n",
        "        self.latent_dim = latent_dim\n",
        "\n",
        "    def on_epoch_end(self, epoch, logs=None):\n",
        "        random_latent_vectors = tf.random.normal(\n",
        "            shape=(self.num_img, self.latent_dim)\n",
        "        )\n",
        "        generated_images = self.model.decoder(random_latent_vectors)\n",
        "        generated_images *= 255\n",
        "        generated_images.numpy()\n",
        "        for i in range(self.num_img):\n",
        "            img = utils.array_to_img(generated_images[i])\n",
        "            img.save(\"./output/generated_img_%03d_%d.png\" % (epoch, i))"
      ],
      "metadata": {
        "id": "jPqye5lrx4S2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load old weights if required\n",
        "if LOAD_MODEL:\n",
        "    vae.load_weights(\"./models/vae\")\n",
        "    tmp = vae.predict(train.take(1))"
      ],
      "metadata": {
        "id": "oJJ5W-Zux4YR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vae.fit(\n",
        "    train,\n",
        "    epochs=EPOCHS,\n",
        "    callbacks=[\n",
        "        model_checkpoint_callback,\n",
        "        tensorboard_callback,\n",
        "        ImageGenerator(num_img=10, latent_dim=Z_DIM),\n",
        "    ],\n",
        ")"
      ],
      "metadata": {
        "id": "rYM-AlUDx4dm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Save the final models\n",
        "vae.save(\"./models/vae\")\n",
        "encoder.save(\"./models/encoder\")\n",
        "decoder.save(\"./models/decoder\")"
      ],
      "metadata": {
        "id": "KGtV3zkqx4i7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Select a subset of the test set\n",
        "batches_to_predict = 1\n",
        "example_images = np.array(\n",
        "    list(train.take(batches_to_predict).get_single_element())\n",
        ")"
      ],
      "metadata": {
        "id": "Fwhy_Qr1x4oX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Create autoencoder predictions and display\n",
        "z_mean, z_log_var, reconstructions = vae.predict(example_images)\n",
        "print(\"Example real faces\")\n",
        "display(example_images)\n",
        "print(\"Reconstructions\")\n",
        "display(reconstructions)"
      ],
      "metadata": {
        "id": "h49hJioBx4ts"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "_, _, z = vae.encoder.predict(example_images)\n",
        "\n",
        "x = np.linspace(-3, 3, 100)\n",
        "\n",
        "fig = plt.figure(figsize=(20, 5))\n",
        "fig.subplots_adjust(hspace=0.6, wspace=0.4)\n",
        "\n",
        "for i in range(50):\n",
        "    ax = fig.add_subplot(5, 10, i + 1)\n",
        "    ax.hist(z[:, i], density=True, bins=20)\n",
        "    ax.axis(\"off\")\n",
        "    ax.text(\n",
        "        0.5, -0.35, str(i), fontsize=10, ha=\"center\", transform=ax.transAxes\n",
        "    )\n",
        "    ax.plot(x, norm.pdf(x))\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "YDofQwf1x4zR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Sample some points in the latent space, from the standard normal distribution\n",
        "grid_width, grid_height = (10, 3)\n",
        "z_sample = np.random.normal(size=(grid_width * grid_height, Z_DIM))"
      ],
      "metadata": {
        "id": "Pq4bPwoxx44o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Decode the sampled points\n",
        "reconstructions = decoder.predict(z_sample)"
      ],
      "metadata": {
        "id": "2n1lftCBx49z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Draw a plot of decoded images\n",
        "fig = plt.figure(figsize=(18, 5))\n",
        "fig.subplots_adjust(hspace=0.4, wspace=0.4)\n",
        "\n",
        "# Output the grid of faces\n",
        "for i in range(grid_width * grid_height):\n",
        "    ax = fig.add_subplot(grid_height, grid_width, i + 1)\n",
        "    ax.axis(\"off\")\n",
        "    ax.imshow(reconstructions[i, :, :])"
      ],
      "metadata": {
        "id": "glaTfWy2x5C4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the label dataset\n",
        "attributes = pd.read_csv(\"/app/data/celeba-dataset/list_attr_celeba.csv\")\n",
        "print(attributes.columns)\n",
        "attributes.head()"
      ],
      "metadata": {
        "id": "cZlheMfWx5H7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the face data with label attached\n",
        "LABEL = \"Blond_Hair\"  # <- Set this label\n",
        "labelled_test = utils.image_dataset_from_directory(\n",
        "    \"/app/data/celeba-dataset/img_align_celeba\",\n",
        "    labels=attributes[LABEL].tolist(),\n",
        "    color_mode=\"rgb\",\n",
        "    image_size=(IMAGE_SIZE, IMAGE_SIZE),\n",
        "    batch_size=BATCH_SIZE,\n",
        "    shuffle=True,\n",
        "    seed=42,\n",
        "    validation_split=0.2,\n",
        "    subset=\"validation\",\n",
        "    interpolation=\"bilinear\",\n",
        ")\n",
        "\n",
        "labelled = labelled_test.map(lambda x, y: (preprocess(x), y))"
      ],
      "metadata": {
        "id": "uc1YFCIby-4g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Find the attribute vector\n",
        "attribute_vec = get_vector_from_label(labelled, vae, Z_DIM, LABEL)"
      ],
      "metadata": {
        "id": "0-xl9raLy-_R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Add vector to images\n",
        "add_vector_to_images(labelled, vae, attribute_vec)"
      ],
      "metadata": {
        "id": "YXgpZHsUy_Fd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "morph_faces(labelled, vae)"
      ],
      "metadata": {
        "id": "dLcJhvuey_LL"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}