Model_infer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "G4tOfpA7aRaB",
    "outputId": "ccec3f96-4799-4bc2-b6ed-702d926a3c19"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/EvgenyKashin/stylegan2-distillation  # clone\n",
    "%cd stylegan2-distillation\n",
    "\n",
    "import imageio\n",
    "import requests\n",
    "import bz2\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "import dlib\n",
    "from pix2pixHD.data.base_dataset import __scale_width\n",
    "from pix2pixHD.models.networks import define_G\n",
    "import pix2pixHD.util.util as util\n",
    "from aligner import align_face\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XdJ0i4JMaRaF"
   },
   "source": [
    "## Load image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aQkiwNX_aRaG"
   },
   "outputs": [],
   "source": [
    "img_url = 'https://upload.wikimedia.org/wikipedia/commons/thumb/7/78/Oona_Chaplin_2016.jpg/800px-Oona_Chaplin_2016.jpg'\n",
    "img_filename = 'image.jpg'\n",
    "imageio.imwrite(img_filename, imageio.imread(img_url))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xyg8dG0laRaH"
   },
   "source": [
    "## Face aligment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TDNtyFGwaRaH"
   },
   "outputs": [],
   "source": [
    "def unpack_bz2(src_path):\n",
    "    data = bz2.BZ2File(src_path).read()\n",
    "    dst_path = src_path[:-4]\n",
    "    with open(dst_path, 'wb') as fp:\n",
    "        fp.write(data)\n",
    "    return dst_path\n",
    "\n",
    "def download(url, file_name):\n",
    "    with open(file_name, \"wb\") as file:\n",
    "        response = requests.get(url)\n",
    "        file.write(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3us3G2CzaRaH"
   },
   "outputs": [],
   "source": [
    "shape_model_url = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'\n",
    "shape_model_path = 'landmarks.dat'\n",
    "download(shape_model_url, shape_model_path)\n",
    "shape_predictor = dlib.shape_predictor(unpack_bz2(shape_model_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wnaNsYxgaRaI",
    "outputId": "2df84359-c080-4f7b-df43-6b21035d8339"
   },
   "outputs": [],
   "source": [
    "aligned_img = align_face(img_filename, shape_predictor)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uruj93d5aRaI"
   },
   "source": [
    "## Prepare pix2pixHD model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-O6WbCpfaRaI"
   },
   "outputs": [],
   "source": [
    "def get_eval_transform(loadSize=512):\n",
    "    transform_list = []\n",
    "    transform_list.append(transforms.Lambda(lambda img: __scale_width(img,\n",
    "                                                                      loadSize,\n",
    "                                                                      Image.BICUBIC)))\n",
    "    transform_list += [transforms.ToTensor()]\n",
    "    transform_list += [transforms.Normalize((0.5, 0.5, 0.5),\n",
    "                                            (0.5, 0.5, 0.5))]\n",
    "    return transforms.Compose(transform_list)\n",
    "\n",
    "transform = get_eval_transform()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rA4Jr6qQbrz5",
    "outputId": "e13b7f82-ebbf-4ffc-856a-fd9f7a82c59e"
   },
   "outputs": [],
   "source": [
    "# To male:\n",
    "!gdown --id 1-6J1CYLsIysk38X9DNN23lIcnvOr8aYh # https://drive.google.com/file/d/1-6J1CYLsIysk38X9DNN23lIcnvOr8aYh/view?usp=sharing\n",
    "\n",
    "# To female:\n",
    "!gdown --id 1frJERJr0WM_R38LnSFQ6XjGQtcXnLco1 # https://drive.google.com/file/d/1frJERJr0WM_R38LnSFQ6XjGQtcXnLco1/view?usp=sharing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "o8HSu2v3aRaJ",
    "outputId": "b1442cea-b867-4657-b344-3e31cf944337"
   },
   "outputs": [],
   "source": [
    "config_G = {\n",
    "    'input_nc': 3,\n",
    "    'output_nc': 3,\n",
    "    'ngf': 64,\n",
    "    'netG': 'global',\n",
    "    'n_downsample_global': 4,\n",
    "    'n_blocks_global': 9,\n",
    "    'n_local_enhancers': 1,\n",
    "    'norm': 'instance',\n",
    "}\n",
    "\n",
    "weights_path = '/content/stylegan2-distillation/to_male_net_G.pth'  # to_female_net_G.pth\n",
    "\n",
    "model = define_G(**config_G)\n",
    "pretrained_dict = torch.load(weights_path)\n",
    "model.load_state_dict(pretrained_dict)\n",
    "model.cuda();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DDKCo12_aRaJ"
   },
   "source": [
    "## Use model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WC83TpoXaRaJ"
   },
   "outputs": [],
   "source": [
    "img = transform(aligned_img).unsqueeze(0)\n",
    "\n",
    "with torch.no_grad():\n",
    "    out = model(img.cuda())\n",
    "\n",
    "out = util.tensor2im(out.data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8dCFBy_SaRaJ"
   },
   "source": [
    "## Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 389
    },
    "id": "o9WtUToyaRaK",
    "outputId": "acc6621e-711c-4544-d408-255f79090898"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(aligned_img)\n",
    "plt.axis('off')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(out)\n",
    "plt.axis('off')\n",
    "\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EYOOBqRhaRaK"
   },
   "outputs": [],
   "source": [
    "imageio.imsave('result.jpg', out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SWfIoGpBaRaL"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Model_infer.ipynb",
   "provenance": []
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}