QuickFaceSwap.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "QuickFaceSwap",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tg-bomze/collection-of-notebooks/blob/master/QuickFaceSwap.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jFf145OKhrZh"
      },
      "source": [
        "<b><font color=\"black\" size=\"+4\">QuickFaceSwap v.1.1</font></b>\n",
        "\n",
        "<b><font color=\"black\" size=\"+2\">Based on:</font></b>\n",
        "\n",
        "**GitHub repository**: [SimSwap](https://github.com/neuralchen/SimSwap), [GFPGAN](https://github.com/TencentARC/GFPGAN)\n",
        "\n",
        "Article: [SimSwap: An Efficient Framework For High Fidelity Face Swapping](https://arxiv.org/pdf/2106.06340v1.pdf), [Towards Real-World Blind Face Restoration with Generative Facial Prior\n",
        "](https://arxiv.org/pdf/2101.04061.pdf)\n",
        "\n",
        "**Creators of SimSwap:** *Renwang Chen, Xuanhong Chen, Bingbing Ni, Yanhao Ge*\n",
        "\n",
        "**Creators of GFPGAN:** *Xintao Wang and Yu Li and Honglun Zhang and Ying Shan*\n",
        "\n",
        "<b><font color=\"black\" size=\"+2\">Colab created by:</font></b>\n",
        "\n",
        "GitHub: [@tg-bomze](https://github.com/tg-bomze),\n",
        "Telegram: [@bomze](https://t.me/bomze),\n",
        "Twitter: [@tg_bomze](https://twitter.com/tg_bomze).\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "```\n",
        "(ENG) To get started, click on the button (where the red arrow indicates). After clicking, wait until the execution is complete.\n",
        "```\n",
        "```\n",
        "(RUS) Чтобы начать, поочередно нажимайте на кнопки (куда указывают красные стрелки), дожидаясь завершения выполнения каждого блока.\n",
        "```\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tr_2fdxirkkb",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Install all necessary libraries</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Установить все необходимые библиотеки</font></b>\n",
        "\n",
        "#@markdown **Follow this link (перейдите по ссылке):**\n",
        "\n",
        "#@markdown **https://drive.google.com/drive/folders/1jV6_0FIMPC53FZ2HzZNJZGMe55bbu17R**\n",
        "\n",
        "#@markdown **1) right click on 'checkpoints' (правой кнопкой по checkpoints)**\n",
        "\n",
        "#@markdown **2) select 'Add shortcut to Drive (выберите Сохранить ярлык на Диск)**\n",
        "\n",
        "#@markdown ![](https://github.com/tg-bomze/collection-of-notebooks/raw/master/dfs.png)\n",
        "\n",
        "#@markdown **3) run this block and follow the further instructions (после этого запустите блок и следуйте инструкции)**\n",
        "\n",
        "#@markdown *Attention! If the weights have already been saved, then run this block and just mount Google Drive.*\n",
        "\n",
        "#@markdown *Внимание! Если веса уже сохранены, то можете сразу запустить блок и смонтировать Google Drive.*\n",
        "\n",
        "from google.colab import drive\n",
        "from google.colab import files\n",
        "from IPython.display import clear_output\n",
        "import os\n",
        "drive.mount('/content/drive')\n",
        "!pip install ffmpeg-python\n",
        "!pip install facexlib\n",
        "\n",
        "%cd /content\n",
        "import torch\n",
        "!BASICSR_EXT=True pip install basicsr\n",
        "!git clone https://github.com/TencentARC/GFPGAN.git\n",
        "%cd /content/GFPGAN\n",
        "!mkdir MODNet\n",
        "!pip install -r requirements.txt\n",
        "!wget https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth -P experiments/pretrained_models\n",
        "!wget https://www.dropbox.com/s/0zwrook2mxxb4vc/modnet_photographic_portrait_matting.ckpt\n",
        "%cd /content/GFPGAN/MODNet\n",
        "!wget https://www.dropbox.com/s/4nwzaaq5x1tm2ze/MODNet.zip\n",
        "!unzip /content/GFPGAN/MODNet/MODNet.zip\n",
        "!rm -rf /content/GFPGAN/MODNet/MODNet.zip\n",
        "\n",
        "%cd /content\n",
        "!git clone https://github.com/neuralchen/SimSwap\n",
        "%cd /content/SimSwap\n",
        "!mkdir -p /usr/local/lib/python3.7/dist-packages/facexlib/weights\n",
        "!wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2\n",
        "!bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2\n",
        "!rm -rf shape_predictor_68_face_landmarks.dat.bz2\n",
        "!mkdir arcface_model checkpoints\n",
        "%cd /content/SimSwap/checkpoints\n",
        "!cp /content/drive/MyDrive/checkpoint/checkpoints.zip ./\n",
        "!unzip /content/SimSwap/checkpoints/checkpoints.zip\n",
        "%cd /content/SimSwap/arcface_model\n",
        "!cp /content/drive/MyDrive/checkpoint/arcface_checkpoint.tar ./\n",
        "%cd /content/SimSwap\n",
        "\n",
        "path_to_frames = '/content/frames'\n",
        "path_to_result = '/content/result'\n",
        "!rm -rf $path_to_frames $path_to_result /content/sample_data\n",
        "!mkdir $path_to_frames $path_to_result\n",
        "\n",
        "clear_output()"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z4ZiRzhY7BYP",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Upload video</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Загрузить видео</font></b>\n",
        "\n",
        "import os\n",
        "import cv2\n",
        "import ffmpeg\n",
        "\n",
        "uploaded = files.upload()\n",
        "for vp in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=vp, length=len(uploaded[vp])))\n",
        "video_path = f'/content/SimSwap/{vp}'\n",
        "\n",
        "!rm -rf /content/frames/*\n",
        "!ffmpeg -y -i $video_path /content/audio.mp3\n",
        "!ffmpeg -i $video_path /content/frames/%7d.png\n",
        "\n",
        "fps_of_video = int(cv2.VideoCapture(video_path).get(cv2.CAP_PROP_FPS))\n",
        "fps_of_video +=1\n",
        "\n",
        "if os.listdir('/content/frames'):\n",
        "  clear_output()\n",
        "  print('Done!')\n",
        "else: print('Error!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ozr3wZfzlqIk",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Upload photo with face</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Загрузить фото с лицом</font></b>\n",
        "uploaded = files.upload()\n",
        "for pp in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=pp, length=len(uploaded[pp])))\n",
        "clear_output()\n",
        "print('Done!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GLbOozniGi_H",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Prepare script</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Подготовить скрипт</font></b>\n",
        "%%writefile /content/SimSwap/test_video.py\n",
        "import torch\n",
        "import fractions\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torchvision import transforms\n",
        "import PIL\n",
        "from PIL import Image\n",
        "from tqdm import tqdm\n",
        "import warnings\n",
        "import numpy as np\n",
        "import imageio\n",
        "import glob\n",
        "import os\n",
        "import sys\n",
        "import cv2\n",
        "import scipy\n",
        "import scipy.ndimage\n",
        "import dlib\n",
        "\n",
        "size = 224\n",
        "path_to_frames = '/content/frames'\n",
        "path_to_result = '/content/result'\n",
        "\n",
        "\n",
        "sys.path.append(\"/content/GFPGAN\")\n",
        "os.chdir('/content/GFPGAN')\n",
        "from MODNet.models.modnet import MODNet\n",
        "from archs.gfpganv1_arch import GFPGANv1\n",
        "from basicsr.utils import img2tensor, imwrite, tensor2img\n",
        "from torchvision.transforms.functional import normalize\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "weight = \"/content/GFPGAN/modnet_photographic_portrait_matting.ckpt\"\n",
        "im_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
        "modnet = MODNet(backbone_pretrained=False)\n",
        "modnet = nn.DataParallel(modnet).cuda()\n",
        "modnet.load_state_dict(torch.load(weight))\n",
        "modnet.eval()\n",
        "\n",
        "\n",
        "print('Loading weights ...')\n",
        "gfpgan = GFPGANv1(\n",
        "        out_size=512,\n",
        "        num_style_feat=512,\n",
        "        channel_multiplier=1,\n",
        "        decoder_load_path=None,\n",
        "        fix_decoder=True,\n",
        "        num_mlp=8,\n",
        "        input_is_latent=True,\n",
        "        different_w=True,\n",
        "        narrow=1,\n",
        "        sft_half=True)\n",
        "gfpgan.to(device)\n",
        "checkpoint = torch.load('/content/GFPGAN/experiments/pretrained_models/GFPGANv1.pth', map_location=lambda storage, loc: storage)\n",
        "gfpgan.load_state_dict(checkpoint['params_ema'])\n",
        "gfpgan.eval()\n",
        "\n",
        "\n",
        "sys.path.append(\"/content/SimSwap\")\n",
        "os.chdir('/content/SimSwap')\n",
        "from models.models import create_model\n",
        "from options.test_options import TestOptions\n",
        "from facexlib.utils.face_restoration_helper import FaceRestoreHelper\n",
        "\n",
        "transformer = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
        "    ])\n",
        "\n",
        "transformer_Arcface = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
        "    ])\n",
        "\n",
        "detransformer = transforms.Compose([\n",
        "        transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),\n",
        "        transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])\n",
        "    ])\n",
        "\n",
        "opt = TestOptions().parse()\n",
        "path_to_face = opt.pic_a_path\n",
        "start_epoch, epoch_iter = 1, 0\n",
        "torch.nn.Module.dump_patches = True\n",
        "model = create_model(opt)\n",
        "model.eval()\n",
        "\n",
        "def post_process(color, orig):\n",
        "        orig_yuv = cv2.cvtColor(color, cv2.COLOR_BGR2YUV)\n",
        "        hires = np.copy(orig_yuv)\n",
        "        hires[:, :, 1:3] = orig[:, :, 1:3]\n",
        "        final = cv2.cvtColor(hires, cv2.COLOR_YUV2BGR)\n",
        "        return final\n",
        "\n",
        "def swapping(face_helper, img_path, latend_id, img_id, save_root):\n",
        "    img_name = os.path.basename(img_path)\n",
        "    basename, _ = os.path.splitext(img_name)\n",
        "    input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)\n",
        "    face_helper.clean_all()\n",
        "\n",
        "    face_helper.read_image(input_img)\n",
        "    face_helper.get_face_landmarks_5(only_center_face=True, pad_blur=False)\n",
        "    save_crop_path = os.path.join(save_root, 'cropped_faces', img_name)\n",
        "    face_helper.align_warp_face(save_crop_path)\n",
        "\n",
        "    for idx, cropped_face in enumerate(face_helper.cropped_faces):\n",
        "      small_cropped_face = cv2.resize(cropped_face, (size,size))\n",
        "      pil_small_cropped_face = Image.fromarray(small_cropped_face).convert('RGB')\n",
        "      img_b = transformer(pil_small_cropped_face)\n",
        "      img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])\n",
        "      img_att = img_att.cuda()\n",
        "\n",
        "      img_fake = model(img_id, img_att, latend_id, latend_id, True)\n",
        "      for i in range(img_id.shape[0]):\n",
        "          if i == 0:\n",
        "              row1 = img_id[i]\n",
        "              row2 = img_att[i]\n",
        "              row3 = img_fake[i]\n",
        "          else:\n",
        "              row1 = torch.cat([row1, img_id[i]], dim=2)\n",
        "              row2 = torch.cat([row2, img_att[i]], dim=2)\n",
        "              row3 = torch.cat([row3, img_fake[i]], dim=2)\n",
        "\n",
        "      full = row3.detach()\n",
        "      full = full.permute(1, 2, 0)\n",
        "      output = full.to('cpu')\n",
        "      output = np.array(output)\n",
        "      output = output[..., ::-1]\n",
        "      output = cv2.resize(output, (512,512))\n",
        "      yuv_swapped_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2YUV)\n",
        "\n",
        "      cropped_face_t = img2tensor(output, bgr2rgb=False, float32=True)\n",
        "      normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)\n",
        "      cropped_face_t = cropped_face_t.unsqueeze(0).to('cuda')\n",
        "\n",
        "      try:\n",
        "        with torch.no_grad():\n",
        "          output = gfpgan(cropped_face_t, return_rgb=False)[0]\n",
        "          restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))\n",
        "      except RuntimeError as error: restored_face = cropped_face\n",
        "\n",
        "      final_face = post_process(restored_face, yuv_swapped_face)\n",
        "\n",
        "      im = im_transform(final_face)\n",
        "      im = im[None, :, :, :]\n",
        "      im_b, im_c, im_h, im_w = im.shape\n",
        "      im_rw = im_w - im_w % 32\n",
        "      im_rh = im_h - im_h % 32\n",
        "      im = F.interpolate(im, size=(im_rh, im_rw), mode='area')\n",
        "      _, _, matte = modnet(im.cuda(), True)\n",
        "      matte = F.interpolate(matte, size=(im_h, im_w), mode='area')\n",
        "      matte = matte[0][0].data.cpu().numpy()\n",
        "      matte = np.repeat(np.asarray(matte)[:, :, None], 3, axis=2)\n",
        "      out = final_face.copy()\n",
        "      out[matte<0.5] = cropped_face[matte<0.5]\n",
        "\n",
        "\n",
        "      face_helper.add_restored_face(out)\n",
        "\n",
        "    face_helper.get_inverse_affine(None)\n",
        "    save_restore_path = os.path.join(save_root, 'restored_imgs', img_name)\n",
        "    face_helper.paste_faces_to_input_image(save_restore_path)\n",
        "\n",
        "source_helper = FaceRestoreHelper(face_size=512, upscale_factor=1, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png')\n",
        "\n",
        "source_helper.read_image(path_to_face)\n",
        "source_helper.get_face_landmarks_5(only_center_face=True, pad_blur=False)\n",
        "save_crop_path = os.path.join(path_to_result, 'cropped_faces', os.path.basename(path_to_face))\n",
        "source_helper.align_warp_face(save_crop_path)\n",
        "\n",
        "path_to_croped_src = '/content/result/cropped_faces/'\n",
        "path_to_cropped_src = path_to_croped_src+os.listdir(path_to_croped_src)[0]\n",
        "cropped_src = Image.open(path_to_cropped_src).convert('RGB')\n",
        "small_cropped_src = cropped_src.resize((size,size), Image.ANTIALIAS)\n",
        "\n",
        "img_a = transformer_Arcface(small_cropped_src)\n",
        "img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])\n",
        "img_id = img_id.cuda()\n",
        "img_id_downsample = F.interpolate(img_id, scale_factor=0.5)\n",
        "latend_id = model.netArc(img_id_downsample)\n",
        "latend_id = latend_id.detach().to('cpu')\n",
        "latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)\n",
        "latend_id = latend_id.to('cuda')\n",
        "\n",
        "face_helper = FaceRestoreHelper(face_size=512, upscale_factor=1, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png')\n",
        "img_list = sorted(glob.glob(os.path.join(path_to_frames, '*')))\n",
        "for img_path in tqdm(img_list):\n",
        "  swapping(face_helper, img_path, latend_id, img_id, path_to_result)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kdmbPNEZ32km",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Process</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Стартуем</font></b>\n",
        "path_to_face_photo = f'/content/SimSwap/{pp}'\n",
        "!rm -rf /content/result/cropped_faces/* /content/result/restored_imgs/* /content/result/cropped_faces/.ipynb_checkpoints\n",
        "!python test_video.py --Arc_path /content/drive/MyDrive/checkpoint/arcface_checkpoint.tar --pic_a_path $path_to_face_photo\n",
        "if os.listdir('/content/result/restored_imgs'):\n",
        "  clear_output()\n",
        "  print('Done!')\n",
        "else: print('Error!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BUin8V0vAirE",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Merge frames into video</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Объединяем кадры в видео</font></b>\n",
        "video_type = \"good for sharing\" #@param [\"good for sharing\", \"good for editing\"]\n",
        "if video_type == \"good for sharing\":\n",
        "  !ffmpeg -y -framerate $fps_of_video -pattern_type glob -i '/content/result/restored_imgs/*.png' -i /content/audio.mp3 -c:a copy -shortest -c:v libx264 -r $fps_of_video -pix_fmt yuv420p /content/swapped_without_audio.mp4\n",
        "  !ffmpeg -y -i /content/swapped_without_audio.mp4 -i /content/audio.mp3 /content/swapped.mp4\n",
        "else:\n",
        "  !ffmpeg -y -r $fps_of_video -i /content/result/restored_imgs/%7d.png -i /content/audio.mp3 /content/swapped.mp4\n",
        "clear_output()\n",
        "files.download('/content/swapped.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l30Kg3KIUZK4"
      },
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "**TODO:**\n",
        "\n",
        "\n",
        "`Done` - Decrease the brightness of the face square\n",
        "\n",
        "`Done` -   Make a round mask\n",
        "\n",
        "`In process` - Make the right Frames2Video\n",
        "\n"
      ]
    }
  ]
}