notebooks/9-DoomVAE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BakUHRa3rmdg"
   },
   "source": [
    "# Doom VAE\n",
    "\n",
    "The idea of this notebook is to construct a small variational auto-encoder that can reproduce images from the \"Doom\" video game.\n",
    "\n",
    "This should all work in Colab, including some fun controls below to investigate the model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LgWRSyQGrN9Q"
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "from keras import layers\n",
    "from keras.layers import Input, Dense, Lambda, Flatten, Reshape, Layer\n",
    "from keras.layers import Conv2D, Conv2DTranspose\n",
    "from keras.models import Model\n",
    "from keras import backend as K\n",
    "from keras.models import Model\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from IPython.display import Image, display\n",
    "from keras.preprocessing import image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mZ6dRfd4u_9Z"
   },
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "url = 'http://folk.uio.no/charlepm/datasets/doom_images.npz'  \n",
    "urllib.request.urlretrieve(url, './doom_images.npz') \n",
    "\n",
    "# Test loading new file.\n",
    "with np.load('doom_images.npz') as data:\n",
    "    x_train = data['arr_0']\n",
    "\n",
    "# View an input\n",
    "#plt.imshow(x_train[0])\n",
    "print(\"Here's a sample image:\")\n",
    "img = image.array_to_img(x_train[np.random.randint(len(x_train))], scale=False)\n",
    "display(img.resize((300, 300)))\n",
    "\n",
    "# The data, split between train and test sets:\n",
    "print('x_train shape:', x_train.shape)\n",
    "print(x_train.shape[0], 'train samples')\n",
    "\n",
    "x_train = x_train.astype('float32') / 255 # scale to [0,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Fei5uiJnrN9X"
   },
   "outputs": [],
   "source": [
    "# Setup neural network hyperparameters\n",
    "img_rows, img_cols, img_chns = 64, 64, 3\n",
    "latent_dim = 16\n",
    "intermediate_dim = 128\n",
    "epsilon_std = 1.0\n",
    "epochs = 100\n",
    "filters = 32\n",
    "num_conv = 3\n",
    "batch_size = 128\n",
    "\n",
    "img_size = (img_rows, img_cols, img_chns)\n",
    "original_dim = img_rows * img_cols * img_chns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oJ_gYRVdrN9Y"
   },
   "outputs": [],
   "source": [
    "# Enc\n",
    "input_img = Input(shape=img_size, name='encoder_input')\n",
    "x = Conv2D(img_chns, kernel_size=(2,2), padding='same', activation='relu')(input_img)\n",
    "x = Conv2D(filters, kernel_size=(2,2), padding='same', activation='relu', strides=(2,2))(x)\n",
    "x = Conv2D(filters, kernel_size=(2,2), padding='same', activation='relu', strides=(2,2))(x)\n",
    "# x = keras.layers.MaxPooling2D(pool_size=(2, 2), strides=None, padding='same')(x) # try a max pooling layer here instead of the previous stride\n",
    "x = Conv2D(filters, kernel_size=num_conv, padding='same', activation='relu', strides=1)(x)\n",
    "shape_before_flattening = K.int_shape(x)\n",
    "x = Flatten()(x)\n",
    "x = Dense(intermediate_dim, activation='relu', name='latent_project')(x)\n",
    "\n",
    "print(\"Shape before flattening:\",shape_before_flattening)\n",
    "\n",
    "# mean and var\n",
    "z_mean = Dense(latent_dim, name='Z_mean')(x)\n",
    "z_log_var = Dense(latent_dim, name='Z_var')(x)\n",
    "\n",
    "# make an encoder model (not used until after training)\n",
    "encoder = Model(input_img, z_mean)\n",
    "\n",
    "# sampling layer\n",
    "def sampling(args):\n",
    "    z_mean, z_log_var = args\n",
    "    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=epsilon_std)\n",
    "    return z_mean + K.exp(z_log_var) * epsilon\n",
    "\n",
    "z = layers.Lambda(sampling, name=\"Z_sample\")([z_mean, z_log_var])\n",
    "\n",
    "# dec\n",
    "decoder_input = layers.Input(K.int_shape(z)[1:])\n",
    "y = Dense(intermediate_dim, activation='relu')(decoder_input) # (z)\n",
    "y = Dense(np.prod(shape_before_flattening[1:]), activation='relu')(y)\n",
    "y = Reshape(shape_before_flattening[1:])(y)\n",
    "y = Conv2DTranspose(filters, kernel_size=num_conv, padding='same', strides=1, activation='relu', name='deconv_1')(y) # deconv 1\n",
    "y = Conv2DTranspose(filters, kernel_size=num_conv, padding='same', strides=(2,2), activation='relu', name='deconv_2')(y) # deconv 2\n",
    "y = Conv2DTranspose(filters, kernel_size=(3, 3), strides=(2, 2), padding='valid', activation='relu', name='deconv_3')(y) # deconv 3, upsamp\n",
    "y = Conv2D(img_chns, kernel_size=2, padding='valid', activation='sigmoid', name=\"mean_squash\")(y) # mean squash\n",
    "decoder = Model(decoder_input, y, name=\"Decoder\")\n",
    "z_decoded = decoder(z) #y\n",
    "\n",
    "def xent(y_true, y_pred):\n",
    "  return keras.metrics.binary_crossentropy(y_true, y_pred)\n",
    "\n",
    "def kl_measure(loc, log_var):\n",
    "  return -0.5 * K.mean(1 + log_var - K.square(loc) - K.exp(log_var), axis=-1)\n",
    "  \n",
    "def kl_custom_metric(y_true, y_pred):\n",
    "  # Ignore input and take from z tensors.\n",
    "  return kl_measure(z_mean, z_log_var)\n",
    "  \n",
    "class VAELayer(keras.layers.Layer):    \n",
    "    def __init__(self, **kwargs):\n",
    "        self.is_placeholder = True\n",
    "        super(VAELayer, self).__init__(**kwargs)\n",
    "      \n",
    "    def vae_loss(self, x, z_decoded):\n",
    "        x = K.flatten(x)\n",
    "        z_decoded = K.flatten(z_decoded)\n",
    "        r_loss = original_dim * xent(x, z_decoded)\n",
    "        kl_loss = kl_measure(z_mean, z_log_var)\n",
    "        print(\"KL Shape:\", K.int_shape(kl_loss))\n",
    "        print(\"Xent shape:\", K.int_shape(r_loss))\n",
    "        return K.mean(r_loss + kl_loss)\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        x = inputs[0]\n",
    "        z_decoded = inputs[1]\n",
    "        loss = self.vae_loss(x, z_decoded)\n",
    "        self.add_loss(loss, inputs=inputs)\n",
    "        return x\n",
    "\n",
    "y = VAELayer()([input_img, z_decoded])\n",
    "\n",
    "vae = Model(input_img, y, name=\"VAE\")\n",
    "vae.compile(optimizer='adam', metrics=['mse','binary_crossentropy'])\n",
    "\n",
    "\n",
    "decoder.summary()\n",
    "vae.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uiWX-hhrrN9b"
   },
   "outputs": [],
   "source": [
    "# Train!\n",
    "history = vae.fit(x_train, epochs=epochs, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NDGHPBEtrN9e"
   },
   "outputs": [],
   "source": [
    "# Plot the training loss.\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(history.history['loss'])\n",
    "plt.xlabel(\"epochs\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "69oDWzgiuAlP"
   },
   "outputs": [],
   "source": [
    "# Alternatively, download and load weights.\n",
    "!wget http://folk.uio.no/charlepm/doom_vae_weights.h5\n",
    "vae.load_weights(\"doom_vae_weights.h5\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uIS_SmmZ9aie"
   },
   "outputs": [],
   "source": [
    "# Let's see how the encoder works\n",
    "# First we'll take a random image from the corpus and encode it to a latent vector:\n",
    "ex = x_train[np.random.randint(len(x_train))]\n",
    "plt.figure(figsize=(5, 5))\n",
    "plt.imshow(ex) # cmap ignored if input is 3D (as it should be here)\n",
    "plt.show()\n",
    "\n",
    "enc_z = encoder.predict(np.array([ex]))\n",
    "display(enc_z[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QLh2FpAa9q2f"
   },
   "outputs": [],
   "source": [
    "# Now we can decode from the same vector to try to reproduce that image:\n",
    "ex_dec = decoder.predict(np.array([enc_z[0]]))\n",
    "# Plot output\n",
    "plt.figure(figsize=(5, 5))\n",
    "plt.imshow(ex_dec[0]) # cmap ignored if input is 3D (as it should be here)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jwpfzmkFrN9f"
   },
   "outputs": [],
   "source": [
    "# Let's try sampling different parts of the latent space to see what we have.\n",
    "n = 10 # num images\n",
    "img_size = 64\n",
    "figure = np.zeros((img_size * n, img_size * n, img_chns))\n",
    "\n",
    "for i in range(n):\n",
    "    for j in range(n):\n",
    "        z_sample = np.array([np.random.uniform(-1,1 ,size=latent_dim)])\n",
    "        x_decoded = decoder.predict(z_sample)\n",
    "        img = x_decoded[0].reshape(img_size, img_size, img_chns)\n",
    "        figure[i * img_size: (i + 1) * img_size,j * img_size: (j + 1) * img_size] = img\n",
    "\n",
    "plt.figure(figsize=(20, 20))\n",
    "plt.imshow(figure)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pHHk1QUdwmXQ"
   },
   "outputs": [],
   "source": [
    "# Save and download models\n",
    "!mkdir models\n",
    "\n",
    "def save_model_three_ways(model, name=\"model\"):\n",
    "  # Save the weights\n",
    "  # model.save(\"./models/\" + name + ('_ld_%d_conv_%d_id_%d_e_%d.h5' % (latent_dim, num_conv, intermediate_dim, epochs)))\n",
    "  model.save_weights(\"./models/\" + name + '_weights.h5')\n",
    "  # Save the model architecture\n",
    "  with open(\"./models/\" + name + '_architecture.json', 'w') as f:\n",
    "    f.write(model.to_json())\n",
    "\n",
    "save_model_three_ways(vae, name=\"vae\")\n",
    "save_model_three_ways(encoder, name=\"encoder\")\n",
    "save_model_three_ways(encoder, name=\"decoder\")\n",
    "!tar -czvf doom_models.tar.gz models\n",
    "\n",
    "#from google.colab import files\n",
    "#files.download('doom_models.tar.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "idmYMHQpgnz9"
   },
   "outputs": [],
   "source": [
    "# Colab only!\n",
    "#@title Interactive Latent Space Exploration { run: \"auto\", vertical-output: true, form-width: \"50px\" }\n",
    "z_1 = 0.88 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_2 = -0.85 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_3 = 0.7 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_4 = 0.51 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_5 = 0.15 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_6 = 0.23 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_7 = 0.51 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_8 = -0.8 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_9 = 0.77 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_10 = -0.99 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_11 = -0.51 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_12 = -0.29 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_13 = -0.48 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_14 = -0.6 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_15 = 0.56 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "z_16 = 0.92 #@param {type:\"slider\", min:-1, max:1, step:0.01}\n",
    "\n",
    "new_z= np.array([z_1,z_2,z_3,z_4,z_5,z_6,z_7,z_8,z_9,z_10,z_11,z_12,z_13,z_14,z_15,z_16])\n",
    "print(new_z)\n",
    "dec = decoder.predict(np.array([new_z]))\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(dec[0]) # cmap ignored if input is 3D (as it should be here)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wlftHnThipcs"
   },
   "outputs": [],
   "source": [
    "ex = x_train[1]\n",
    "plt.figure(figsize=(5, 5))\n",
    "plt.imshow(ex) # cmap ignored if input is 3D (as it should be here)\n",
    "plt.show()\n",
    "\n",
    "enc_z = encoder.predict(np.array([ex]))[0]\n",
    "print(enc_z)\n",
    "z_1 = enc_z[0]\n",
    "z_2 = enc_z[1]\n",
    "z_3 = enc_z[2]\n",
    "z_4 = enc_z[3]\n",
    "z_5 = enc_z[4]\n",
    "z_6 = enc_z[5]\n",
    "z_7 = enc_z[6]\n",
    "z_8 = enc_z[7]\n",
    "z_9 = enc_z[8]\n",
    "z_10 = enc_z[9]\n",
    "z_11 = enc_z[10]\n",
    "z_12 = enc_z[11]\n",
    "z_13 = enc_z[12]\n",
    "z_14 = enc_z[13]\n",
    "z_15 = enc_z[14]\n",
    "z_16 = enc_z[15]"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "VAE-Doom",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}