{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3CCSQIBMAYYZ"
      },
      "source": [
        "# FIFA Virtual Try-On Inference ⚽\n",
        "\n",
        "This notebook generates a synthetic image of a person wearing a target clothing. It requires an image of a person and a target clothing as inputs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7WJtP2PfBcPN"
      },
      "source": [
        "### Setup dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GVm5QFBMDBbT"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/hasibzunair/fifa-demo.git\n",
        "%cd fifa-demo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cnk7syY0rPKp"
      },
      "outputs": [],
      "source": [
        "!pip install ninja\n",
        "!pip install pymatting==1.1.5\n",
        "!pip install onnxruntime==1.10.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xQwI--uhoH6R"
      },
      "outputs": [],
      "source": [
        "import gdown\n",
        "import numpy as np\n",
        "from PIL import Image, ImageOps\n",
        "import IPython\n",
        "import gdown\n",
        "import os\n",
        "import sys\n",
        "\n",
        "from rembg import remove\n",
        "from predict_pose import generate_pose_keypoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "86Ll7LIW1kGw"
      },
      "outputs": [],
      "source": [
        "# Create dummy folders same as train/test data dir tree\n",
        "\n",
        "!mkdir Data_preprocessing/test_color\n",
        "!mkdir Data_preprocessing/test_colormask\n",
        "!mkdir Data_preprocessing/test_edge\n",
        "!mkdir Data_preprocessing/test_img\n",
        "!mkdir Data_preprocessing/test_label\n",
        "!mkdir Data_preprocessing/test_mask\n",
        "!mkdir Data_preprocessing/test_pose\n",
        "!mkdir inputs\n",
        "!mkdir inputs/img\n",
        "!mkdir inputs/cloth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K8hYM6XqCnxC"
      },
      "outputs": [],
      "source": [
        "# Get pose model\n",
        "\n",
        "%cd pose\n",
        "!wget https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/pose_deploy_linevec.prototxt\n",
        "!wget https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/pose_iter_440000.caffemodel\n",
        "%cd .."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HbGDB31KrKHV"
      },
      "outputs": [],
      "source": [
        "# Get parser and segmentation model\n",
        "\n",
        "!git clone https://github.com/hasibzunair/Self-Correction-Human-Parsing-for-ACGPN.git\n",
        "!git clone https://github.com/hasibzunair/U-2-Net.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3hSJI347rZtQ"
      },
      "outputs": [],
      "source": [
        "# For segmentation mask generation\n",
        "\n",
        "!wget https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/lip_final.pth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ooARXEZZYnGD"
      },
      "outputs": [],
      "source": [
        "# Get U-2-Net weights\n",
        "\n",
        "%cd U-2-Net\n",
        "!mkdir saved_models\n",
        "!mkdir saved_models/u2net\n",
        "!mkdir saved_models/u2netp\n",
        "!wget -P saved_models/u2netp/ https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/u2netp.pth\n",
        "!wget -P saved_models/u2net/ https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/u2net.pth\n",
        "import u2net_load\n",
        "import u2net_run\n",
        "u2net = u2net_load.model(model_name = 'u2netp')\n",
        "%cd .."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r1VknOqswSTW"
      },
      "outputs": [],
      "source": [
        "# Get model checkpoints\n",
        "\n",
        "!mkdir checkpoints\n",
        "%cd checkpoints\n",
        "!wget https://github.com/hasibzunair/fifa-tryon/releases/download/v1.0-models/fifa_viton.zip\n",
        "!unzip fifa_viton.zip\n",
        "%cd .."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eD-DMczQthkd"
      },
      "source": [
        "### Please upload your cloth image below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eAMmTiQ7zibP"
      },
      "outputs": [],
      "source": [
        "%cd inputs/cloth\n",
        "from google.colab import files\n",
        "uploaded = files.upload()\n",
        "%cd ..\n",
        "%cd .."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dQP04cGctrOW"
      },
      "source": [
        "### Please upload your person image below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_wmUOR92170v"
      },
      "outputs": [],
      "source": [
        "%cd inputs/img\n",
        "from google.colab import files\n",
        "uploaded = files.upload()\n",
        "%cd ..\n",
        "%cd .."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u-zkuJWVWgrX"
      },
      "source": [
        "### Preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tm-9Up6z0Zpg"
      },
      "outputs": [],
      "source": [
        "cloth_name = 'cloth.png'\n",
        "cloth_path = os.path.join('inputs/cloth', sorted(os.listdir('inputs/cloth'))[0])\n",
        "cloth = Image.open(cloth_path)\n",
        "\n",
        "# Resize cloth image\n",
        "cloth = ImageOps.fit(cloth, (192, 256), Image.BICUBIC).convert(\"RGB\")\n",
        "\n",
        "# Save resized cloth image\n",
        "cloth.save(os.path.join('Data_preprocessing/test_color', cloth_name))\n",
        "\n",
        "# 1. Get binary mask for clothing image\n",
        "u2net_run.infer(u2net, 'Data_preprocessing/test_color', 'Data_preprocessing/test_edge')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y3aHah45D655"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "start_time = time.time()\n",
        "\n",
        "# Remove background from person image\n",
        "remove_bg = False\n",
        "# Person image\n",
        "img_name = 'person.png'\n",
        "img_path = os.path.join('inputs/img', sorted(os.listdir('inputs/img'))[0])\n",
        "img = Image.open(img_path)\n",
        "if remove_bg:\n",
        "  # Remove background\n",
        "  img = remove(img, alpha_matting=True, alpha_matting_erode_size=15)\n",
        "  print(\"Removing background from person image..\")\n",
        "img = ImageOps.fit(img, (192, 256), Image.BICUBIC).convert(\"RGB\")\n",
        "# Get binary from person image\n",
        "img_mask = remove(img, alpha_matting=True, alpha_matting_erode_size=15, only_mask=True)\n",
        "img_path = os.path.join('Data_preprocessing/test_img', img_name)\n",
        "img.save(img_path)\n",
        "resize_time = time.time()\n",
        "print('Resized image in {}s'.format(resize_time-start_time))\n",
        "\n",
        "# 2. Get parsed person image (test_label), uses person image\n",
        "!python3 Self-Correction-Human-Parsing-for-ACGPN/simple_extractor.py --dataset 'lip' --model-restore 'lip_final.pth' --input-dir 'Data_preprocessing/test_img' --output-dir 'Data_preprocessing/test_label'\n",
        "parse_time = time.time()\n",
        "print('Parsing generated in {}s'.format(parse_time-resize_time))\n",
        "\n",
        "# 3. Get pose map from person image\n",
        "pose_path = os.path.join('Data_preprocessing/test_pose', img_name.replace('.png', '_keypoints.json'))\n",
        "generate_pose_keypoints(img_path, pose_path)\n",
        "pose_time = time.time()\n",
        "print('Pose map generated in {}s'.format(pose_time-parse_time))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EgMi912KAUNs"
      },
      "outputs": [],
      "source": [
        "!rm -rf Data_preprocessing/test_pairs.txt\n",
        "\n",
        "# Format: person, cloth image\n",
        "with open('Data_preprocessing/test_pairs.txt','w') as f:\n",
        "    f.write('person.png cloth.png')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NTfjVq44X9dg"
      },
      "source": [
        "### Try it on!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EIcoP4ll14Ia"
      },
      "outputs": [],
      "source": [
        "# Run test.py using the preferred configuration (e.g. changes in architecture etc.)\n",
        "!python test.py --name fifa_viton"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_vfYoKcQbRM0"
      },
      "source": [
        "### Post processing\n",
        "\n",
        "Put background from the input image back into the tryon output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7DM4n7OdhXas"
      },
      "outputs": [],
      "source": [
        "def composite_background(person_image_path, tryon_image_path):\n",
        "  \"\"\"Put background back on the person image after tryon.\"\"\"\n",
        "  person = np.array(Image.open(person_image_path))\n",
        "  # tryon image\n",
        "  tryon = np.array(Image.open(tryon_image_path))\n",
        "  # persom image mask from rembg\n",
        "  p_mask = np.array(img_mask)\n",
        "  # make binary mask\n",
        "  p_mask = np.where(p_mask>0, 1, 0)\n",
        "  # invert mask\n",
        "  p_mask_inv = np.logical_not(p_mask)\n",
        "  # make bg without person\n",
        "  background = person * np.stack((p_mask_inv, p_mask_inv, p_mask_inv), axis=2)\n",
        "  # make tryon image without background\n",
        "  tryon_nobg = tryon * np.stack((p_mask, p_mask, p_mask), axis=2)\n",
        "  tryon_nobg = tryon_nobg.astype(\"uint8\")\n",
        "  # composite \n",
        "  tryon_with_bg = np.add(tryon_nobg, background)\n",
        "  tryon_with_bg_pil = Image.fromarray(np.uint8(tryon_with_bg)).convert('RGB')\n",
        "  tryon_with_bg_pil.save(\"results/test/try-on/tryon_with_bg.png\")\n",
        "\n",
        "\n",
        "composite_background('Data_preprocessing/test_img/person.png',\n",
        "                     'results/test/try-on/person.png')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0IfZakpL1xX9"
      },
      "source": [
        "### Show results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e7gM3vwWZHGj"
      },
      "outputs": [],
      "source": [
        "# See 'results' directory\n",
        "output_grid = np.concatenate([np.array(Image.open('Data_preprocessing/test_img/person.png')),\n",
        "                np.array(Image.open('Data_preprocessing/test_color/cloth.png')),\n",
        "                np.array(Image.open('results/test/try-on/person.png')),\n",
        "                np.array(Image.open('results/test/try-on/tryon_with_bg.png'))], axis=1)\n",
        "image_grid = Image.fromarray(output_grid)\n",
        "image_grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V_zzoKdCOTNn"
      },
      "source": [
        "**To try a new person and cloth pair, remove current files and start again by uploading a new cloth image!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9b-Fx2TyKbZ9"
      },
      "outputs": [],
      "source": [
        "!rm -rf Data_preprocessing/test_color/*\n",
        "!rm -rf Data_preprocessing/test_colormask/*\n",
        "!rm -rf Data_preprocessing/test_edge/*\n",
        "!rm -rf Data_preprocessing/test_img/*\n",
        "!rm -rf Data_preprocessing/test_label/*\n",
        "!rm -rf Data_preprocessing/test_mask/*\n",
        "!rm -rf Data_preprocessing/test_pose/*\n",
        "!rm -rf inputs/cloth/*\n",
        "!rm -rf inputs/img/*\n",
        "!rm -rf results/*\n",
        "%cd /content/fifa-demo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0wpSoOskH2qq"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "7WJtP2PfBcPN",
        "u-zkuJWVWgrX",
        "NTfjVq44X9dg"
      ],
      "name": "demo.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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}