ResNetVAE_reconstruction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "import torch.utils.data as data\n",
    "import torchvision\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "from modules import *\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pickle\n",
    "from sklearn.datasets import fetch_olivetti_faces\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "from skimage.transform import resize\n",
    "\n",
    "\n",
    "def decoder(model, device, z):\n",
    "    model.eval()\n",
    "    z = Variable(torch.FloatTensor(z)).to(device)\n",
    "    new_images = model.decode(z).squeeze_().data.cpu().numpy().transpose((1, 2, 0))\n",
    "    return new_images\n",
    "\n",
    "saved_model_path = './results_Olivetti_face'\n",
    "# saved_model_path = './results_MNIST'\n",
    "\n",
    "exp = 'Olivetti'\n",
    "# exp = 'MNIST'\n",
    "\n",
    "# use same ResNet Encoder saved earlier!\n",
    "CNN_fc_hidden1, CNN_fc_hidden2 = 1024, 1024\n",
    "CNN_embed_dim = 256\n",
    "res_size = 224        # ResNet image size\n",
    "dropout_p = 0.2       # dropout probability\n",
    "\n",
    "epoch = 20\n",
    "\n",
    "\n",
    "use_cuda = torch.cuda.is_available()                   # check if GPU exists\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")   # use CPU or GPU\n",
    "\n",
    "# reload ResNetVAE model\n",
    "resnet_vae = ResNet_VAE(fc_hidden1=CNN_fc_hidden1, fc_hidden2=CNN_fc_hidden2, drop_p=dropout_p, CNN_embed_dim=CNN_embed_dim).to(device)\n",
    "resnet_vae.load_state_dict(torch.load(os.path.join(saved_model_path, 'model_epoch{}.pth'.format(epoch))))\n",
    "\n",
    "print('ResNetVAE epoch {} model reloaded!'.format(epoch))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reconstruction "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_train = np.load(os.path.join(saved_model_path, 'z_{}_train_epoch{}.npy').format(exp, epoch))\n",
    "X_train = np.load(os.path.join(saved_model_path, 'X_{}_train_epoch{}.npy').format(exp, epoch))\n",
    "\n",
    "ind = 1\n",
    "zz = torch.from_numpy(z_train[ind]).view(1, -1)\n",
    "X = np.transpose(X_train[ind], (1, 2, 0))\n",
    "\n",
    "new_imgs = decoder(resnet_vae, device, zz)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(X)\n",
    "plt.title('original')\n",
    "plt.axis('off')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(new_imgs)\n",
    "plt.title('reconstructed')\n",
    "plt.axis('off')\n",
    "plt.savefig(\"./reconstruction_{}.png\".format(exp), bbox_inches='tight', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate new images from latent points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# choose two original images\n",
    "sample1, sample2 = 0, 1\n",
    "w = 0.4 # weight for fusing two images\n",
    "\n",
    "X1 = np.transpose(X_train[-sample1], (1, 2, 0))\n",
    "X2 = np.transpose(X_train[-sample2], (1, 2, 0))\n",
    "\n",
    "# generate image using decoder\n",
    "z_train = np.load(os.path.join(saved_model_path, 'z_{}_train_epoch{}.npy').format(exp, epoch))\n",
    "z = z_train[-sample1] * w + z_train[-sample2] * (1 - w)\n",
    "new_imgs = decoder(resnet_vae, device, torch.from_numpy(z).view(1, -1))\n",
    "\n",
    "fig = plt.figure(figsize=(15, 15))\n",
    "\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.imshow(X1)\n",
    "plt.title('original 1')\n",
    "plt.axis('off')\n",
    "\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.imshow(X2)\n",
    "plt.title('original 2')\n",
    "plt.axis('off')\n",
    "\n",
    "\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.imshow(new_imgs)\n",
    "plt.title('new image')\n",
    "plt.axis('off')\n",
    "plt.savefig(\"./generated_{}.png\".format(exp), bbox_inches='tight', dpi=600)\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}