scripts/Test_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.utils as vutils\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "from PIL import Image\n",
    "\n",
    "import gensim\n",
    "\n",
    "model = gensim.models.KeyedVectors.load_word2vec_format('../models/word2vec_model.bin', binary=True) \n",
    "\n",
    "def get_feature_space_representaion(sentence):\n",
    "    doc_tokens = sentence.split()\n",
    "    words = [token for token in doc_tokens if token in model.vocab]\n",
    "    features = model[words]\n",
    "    mins = np.min(features, axis=0)\n",
    "    maxs = np.max(features, axis=0)\n",
    "    stack = np.hstack((mins, maxs))\n",
    "    return stack\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "from pprint import pprint\n",
    "from enum import Enum\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(Generator, self).__init__()\n",
    "        self.ngpu = config.N_GPU.value\n",
    "        nz = config.N_LATENT_VECTOR.value\n",
    "        nzc = config.N_TEXT_EMBEDDING.value\n",
    "        ngf = config.N_GENERATOR_FEATURE_MAP.value\n",
    "        nc = config.N_COLOR_CHANNELS.value\n",
    "        self.main = nn.Sequential(\n",
    "            # input is Z, going into a convolution\n",
    "            nn.ConvTranspose2d( nz+nzc, ngf * 8, 4, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 8),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf*8) x 4 x 4\n",
    "            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 4),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf*4) x 8 x 8\n",
    "            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 2),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf*2) x 16 x 16\n",
    "            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf) x 32 x 32\n",
    "            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),\n",
    "            nn.Tanh()\n",
    "            # state size. (nc) x 64 x 64\n",
    "        )\n",
    "\n",
    "    def forward(self, image, comment):\n",
    "        comment = comment.reshape((comment.shape[0],comment.shape[1],1,1))\n",
    "        cat = torch.cat((image, comment), dim=1)\n",
    "        return self.main(cat)\n",
    "\n",
    "#load image model\n",
    "\n",
    "config_path = '../config/config.yml'\n",
    "with open(config_path, 'r') as stream:\n",
    "    try:\n",
    "        config = yaml.safe_load(stream)\n",
    "    except yaml.YAMLError as exc:\n",
    "        print(exc)\n",
    "print('Config loaded from: {}'.format(config_path))\n",
    "\n",
    "#sentence = 'woman in playing ball on beach'\n",
    "sentence = 'A man with a jackhammer demolishing cement'\n",
    "\n",
    "comment = get_feature_space_representaion(sentence)\n",
    "\n",
    "config = Enum('config', config)\n",
    "device = torch.device(\"cuda:0\" if (config.GPU_MODE.value and torch.cuda.is_available() and config.N_GPU.value > 0) else \"cpu\")\n",
    "fixed_noise = torch.randn(1, config.N_LATENT_VECTOR.value, 1, 1, device=device)\n",
    "\n",
    "path_to_model = '../models/generator_50.pth'\n",
    "generator = Generator(config).to(device)\n",
    "generator.load_state_dict(torch.load(path_to_model))\n",
    "generator.eval()\n",
    "\n",
    "comment = torch.tensor(comment.reshape((1,600,1))).to(device)\n",
    "img = generator(fixed_noise, comment).detach().cpu()\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.axis(\"off\")\n",
    "#plt.title(\"Result of sentence: {}\".format(sentence))\n",
    "img = vutils.make_grid(img, padding=0, normalize=True)\n",
    "plt.imshow(np.transpose(img,(1,2,0)))"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}