inference_FUNIT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating images with FUNIT model from a reference style\n",
    "In this example, we'll generate images with trained FUNIT model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Loading packages\n",
    "* First, load the packages used in this code.\n",
    "* All of the packages are avilable in `pip`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "\n",
    "import torch\n",
    "from sconf import Config\n",
    "from torchvision import transforms\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((128, 128)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* These modules are defined in this repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from base.dataset import read_font, render, sample\n",
    "from base.utils import save_tensor_to_image, load_reference\n",
    "from FUNIT.models.networks import FewShotGen\n",
    "from inference import infer_FUNIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Build model\n",
    "* Build and load the trained model.\n",
    "* `weight_path` : \n",
    "    * The location of the trained model weight."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "###############################################################\n",
    "weight_path = \"weights/kor/funit.pth\"  # path to weight to infer\n",
    "###############################################################\n",
    "\n",
    "# building and loading the model (not recommended to modify)\n",
    "cfg = Config(\"cfgs/FUNIT/default.yaml\")\n",
    "cfg.gen.n_mlp_blks = 2\n",
    "\n",
    "gen = FewShotGen(**cfg.gen).cuda().eval()\n",
    "weight = torch.load(weight_path)\n",
    "gen.load_state_dict(weight[\"generator_ema\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Load reference images.\n",
    "* `ref_path`: \n",
    "    * The path of reference font or images.\n",
    "    * If you are using a ttf file, set this to the location of the ttf file.\n",
    "    * If you want to use rendered images, set this to the path to the directory which contains the reference images.\n",
    "* `ref_chars`:\n",
    "    * The characters of reference images.\n",
    "    * If this is `None`, all the available images will be loaded.\n",
    "* `extension`:\n",
    "    * If you are using ttf files, set this to \"ttf\".\n",
    "    * If you are using image files, set this to their extension(png, jpg, etc..)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "###############################################################\n",
    "ref_path = \"data_example/kor/png\"\n",
    "extension = \"png\"\n",
    "ref_chars = \"값넋닻볘츄퀭핥훟\"\n",
    "## Comment upper lines and uncomment lower lines to test with ttf files.\n",
    "# extension = \"ttf\"\n",
    "# ref_chars = \"값같곬곶깎넋늪닫닭닻됩뗌략몃밟볘뺐뽈솩쐐앉않얘얾엌옳읊죡쮜춰츄퀭틔핀핥훟\"\n",
    "###############################################################\n",
    "\n",
    "ref_dict, load_img = load_reference(ref_path, extension, ref_chars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Generate and save the images.\n",
    "* `gen_chars`: The characters to generate.\n",
    "* `save_dir`: Path to save the generated images.\n",
    "* `source_path`: Path to the source.\n",
    "* `source_ext`: Extension of the source file. If you are using image files, set this to the image's extension.\n",
    "* `batch_size`: The number of images inferred at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UhBee_HanByeol\n",
      "UhBee_charming\n"
     ]
    }
   ],
   "source": [
    "###############################################################\n",
    "gen_chars = \"좋은하루되세요\"  # characters to generate\n",
    "save_dir = \"./result/funit\"\n",
    "source_path = \"data/kor/source.ttf\"\n",
    "source_ext = \"ttf\"\n",
    "batch_size = 16\n",
    "###############################################################\n",
    "\n",
    "outs = infer_FUNIT(gen, save_dir, source_path, source_ext, gen_chars, ref_dict, load_img, batch_size, True)"
   ]
  },
  {
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}