SimSwap_videos.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SimSwap_videos.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "0Qzzx2UpDkqw",
        "RVcsLmEwrUhO",
        "hBecZ0Trr9MD",
        "s6vGR7Ky9Pmf",
        "KSiwfjsy9Y_y"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MdUGtgKzpWf3"
      },
      "source": [
        "\n",
        "# SimSwap for videos\n",
        "\n",
        "Reference: [my changes to the official notebook](https://gist.github.com/woctezuma/78a98b73cbba8cba478d99c8c50bc359)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Qzzx2UpDkqw"
      },
      "source": [
        "## Prepare code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VA_4CeWZCHLP"
      },
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/woctezuma/SimSwap.git\n",
        "%cd /content/SimSwap/\n",
        "!git checkout upgrade-insightface"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y5K4au_UCkKn"
      },
      "source": [
        "!pip install insightface==0.7.3 onnxruntime moviepy > /dev/null\n",
        "!pip install googledrivedownloader > /dev/null\n",
        "!pip install imageio==2.34.0 > /dev/null"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RVcsLmEwrUhO"
      },
      "source": [
        "## Prepare models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ATb7G4P8qC4X"
      },
      "source": [
        "%cd /content/SimSwap"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gLti1J0pEFjJ"
      },
      "source": [
        "from google_drive_downloader import GoogleDriveDownloader\n",
        "\n",
        "!wget -P ./arcface_model https://github.com/woctezuma/SimSwap-colab/releases/download/1.0/arcface_checkpoint.tar\n",
        "!wget https://github.com/neuralchen/SimSwap/releases/download/1.0/checkpoints.zip\n",
        "!wget -P ./parsing_model/checkpoint https://github.com/neuralchen/SimSwap/releases/download/1.0/79999_iter.pth\n",
        "!wget https://github.com/neuralchen/SimSwap/releases/download/512_beta/512.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aSRnK5V4HI-k"
      },
      "source": [
        "!wget --no-check-certificate \\\n",
        " https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w \\\n",
        " -O antelope.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q-pyAF6lrbE2"
      },
      "source": [
        "!unzip ./checkpoints.zip  -d ./checkpoints\n",
        "\n",
        "!unzip 512.zip -d ./checkpoints\n",
        "\n",
        "!unzip antelope.zip -d ./insightface_func/models/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jRJt2CxBrfZx"
      },
      "source": [
        "## Prepare data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Qr3ohTtrmxX"
      },
      "source": [
        "### Download"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zGAEo4uprGOK"
      },
      "source": [
        "%cd /content\n",
        "\n",
        "!wget https://i.imgur.com/iQtmj1N.png -O photo.png\n",
        "!wget http://i.imgur.com/yCD2Uxm.gif -O video.gif"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9brzTldrrAW"
      },
      "source": [
        "input_image_fname = '/content/photo.png'\n",
        "input_video_fname = '/content/video.gif'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BsGmIMxLVxyO"
      },
      "source": [
        "## Run"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBecZ0Trr9MD"
      },
      "source": [
        "### Official code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QL9fX8Fl588y"
      },
      "source": [
        "%cd /content/SimSwap/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PfSsND36EMvn"
      },
      "source": [
        "import cv2\n",
        "import torch\n",
        "import fractions\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import torch.nn.functional as F\n",
        "from torchvision import transforms\n",
        "from models.models import create_model\n",
        "from options.test_options import TestOptions\n",
        "from insightface_func.face_detect_crop_multi import Face_detect_crop\n",
        "from util.videoswap import video_swap\n",
        "from util.add_watermark import watermark_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rxSbZ2EDNDlf"
      },
      "source": [
        "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",
        "    ])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lyhc4v8Hud-1"
      },
      "source": [
        "# If the algorithm misses some faces, you could lower the detection threshold.\n",
        "# Reference: https://github.com/neuralchen/SimSwap/issues/39#issuecomment-873758730\n",
        "\n",
        "det_thresh = 0.6\n",
        "\n",
        "# You could also decrease the image size used for face detection:\n",
        "\n",
        "det_size = (640,640)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Either 224 (okay) or 512 (experimental)\n",
        "crop_size = 224"
      ],
      "metadata": {
        "id": "BGyMD6HB0IqL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wwJOwR9LNKRz"
      },
      "source": [
        "start_epoch, epoch_iter = 1, 0\n",
        "\n",
        "opt = TestOptions()\n",
        "opt.initialize()\n",
        "opt.parser.add_argument('-f') ## dummy arg to avoid bug\n",
        "opt = opt.parse()\n",
        "opt.pic_a_path = input_image_fname ## or replace it with image from your own google drive\n",
        "opt.video_path = input_video_fname ## or replace it with video from your own google drive\n",
        "opt.output_path = '/content/output.mp4'\n",
        "opt.temp_path = './tmp'\n",
        "opt.Arc_path = './arcface_model/arcface_checkpoint.tar'\n",
        "opt.isTrain = False\n",
        "opt.no_simswaplogo = True\n",
        "opt.use_mask = True\n",
        "opt.crop_size = crop_size\n",
        "\n",
        "if crop_size == 512:\n",
        "    opt.which_epoch = 550000\n",
        "    opt.name = '512'\n",
        "    mode = 'ffhq'\n",
        "else:\n",
        "    mode = 'None'\n",
        "\n",
        "torch.nn.Module.dump_patches = True\n",
        "model = create_model(opt)\n",
        "model.eval()\n",
        "\n",
        "\n",
        "app = Face_detect_crop(name='antelope', root='./insightface_func/models')\n",
        "app.prepare(ctx_id= 0, det_thresh=det_thresh, det_size=det_size, mode=mode)\n",
        "\n",
        "pic_a = opt.pic_a_path\n",
        "# img_a = Image.open(pic_a).convert('RGB')\n",
        "img_a_whole = cv2.imread(pic_a)\n",
        "img_a_align_crop, _ = app.get(img_a_whole,crop_size)\n",
        "img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) \n",
        "img_a = transformer_Arcface(img_a_align_crop_pil)\n",
        "img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])\n",
        "\n",
        "# convert numpy to tensor\n",
        "img_id = img_id.cuda()\n",
        "\n",
        "#create latent id\n",
        "img_id_downsample = F.interpolate(img_id, size=(112,112))\n",
        "latend_id = model.netArc(img_id_downsample)\n",
        "latend_id = F.normalize(latend_id, p=2, dim=1)\n",
        "\n",
        "try:\n",
        "  video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\\\n",
        "             no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask,crop_size=crop_size)\n",
        "except IndexError:\n",
        "  print('[error] This is most likely due to the absence of audio from a GIF input.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tJflBIWVr_Z3"
      },
      "source": [
        "### My fix for GIF"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "myqMMla8pLgf"
      },
      "source": [
        "If you want to apply SimSwap to a GIF, there will be an error because the input video has no audio.\n",
        "\n",
        "To fix this issue, aggregate the temporary output by yourself by running:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kiOBpF_T66m5"
      },
      "source": [
        "import cv2\n",
        "\n",
        "def get_fps(video_path):\n",
        "  video = cv2.VideoCapture(video_path)\n",
        "  fps = video.get(cv2.CAP_PROP_FPS)\n",
        "\n",
        "  return fps"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dj3gHzQyj-GX"
      },
      "source": [
        "import os \n",
        "import glob\n",
        "from moviepy.video.io.ImageSequenceClip import ImageSequenceClip\n",
        "\n",
        "def collate_into_gif(temp_results_dir, output_fname, fps):\n",
        "  path = os.path.join(temp_results_dir,'*.jpg')\n",
        "  image_filenames = sorted(glob.glob(path))\n",
        "\n",
        "  clips = ImageSequenceClip(image_filenames,fps = fps)\n",
        "  clips.write_gif(output_fname)\n",
        "\n",
        "  return"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_7_0Af-V6yic"
      },
      "source": [
        "temp_results_dir = '/content/SimSwap/tmp/'\n",
        "output_fname = '/content/output.gif'\n",
        "\n",
        "collate_into_gif(temp_results_dir,\n",
        "                 output_fname,\n",
        "                 fps=get_fps(input_video_fname))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "el1eu0gmpNsB"
      },
      "source": [
        "To optimize the file size (in MB) of the GIF, you can upload it to a website like https://ezgif.com/optimize"
      ]
    }
  ]
}