play_text_to_face.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import PIL\n",
    "from keras.models import load_model\n",
    "from model import textTodlatentsModel\n",
    "from bert_serving.client import BertClient\n",
    "import dnnlib\n",
    "import pickle\n",
    "import dnnlib.tflib as tflib\n",
    "from textwrap import wrap\n",
    "from encoder.generator_model import Generator\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['font.family'] = ['sans-serif']\n",
    "plt.rcParams['font.sans-serif'] = ['simhei']\n",
    "plt.rcParams['axes.unicode_minus'] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tflib.init_tf()\n",
    "with open('./model/stylegan.pkl', 'rb') as f:\n",
    "    generator_network, discriminator_network, Gs_network = pickle.load(f)\n",
    "\n",
    "generator = Generator(Gs_network, batch_size=1, randomize_noise=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(inputText):\n",
    "    sentenceEmbedding = np.zeros(shape = [1, 768], dtype = np.float)\n",
    "    bc = BertClient()\n",
    "    sentenceEmbedding = bc.encode([inputText])\n",
    "    model = load_model(\"./model/textEmbeddingDlatents.h5\")\n",
    "    sentenceEmbedding = np.expand_dims(sentenceEmbedding, 0)\n",
    "    result = model.predict(sentenceEmbedding)\n",
    "    return result\n",
    "def generate_image(latent_vector):\n",
    "    latent_vector = latent_vector.reshape((1, 18, 512))\n",
    "    generator.set_dlatents(latent_vector)\n",
    "    img_array = generator.generate_images()[0]\n",
    "    img = PIL.Image.fromarray(img_array, 'RGB')\n",
    "    return img.resize((256, 256))\n",
    "text1 = \"黑色短发年轻女子，有刘海，开心的笑，露齿，大眼睛\"\n",
    "text2 = \"金色短发年轻女子，有刘海，开心的笑，露齿，大眼睛\"\n",
    "text3 = \"金色长发年轻女子，有刘海，开心的笑，露齿，大眼睛\"\n",
    "result1 = predict(text1)\n",
    "image1 = generate_image(result1)\n",
    "result2 = predict(text2)\n",
    "image2 = generate_image(result2)\n",
    "result3 = predict(text3)\n",
    "image3 = generate_image(result3)\n",
    "plt.figure(figsize=(15,15), dpi=100)\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.title(\"\\n\".join(wrap(text1, 14)), fontsize=12)\n",
    "plt.imshow(image1)\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.title(\"\\n\".join(wrap(text2, 14)), fontsize=12)\n",
    "plt.imshow(image2)\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.title(\"\\n\".join(wrap(text3, 14)), fontsize=12)\n",
    "plt.imshow(image3)\n",
    "plt.savefig('example.png',bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}