GHOST_inference.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VYMolZvRtrqx"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/lorenzinigiovanni/taaocv-project/blob/main/GHOST_inference.ipynb\" target=\"_parent\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>\n",
        "\n",
        "# GHOST: Generative High-fidelity One Shot Transfer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5tVMVEIWxycf"
      },
      "source": [
        "<img src=\"https://drive.google.com/uc?export=view&id=1MdOWxP9CqyqmW6t9MQ6-gcfEux54zbqP\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "34JxjB_1sifK",
        "outputId": "f6829e3d-50d3-45a7-a706-52f0e794c65a"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Check GPU and CUDA version**\n",
        "\n",
        "!nvidia-smi\n",
        "\n",
        "!nvcc --version"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nDRBmMbDiR6R",
        "outputId": "794b57b1-3219-46ec-bfc6-4ea355b9dab5"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Clone github & download models**\n",
        "\n",
        "!git clone https://github.com/sberbank-ai/sber-swap.git\n",
        "%cd sber-swap\n",
        "\n",
        "# load arcface\n",
        "!wget -P ./arcface_model https://github.com/sberbank-ai/sber-swap/releases/download/arcface/backbone.pth\n",
        "!wget -P ./arcface_model https://github.com/sberbank-ai/sber-swap/releases/download/arcface/iresnet.py\n",
        "\n",
        "# load landmarks detector\n",
        "!wget -P ./insightface_func/models/antelope https://github.com/sberbank-ai/sber-swap/releases/download/antelope/glintr100.onnx\n",
        "!wget -P ./insightface_func/models/antelope https://github.com/sberbank-ai/sber-swap/releases/download/antelope/scrfd_10g_bnkps.onnx\n",
        "\n",
        "# load model itself\n",
        "!wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/G_unet_2blocks.pth\n",
        "\n",
        "# load super res model\n",
        "!wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/super-res/10_net_G.pth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sL5OeF3IqBNa",
        "outputId": "02b0e5fd-2585-4cb6-8bf0-698fac75931c"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Install required libraries**\n",
        "\n",
        "!pip install mxnet-cu112\n",
        "!pip install onnxruntime-gpu==1.8\n",
        "!pip install insightface==0.2.1\n",
        "!pip install kornia==0.5.4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H6pX8e9JtBhW",
        "outputId": "a6412e19-f42a-4782-9866-f4f17947386b"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Preparation**\n",
        "\n",
        "import cv2\n",
        "import torch\n",
        "import time\n",
        "import os\n",
        "\n",
        "from utils.inference.image_processing import crop_face, get_final_image, show_images\n",
        "from utils.inference.video_processing import read_video, get_target, get_final_video, add_audio_from_another_video, face_enhancement\n",
        "from utils.inference.core import model_inference\n",
        "\n",
        "from network.AEI_Net import AEI_Net\n",
        "from coordinate_reg.image_infer import Handler\n",
        "from insightface_func.face_detect_crop_multi import Face_detect_crop\n",
        "from arcface_model.iresnet import iresnet100\n",
        "from models.pix2pix_model import Pix2PixModel\n",
        "from models.config_sr import TestOptions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-up7FWmYtDL4",
        "outputId": "a3cb3bc2-1a0b-404b-b3c5-e9e6e2013125"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Initialize models**\n",
        "\n",
        "app = Face_detect_crop(name='antelope', root='./insightface_func/models')\n",
        "app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))\n",
        "\n",
        "# main model for generation\n",
        "G = AEI_Net(backbone='unet', num_blocks=2, c_id=512)\n",
        "G.eval()\n",
        "G.load_state_dict(torch.load('weights/G_unet_2blocks.pth', map_location=torch.device('cpu')))\n",
        "G = G.cuda()\n",
        "G = G.half()\n",
        "\n",
        "# arcface model to get face embedding\n",
        "netArc = iresnet100(fp16=False)\n",
        "netArc.load_state_dict(torch.load('arcface_model/backbone.pth'))\n",
        "netArc=netArc.cuda()\n",
        "netArc.eval()\n",
        "\n",
        "# model to get face landmarks\n",
        "handler = Handler('./coordinate_reg/model/2d106det', 0, ctx_id=0, det_size=640)\n",
        "\n",
        "# model to make superres of face, set use_sr=True if you want to use super resolution or use_sr=False if you don't\n",
        "use_sr = True\n",
        "if use_sr:\n",
        "    os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
        "    torch.backends.cudnn.benchmark = True\n",
        "    opt = TestOptions()\n",
        "    #opt.which_epoch ='10_7'\n",
        "    model = Pix2PixModel(opt)\n",
        "    model.netG.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gRNla-XdtDO5",
        "outputId": "aad2889d-cf0c-4884-9a28-c2d8cae1a314"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Upload source image and video**\n",
        "\n",
        "#@markdown choose not really long videos, coz it can take a lot of time otherwise  \n",
        "\n",
        "#@markdown choose source image as a photo -- preferable a selfie of a person\n",
        "\n",
        "source_path = 'examples/images/elon_musk.jpg' #@param {type:\"string\"}\n",
        "path_to_video = 'examples/videos/nggyup.mp4' #@param {type:\"string\"}\n",
        "\n",
        "source_full = cv2.imread(source_path)\n",
        "OUT_VIDEO_NAME = \"examples/results/result.mp4\"\n",
        "crop_size = 224 # don't change this\n",
        "\n",
        "\n",
        "# check, if we can detect face on the source image\n",
        "\n",
        "try:    \n",
        "    source = crop_face(source_full, app, crop_size)[0]\n",
        "    source = [source[:, :, ::-1]]\n",
        "    print(\"Everything is ok!\")\n",
        "except TypeError:\n",
        "    print(\"Bad source images\")\n",
        "\n",
        "# read video\n",
        "full_frames, fps = read_video(path_to_video)\n",
        "target = get_target(full_frames, app, crop_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PzPhKk5PAQHe",
        "outputId": "375705be-e2e5-4051-880b-b7d90b755231"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Inference**\n",
        "\n",
        "\n",
        "batch_size =  40#@param {type:\"integer\"}\n",
        "\n",
        "START_TIME = time.time()\n",
        "\n",
        "final_frames_list, crop_frames_list, full_frames, tfm_array_list = model_inference(full_frames,\n",
        "                                                                                   source,\n",
        "                                                                                   target,\n",
        "                                                                                   netArc,\n",
        "                                                                                   G,\n",
        "                                                                                   app,\n",
        "                                                                                   set_target = False,\n",
        "                                                                                   crop_size=crop_size,\n",
        "                                                                                   BS=batch_size)\n",
        "\n",
        "if use_sr:\n",
        "    final_frames_list = face_enhancement(final_frames_list, model)\n",
        "\n",
        "get_final_video(final_frames_list,\n",
        "                crop_frames_list,\n",
        "                full_frames,\n",
        "                tfm_array_list,\n",
        "                OUT_VIDEO_NAME,\n",
        "                fps, \n",
        "                handler)\n",
        "\n",
        "add_audio_from_another_video(path_to_video, OUT_VIDEO_NAME, \"audio\")\n",
        "\n",
        "print(f'Full pipeline took {time.time() - START_TIME}')\n",
        "print(f\"Video saved with path {OUT_VIDEO_NAME}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 471
        },
        "id": "iqOjUHPPbVZI",
        "outputId": "f94f5a9f-097a-4d52-8bdd-96e737f8bfe0"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Visualize Video Swap**\n",
        "\n",
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "\n",
        "video_file = open(OUT_VIDEO_NAME, \"r+b\").read()\n",
        "video_url = f\"data:video/mp4;base64,{b64encode(video_file).decode()}\"\n",
        "\n",
        "HTML(f\"\"\"<video width={800} controls><source src=\"{video_url}\"></video>\"\"\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.9.13 ('base')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.13"
    },
    "vscode": {
      "interpreter": {
        "hash": "19d1d53a962d236aa061289c2ac16dc8e6d9648c89fe79f459ae9a3493bc67b4"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}