toonify.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "inference_playground.ipynb",
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "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.3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/eyaler/avatars4all/blob/master/toonify.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uuviq3qQkUFy"
      },
      "source": [
        "import os\n",
        "os.chdir('/content')\n",
        "CODE_DIR = 'restyle-encoder'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QQ6XEmlHlXbk",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e66b792b-0bba-456c-df58-396a7a47a314"
      },
      "source": [
        "!git clone https://github.com/yuval-alaluf/restyle-encoder.git $CODE_DIR"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'restyle-encoder'...\n",
            "remote: Enumerating objects: 186, done.\u001b[K\n",
            "remote: Counting objects: 100% (186/186), done.\u001b[K\n",
            "remote: Compressing objects: 100% (153/153), done.\u001b[K\n",
            "remote: Total 186 (delta 42), reused 162 (delta 29), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (186/186), 24.78 MiB | 21.25 MiB/s, done.\n",
            "Resolving deltas: 100% (42/42), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JaRUFuVHkzye",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "47c6f353-69eb-436b-c9df-a69a11a21d6a"
      },
      "source": [
        "!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
        "!sudo unzip ninja-linux.zip -d /usr/local/bin/\n",
        "!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2021-05-18 19:51:17--  https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
            "Resolving github.com (github.com)... 52.69.186.44\n",
            "Connecting to github.com (github.com)|52.69.186.44|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://github-releases.githubusercontent.com/1335132/d2f252e2-9801-11e7-9fbf-bc7b4e4b5c83?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210518%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210518T195117Z&X-Amz-Expires=300&X-Amz-Signature=16676b904f13c3399bac5989c306a3944a1bb9ce67d5386caaa4aa716318e2b2&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=1335132&response-content-disposition=attachment%3B%20filename%3Dninja-linux.zip&response-content-type=application%2Foctet-stream [following]\n",
            "--2021-05-18 19:51:18--  https://github-releases.githubusercontent.com/1335132/d2f252e2-9801-11e7-9fbf-bc7b4e4b5c83?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210518%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210518T195117Z&X-Amz-Expires=300&X-Amz-Signature=16676b904f13c3399bac5989c306a3944a1bb9ce67d5386caaa4aa716318e2b2&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=1335132&response-content-disposition=attachment%3B%20filename%3Dninja-linux.zip&response-content-type=application%2Foctet-stream\n",
            "Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.108.154, 185.199.110.154, 185.199.111.154, ...\n",
            "Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.108.154|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 77854 (76K) [application/octet-stream]\n",
            "Saving to: ‘ninja-linux.zip’\n",
            "\n",
            "ninja-linux.zip     100%[===================>]  76.03K  --.-KB/s    in 0.01s   \n",
            "\n",
            "2021-05-18 19:51:18 (5.00 MB/s) - ‘ninja-linux.zip’ saved [77854/77854]\n",
            "\n",
            "Archive:  ninja-linux.zip\n",
            "  inflating: /usr/local/bin/ninja    \n",
            "update-alternatives: using /usr/local/bin/ninja to provide /usr/bin/ninja (ninja) in auto mode\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "23baccYQlU9E"
      },
      "source": [
        "os.chdir(f'./{CODE_DIR}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d13v7In0kTJn"
      },
      "source": [
        "from argparse import Namespace\n",
        "import time\n",
        "import os\n",
        "import sys\n",
        "import pprint\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import torch\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "sys.path.append(\".\")\n",
        "sys.path.append(\"..\")\n",
        "\n",
        "from utils.common import tensor2im\n",
        "from models.psp import pSp\n",
        "from models.e4e import e4e\n",
        "\n",
        "%load_ext autoreload\n",
        "%autoreload 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRjtz6uLkTJs"
      },
      "source": [
        "## Step 1: Select Experiment Type\n",
        "Select which experiment you wish to perform inference on:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XESWAO65kTJt"
      },
      "source": [
        "#@title Select which experiment you wish to perform inference on: { run: \"auto\" }\n",
        "experiment_type = 'toonify' #@param ['ffhq_encode', 'cars_encode', 'church_encode', 'horse_encode', 'afhq_wild_encode', 'toonify']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4etDz82xkTJz"
      },
      "source": [
        "## Step 2: Prepare to Download Pretrained Models \n",
        "As part of this repository, we provide pretrained models for each of the above experiments. Here, we'll create the download command needed for downloading the desired model.\n",
        "\n",
        "Note: in this notebook, we'll be using ReStyle applied over pSp for all domains except for the horses domain where we'll be using e4e. This is done since e4e is generally able to generate more realistic reconstructions on this domain. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KSnjlBZOkTJ0"
      },
      "source": [
        "def get_download_model_command(file_id, file_name):\n",
        "    \"\"\" Get wget download command for downloading the desired model and save to directory ../pretrained_models. \"\"\"\n",
        "    current_directory = os.getcwd()\n",
        "    save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR, \"pretrained_models\")\n",
        "    if not os.path.exists(save_path):\n",
        "        os.makedirs(save_path)\n",
        "    url = r\"\"\"wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id={FILE_ID}' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id={FILE_ID}\" -O {SAVE_PATH}/{FILE_NAME} && rm -rf /tmp/cookies.txt\"\"\".format(FILE_ID=file_id, FILE_NAME=file_name, SAVE_PATH=save_path)\n",
        "    return url    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m4sjldFMkTJ5"
      },
      "source": [
        "MODEL_PATHS = {\n",
        "    \"ffhq_encode\": {\"id\": \"1sw6I2lRIB0MpuJkpc8F5BJiSZrc0hjfE\", \"name\": \"restyle_psp_ffhq_encode.pt\"},\n",
        "    \"cars_encode\": {\"id\": \"1zJHqHRQ8NOnVohVVCGbeYMMr6PDhRpPR\", \"name\": \"restyle_psp_cars_encode.pt\"},\n",
        "    \"church_encode\": {\"id\": \"1bcxx7mw-1z7dzbJI_z7oGpWG1oQAvMaD\", \"name\": \"restyle_psp_church_encode.pt\"},\n",
        "    \"horse_encode\": {\"id\": \"19_sUpTYtJmhSAolKLm3VgI-ptYqd-hgY\", \"name\": \"restyle_e4e_horse_encode.pt\"},\n",
        "    \"afhq_wild_encode\": {\"id\": \"1GyFXVTNDUw3IIGHmGS71ChhJ1Rmslhk7\", \"name\": \"restyle_psp_afhq_wild_encode.pt\"},\n",
        "    \"toonify\": {\"id\": \"1GtudVDig59d4HJ_8bGEniz5huaTSGO_0\", \"name\": \"restyle_psp_toonify.pt\"}\n",
        "}\n",
        "\n",
        "path = MODEL_PATHS[experiment_type]\n",
        "download_command = get_download_model_command(file_id=path[\"id\"], file_name=path[\"name\"]) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Tozsg81kTKA"
      },
      "source": [
        "## Step 3: Define Inference Parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XIhyc7RqkTKB"
      },
      "source": [
        "Below we have a dictionary defining parameters such as the path to the pretrained model to use and the path to the image to perform inference on.  \n",
        "While we provide default values to run this script, feel free to change as needed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2kE5y1-skTKC"
      },
      "source": [
        "EXPERIMENT_DATA_ARGS = {\n",
        "    \"ffhq_encode\": {\n",
        "        \"model_path\": \"pretrained_models/restyle_psp_ffhq_encode.pt\",\n",
        "        \"image_path\": \"notebooks/images/face_img.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "    \"cars_encode\": {\n",
        "        \"model_path\": \"pretrained_models/restyle_psp_cars_encode.pt\",\n",
        "        \"image_path\": \"notebooks/images/car_img.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((192, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "    \"church_encode\": {\n",
        "        \"model_path\": \"pretrained_models/restyle_psp_church_encode.pt\",\n",
        "        \"image_path\": \"notebooks/images/church_img.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "    \"horse_encode\": {\n",
        "        \"model_path\": \"pretrained_models/restyle_e4e_horse_encode.pt\",\n",
        "        \"image_path\": \"notebooks/images/horse_img.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "    \"afhq_wild_encode\": {\n",
        "        \"model_path\": \"pretrained_models/restyle_psp_afhq_wild_encode.pt\",\n",
        "        \"image_path\": \"notebooks/images/afhq_wild_img.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "    \"toonify\": {\n",
        "        \"model_path\": \"pretrained_models/restyle_psp_toonify.pt\",\n",
        "        \"image_path\": \"notebooks/images/toonify_img.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IzUHoD9ukTKG"
      },
      "source": [
        "EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ov6h6gkeZ0VI"
      },
      "source": [
        "To reduce the number of requests to fetch the model, we'll check if the model was previously downloaded and saved before downloading the model.  \n",
        "We'll download the model for the selected experiment and save it to the folder `../pretrained_models`.\n",
        "\n",
        "We also need to verify that the model was downloaded correctly. All of our models should weigh approximately 800MB - 1GB.  \n",
        "Note that if the file weighs several KBs, you most likely encounter a \"quota exceeded\" error from Google Drive. In that case, you should try downloading the model again after a few hours."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jQ31J_m7kTJ8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "08f7dfe6-f332-4ae9-9139-bb93fda38f73"
      },
      "source": [
        "if not os.path.exists(EXPERIMENT_ARGS['model_path']) or os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:\n",
        "    print(f'Downloading ReStyle model for {experiment_type}...')\n",
        "    os.system(f\"wget {download_command}\")\n",
        "    # if google drive receives too many requests, we'll reach the quota limit and be unable to download the model\n",
        "    if os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:\n",
        "        model_name = EXPERIMENT_ARGS['model_path'].split('/')[-1]\n",
        "        !wget --no-check-certificate -nc https://eyalgruss.com/fomm/$model_name -O /content/restyle-encoder/pretrained_models/$model_name\n",
        "else:\n",
        "    print(f'ReStyle model for {experiment_type} already exists!')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading ReStyle model for toonify...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TAWrUehTkTKJ"
      },
      "source": [
        "## Step 4: Load Pretrained Model\n",
        "We assume that you have downloaded all relevant models and placed them in the directory defined by the above dictionary."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1t-AOhP1kTKJ"
      },
      "source": [
        "model_path = EXPERIMENT_ARGS['model_path']\n",
        "ckpt = torch.load(model_path, map_location='cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2UBwJ3dJkTKM",
        "outputId": "4a686dcd-8120-4515-8cbe-bdc1ba08e04c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "opts = ckpt['opts']\n",
        "pprint.pprint(opts)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "{'batch_size': 8,\n",
            " 'board_interval': 50,\n",
            " 'checkpoint_path': '',\n",
            " 'dataset_type': 'ffhq_encode',\n",
            " 'device': 'cuda:0',\n",
            " 'encoder_type': 'BackboneEncoder',\n",
            " 'exp_dir': '',\n",
            " 'id_lambda': 1.0,\n",
            " 'image_interval': 100,\n",
            " 'input_nc': 6,\n",
            " 'l2_lambda': 1.0,\n",
            " 'learning_rate': 0.0001,\n",
            " 'lpips_lambda': 0.8,\n",
            " 'max_steps': 500000,\n",
            " 'moco_lambda': 0,\n",
            " 'n_iters_per_batch': 5,\n",
            " 'optim_name': 'ranger',\n",
            " 'output_size': 1024,\n",
            " 'save_interval': 1000,\n",
            " 'start_from_latent_avg': True,\n",
            " 'stylegan_weights': '',\n",
            " 'test_batch_size': 8,\n",
            " 'test_workers': 8,\n",
            " 'train_decoder': False,\n",
            " 'val_interval': 1000,\n",
            " 'w_norm_lambda': 0.025,\n",
            " 'workers': 8}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EMKhWoFKkTKS"
      },
      "source": [
        "# update the training options\n",
        "opts['checkpoint_path'] = model_path"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6hccfNizkTKW",
        "outputId": "81a4513d-ebe6-4746-dce6-9f7593a33f39",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "opts = Namespace(**opts)\n",
        "if experiment_type == 'horse_encode': \n",
        "    net = e4e(opts)\n",
        "else:\n",
        "    net = pSp(opts)\n",
        "    \n",
        "net.eval()\n",
        "net.cuda()\n",
        "print('Model successfully loaded!')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loading ReStyle pSp from checkpoint: pretrained_models/restyle_psp_toonify.pt\n",
            "Model successfully loaded!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_nJiSjSLsTPD",
        "outputId": "be4fc5b4-2708-4afa-c6e7-29cf9c9ab444",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "#@title Get the image/video from the web\n",
        "#@markdown 1. You can change the URLs to your **own** stuff!\n",
        "#@markdown 2. Alternatively, you can upload **local** files in the next cells\n",
        "\n",
        "foreground_url = 'https://storage.googleapis.com/startup-experts-directory/Eyal.Gruss.jpg' #@param {type:\"string\"}\n",
        "\n",
        "!pip install -U git+https://github.com/ytdl-org/youtube-dl\n",
        "import os\n",
        "import youtube_dl\n",
        "def is_supported(url):\n",
        "    if url.lower().endswith(('.png','.jpg','.jpeg','.bmp')):\n",
        "      return False\n",
        "    extractors = youtube_dl.extractor.gen_extractors()\n",
        "    for e in extractors:\n",
        "        if e.suitable(url) and e.IE_NAME != 'generic':\n",
        "            return True\n",
        "    return False\n",
        "\n",
        "if foreground_url:\n",
        "  !rm -f /content/foreground\n",
        "  if is_supported(foreground_url):\n",
        "    !rm -f /content/foreground.mp4\n",
        "    !youtube-dl --no-playlist -f \"bestvideo[ext=mp4][vcodec!*=av01]+bestaudio[ext=m4a]/mp4\" \"$foreground_url\" --merge-output-format mp4 -o /content/foreground\n",
        "    !mv /content/foreground.mp4 /content/foreground \n",
        "    fg_time_params = ''\n",
        "  if not os.path.exists('/content/foreground'):\n",
        "    !wget \"$foreground_url\" -O /content/foreground"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: youtube_dl in /usr/local/lib/python3.7/dist-packages (2021.5.16)\n",
            "--2021-05-18 20:17:47--  https://storage.googleapis.com/startup-experts-directory/Eyal.Gruss.jpg\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 74.125.23.128, 74.125.203.128, 74.125.204.128, ...\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|74.125.23.128|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 551080 (538K) [image/jpeg]\n",
            "Saving to: ‘/content/foreground’\n",
            "\n",
            "/content/foreground 100%[===================>] 538.16K  --.-KB/s    in 0.004s  \n",
            "\n",
            "2021-05-18 20:17:47 (143 MB/s) - ‘/content/foreground’ saved [551080/551080]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KeVURhansjmQ"
      },
      "source": [
        "#@title Optionally upload foreground image/video { run: \"auto\" }\n",
        "manually_upload_foreground = False #@param {type:\"boolean\"}\n",
        "if manually_upload_foreground:\n",
        "  from google.colab import files\n",
        "  import shutil\n",
        "\n",
        "  %cd /content/sample_data\n",
        "  try:\n",
        "    uploaded = files.upload()\n",
        "  except Exception as e:\n",
        "    %cd /content\n",
        "    raise e\n",
        "\n",
        "  for fn in uploaded:\n",
        "    shutil.move('/content/sample_data/'+fn, '/content/foreground')\n",
        "    break\n",
        "  foreground_url = None\n",
        "  fg_time_params = ''\n",
        "\n",
        "  %cd /content"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2KbIOyYIsxTG"
      },
      "source": [
        "#@title Optionally shorten foreground video\n",
        "start_seconds =  0#@param {type:\"number\"}\n",
        "duration_seconds =  60#@param {type:\"number\"}\n",
        "start_seconds = max(start_seconds,0)\n",
        "duration_seconds = max(duration_seconds,0)\n",
        "fg_time_params = ''\n",
        "if duration_seconds: \n",
        "  fg_time_params = '-ss %f -t %f'%(start_seconds, duration_seconds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GhV9c6zPtSWT",
        "outputId": "f5fa118e-3e32-4c5c-c424-dc09feeaf407",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "mirror_foreground = False #@param {type:\"boolean\"} \n",
        "copy_audio = True #@param {type:\"boolean\"}\n",
        "\n",
        "%cd /content\n",
        "fg_dir = '/content/U-2-Net/test_data/test_images'\n",
        "result_dir = '/content/out_frames'\n",
        "!rm -rf $fg_dir\n",
        "!mkdir -p $fg_dir\n",
        "!rm -rf $result_dir\n",
        "!mkdir -p $result_dir\n",
        "\n",
        "import imageio\n",
        "imageio.plugins.freeimage.download()\n",
        "import cv2\n",
        "import numpy as np\n",
        "import os\n",
        "import io\n",
        "import PIL\n",
        "\n",
        "def fix_dims(im):\n",
        "    if im.ndim == 2:\n",
        "        im = np.tile(im[..., None], [1, 1, 3])\n",
        "    return im[...,:3]\n",
        "\n",
        "def crop_resize(im, size, crop=False):\n",
        "  if im.shape[:2] == size:\n",
        "    return im\n",
        "  if size[0]<im.shape[0] or size[1]<im.shape[1]:\n",
        "    interp = cv2.INTER_AREA\n",
        "  else:\n",
        "    interp = cv2.INTER_CUBIC\n",
        "  if not crop:\n",
        "    return np.clip(cv2.resize(im, size[::-1], interpolation=interp),0,1)\n",
        "  ratio = max(size[0]/im.shape[0], size[1]/im.shape[1])\n",
        "  im = np.clip(cv2.resize(im, (int(np.ceil(im.shape[1]*ratio)), int(np.ceil(im.shape[0]*ratio))), interpolation=interp),0,1)\n",
        "  return im[(im.shape[0]-size[0])//2:(im.shape[0]-size[0])//2+size[0], (im.shape[1]-size[1])//2:(im.shape[1]-size[1])//2+size[1]]\n",
        "\n",
        "fg_image = fg_dir+'/frame_%05d.png'%1\n",
        "try:\n",
        "    fg_now = imageio.imread('/content/foreground')\n",
        "    fg_now = fix_dims(fg_now)\n",
        "    imageio.imwrite(fg_image, fg_now, format='PNG-FI')\n",
        "    fg_now = fg_now/255\n",
        "except Exception:\n",
        "    !ffmpeg $fg_time_params -i /content/foreground $fg_dir/frame_%05d.png    \n",
        "    fg_now = imageio.imread(fg_image, format='PNG-FI')\n",
        "fg_files = [x for x in sorted(os.listdir(fg_dir)) if x.endswith('.png')]\n",
        "!cp $fg_image /content/restyle-encoder/notebooks/images/toonify_img.jpg\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/content\n",
            "Imageio: 'libfreeimage-3.16.0-linux64.so' was not found on your computer; downloading it now.\n",
            "Try 1. Download from https://github.com/imageio/imageio-binaries/raw/master/freeimage/libfreeimage-3.16.0-linux64.so (4.6 MB)\n",
            "Downloading: 8192/4830080 bytes (0.2%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b2318336/4830080 bytes (48.0%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b4830080/4830080 bytes (100.0%)\n",
            "  Done\n",
            "File saved as /root/.imageio/freeimage/libfreeimage-3.16.0-linux64.so.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:root:imageio.freeimage warning: Could not determine tag type of 'ExifVersion'.\n",
            "WARNING:root:imageio.freeimage warning: Could not determine tag type of 'ComponentsConfiguration'.\n",
            "WARNING:root:imageio.freeimage warning: Could not determine tag type of 'FlashPixVersion'.\n",
            "WARNING:root:imagio.freeimage warning: Could not set tag 34953: No known reason., 'int' object has no attribute 'encode'\n",
            "WARNING:root:imageio.freeimage warning: Could not determine tag type of 59932.\n",
            "WARNING:root:imageio.freeimage warning: Could not determine tag type of 'SceneType'.\n",
            "WARNING:root:imageio.freeimage warning: Could not determine tag type of 'GPSInfo'.\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4weLFoPbkTKZ"
      },
      "source": [
        "## Step 5: Visualize Input"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r2H9zFLJkTKa",
        "outputId": "2860d9a0-c71c-454d-d1c6-849f7fc4801e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 319
        }
      },
      "source": [
        "image_path = EXPERIMENT_DATA_ARGS[experiment_type][\"image_path\"]\n",
        "original_image = Image.open(image_path).convert(\"RGB\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "FileNotFoundError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-22-7f6fb35abf16>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mimage_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mEXPERIMENT_DATA_ARGS\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mexperiment_type\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"image_path\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0moriginal_image\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mImage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconvert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"RGB\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/PIL/Image.py\u001b[0m in \u001b[0;36mopen\u001b[0;34m(fp, mode)\u001b[0m\n\u001b[1;32m   2841\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2842\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mfilename\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2843\u001b[0;31m         \u001b[0mfp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbuiltins\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilename\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"rb\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2844\u001b[0m         \u001b[0mexclusive_fp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2845\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'notebooks/images/toonify_img.jpg'"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-lbLKtl-kTKc"
      },
      "source": [
        "if experiment_type == 'cars_encode':\n",
        "    original_image = original_image.resize((192, 256))\n",
        "else:\n",
        "    original_image = original_image.resize((256, 256))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n0CWAjgxZ0VN"
      },
      "source": [
        "original_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o6oqf8JwzK0K"
      },
      "source": [
        "### Align Image\n",
        "\n",
        "Note: in this notebook we'll run alignment on the input image when working on the human facial domain."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hJ9Ce1aYzmFF"
      },
      "source": [
        "def run_alignment(image_path):\n",
        "    import dlib\n",
        "    from scripts.align_faces_parallel import align_face\n",
        "    if not os.path.exists(\"shape_predictor_68_face_landmarks.dat\"):\n",
        "        print('Downloading files for aligning face image...')\n",
        "        os.system('wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2')\n",
        "        os.system('bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2')\n",
        "        print('Done.')\n",
        "    predictor = dlib.shape_predictor(\"shape_predictor_68_face_landmarks.dat\")\n",
        "    aligned_image = align_face(filepath=image_path, predictor=predictor) \n",
        "    print(\"Aligned image has shape: {}\".format(aligned_image.size))\n",
        "    return aligned_image "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aTZcKMdK8y77"
      },
      "source": [
        "if experiment_type in ['ffhq_encode', 'toonify']:\n",
        "    input_image = run_alignment(image_path)\n",
        "else:\n",
        "    input_image = original_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hUBAfodh5PaM"
      },
      "source": [
        "input_image.resize((256, 256))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D0BmXzu1kTKg"
      },
      "source": [
        "## Step 6: Perform Inference"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T3h3E7VLkTKg"
      },
      "source": [
        "img_transforms = EXPERIMENT_ARGS['transform']\n",
        "transformed_image = img_transforms(input_image)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_fNBlRU8OSDL"
      },
      "source": [
        "Before running inference, we need to generate the image corresponding to the average latent code. These will be used to initialize the iterative refinement process."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fmpzoODNOSDL"
      },
      "source": [
        "def get_avg_image(net):\n",
        "    avg_image = net(net.latent_avg.unsqueeze(0),\n",
        "                    input_code=True,\n",
        "                    randomize_noise=False,\n",
        "                    return_latents=False,\n",
        "                    average_code=True)[0]\n",
        "    avg_image = avg_image.to('cuda').float().detach()\n",
        "    if experiment_type == \"cars_encode\":\n",
        "        avg_image = avg_image[:, 32:224, :]\n",
        "    return avg_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M5eWR2S4OSDM"
      },
      "source": [
        "Now we'll run inference. By default, we'll run using 5 inference steps. You can change the parameter in the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ct_jm0obOSDM"
      },
      "source": [
        "opts.n_iters_per_batch = 5\n",
        "opts.resize_outputs = False  # generate outputs at full resolution"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ls5zb0fRkTKs"
      },
      "source": [
        "from utils.inference_utils import run_on_batch\n",
        "\n",
        "with torch.no_grad():\n",
        "    avg_image = get_avg_image(net)\n",
        "    tic = time.time()\n",
        "    result_batch, result_latents = run_on_batch(transformed_image.unsqueeze(0).cuda(), net, opts, avg_image)\n",
        "    toc = time.time()\n",
        "    print('Inference took {:.4f} seconds.'.format(toc - tic))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nq0dkSz6kTKv"
      },
      "source": [
        "### Visualize Result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UVR03XT_kTK0"
      },
      "source": [
        "We'll visualize the step-by-step outputs side by side."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ca5BtxdUOSDN"
      },
      "source": [
        "if opts.dataset_type == \"cars_encode\":\n",
        "    resize_amount = (256, 192) if opts.resize_outputs else (512, 384)\n",
        "else:\n",
        "    resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WdR51hOROSDN"
      },
      "source": [
        "def get_coupled_results(result_batch, transformed_image):\n",
        "    \"\"\"\n",
        "    Visualize output images from left to right (the input image is on the right)\n",
        "    \"\"\"\n",
        "    result_tensors = result_batch[0]  # there's one image in our batch\n",
        "    result_images = [tensor2im(result_tensors[iter_idx]) for iter_idx in range(opts.n_iters_per_batch)]\n",
        "    input_im = tensor2im(transformed_image)\n",
        "    res = np.array(result_images[0].resize(resize_amount))\n",
        "    for idx, result in enumerate(result_images[1:]):\n",
        "        res = np.concatenate([res, np.array(result.resize(resize_amount))], axis=1)\n",
        "    res = np.concatenate([res, input_im.resize(resize_amount)], axis=1)\n",
        "    res = Image.fromarray(res)\n",
        "    return res"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uSDCvtTMOSDN"
      },
      "source": [
        "Note that the step-by-step outputs are shown left-to-right with the original input on the right-hand side."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lb3raAKFOSDN"
      },
      "source": [
        "res = get_coupled_results(result_batch, transformed_image)\n",
        "res"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qaB7RN7cOSDN"
      },
      "source": [
        "# save image \n",
        "res.save(f'./{experiment_type}_results.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ISEMFxmekTK7"
      },
      "source": [
        "# Encoder Bootstrapping\n",
        "\n",
        "In the paper, we introduce an encoder bootstrapping technique that can be used to solve the image toonification task by pairing an FFHQ-based encoder with a Toon-based encoder.  \n",
        "\n",
        "We demonstrate this idea below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sv284Ox8OSDO"
      },
      "source": [
        "# download the ffhq-based encoder if not previously downloaded\n",
        "path = MODEL_PATHS['ffhq_encode']\n",
        "EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS['ffhq_encode']\n",
        "ffhq_model_path = EXPERIMENT_ARGS['model_path']\n",
        "download_command = get_download_model_command(file_id=path[\"id\"], file_name=path[\"name\"]) \n",
        "if not os.path.exists(ffhq_model_path) or os.path.getsize(ffhq_model_path) < 1000000:\n",
        "    print('Downloading FFHQ ReStyle encoder...')\n",
        "    os.system(f\"wget {download_command}\")\n",
        "    # if google drive receives too many requests, we'll reach the quota limit and be unable to download the model\n",
        "    if os.path.getsize(ffhq_model_path) < 1000000:\n",
        "        !wget --no-check-certificate -nc https://eyalgruss.com/fomm/restyle_psp_ffhq_encode.pt -O /content/restyle-encoder/pretrained_models/restyle_psp_ffhq_encode.pt\n",
        "else:\n",
        "    print('FFHQ ReStyle encoder already exists!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FKbAFK7_OSDO"
      },
      "source": [
        "# download the toon-based encoder if not previously downloaded\n",
        "path = MODEL_PATHS['toonify']\n",
        "EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS['toonify']\n",
        "toonify_model_path = EXPERIMENT_ARGS['model_path']\n",
        "download_command = get_download_model_command(file_id=path[\"id\"], file_name=path[\"name\"]) \n",
        "# download the ffhq-based encoder if not previously downloaded\n",
        "path = MODEL_PATHS['ffhq_encode']\n",
        "EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS['ffhq_encode']\n",
        "ffhq_model_path = EXPERIMENT_ARGS['model_path']\n",
        "download_command = get_download_model_command(file_id=path[\"id\"], file_name=path[\"name\"]) \n",
        "if not os.path.exists(toonify_model_pathh) or os.path.getsize(toonify_model_path) < 1000000:\n",
        "    print('Downloading Toonify ReStyle encoder...')\n",
        "    os.system(f\"wget {download_command}\")\n",
        "    # if google drive receives too many requests, we'll reach the quota limit and be unable to download the model\n",
        "    if os.path.getsize(toonify_model_path) < 1000000:\n",
        "        !wget --no-check-certificate -nc https://eyalgruss.com/fomm/restyle_psp_toonify.pt -O /content/restyle-encoder/pretrained_models/restyle_psp_toonify.pt\n",
        "else:\n",
        "    print('Toonify ReStyle encoder already exists!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K3v0X3ZWkTK8"
      },
      "source": [
        "# load models \n",
        "ckpt = torch.load(ffhq_model_path, map_location='cpu')\n",
        "opts = ckpt['opts']\n",
        "opts['checkpoint_path'] = ffhq_model_path\n",
        "opts = Namespace(**opts)\n",
        "net1 = pSp(opts)\n",
        "net1.eval()\n",
        "net1.cuda()\n",
        "print('FFHQ Model successfully loaded!')\n",
        "\n",
        "ckpt = torch.load(toonify_model_path, map_location='cpu')\n",
        "opts = ckpt['opts']\n",
        "opts['checkpoint_path'] = toonify_model_path\n",
        "opts = Namespace(**opts)\n",
        "net2 = pSp(opts)\n",
        "net2.eval()\n",
        "net2.cuda()\n",
        "print('Toonify Model successfully loaded!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XW-CJsuwOSDO"
      },
      "source": [
        "# load image \n",
        "image_path = EXPERIMENT_DATA_ARGS['toonify'][\"image_path\"]\n",
        "original_image = Image.open(image_path).convert(\"RGB\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MmPWPODaOSDP"
      },
      "source": [
        "# transform image\n",
        "img_transforms = EXPERIMENT_ARGS['transform']\n",
        "transformed_image = img_transforms(original_image)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BiMjTyMzOSDP"
      },
      "source": [
        "opts.n_iters_per_batch = 5\n",
        "opts.resize_outputs = False  # generate outputs at full resolution"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o81i-MtOOSDQ"
      },
      "source": [
        "from scripts.encoder_bootstrapping_inference import run_on_batch\n",
        "\n",
        "with torch.no_grad():\n",
        "    avg_image = get_avg_image(net1)\n",
        "    tic = time.time()\n",
        "    result_batch = run_on_batch(transformed_image.unsqueeze(0).cuda(), net1, net2, opts, avg_image)\n",
        "    toc = time.time()\n",
        "    print('Inference took {:.4f} seconds.'.format(toc - tic))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1AGWnm9BOSDQ"
      },
      "source": [
        "Again we'll visualize the results from left to right. Here, the leftmost image is the inverted FFHQ image that is used to initialize the toonify ReStyle encoder. The following images show iterative results outputted by the toonify model.\n",
        "Finally, the rightmost image is the original input image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FX-_45rxOSDQ"
      },
      "source": [
        "res = get_coupled_results(result_batch, transformed_image)\n",
        "res"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pdk_QLRFOSDQ"
      },
      "source": [
        "# save image \n",
        "res.save(f'./encoder_bootstrapping_results.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ScKIEXTyt3cL"
      },
      "source": [
        "'''\n",
        "!pip install imageio-ffmpeg\n",
        "from IPython.display import HTML, clear_output, Image\n",
        "from base64 import b64encode\n",
        "import shutil\n",
        "!rm -f /content/final.mp4\n",
        "!rm -f /content/final.png\n",
        "if len(fg_files)>1 or len(bg_files)>1:\n",
        "  if is_fg:\n",
        "    with imageio.get_reader('/content/foreground', format='mp4') as reader:\n",
        "      fps = reader.get_meta_data()['fps']\n",
        "    if copy_audio:\n",
        "      !ffmpeg -framerate $fps -i $result_dir/frame_%05d.png $fg_time_params -i /content/foreground -c:v libx264 -c:a aac -map 0:v -map 1:a? -vf \"crop=trunc(iw/2)*2:trunc(ih/2)*2\" -pix_fmt yuv420p -profile:v baseline -movflags +faststart /content/final.mp4 -y\n",
        "    else:\n",
        "      !ffmpeg -framerate $fps -i $result_dir/frame_%05d.png -c:v libx264 -vf \"crop=trunc(iw/2)*2:trunc(ih/2)*2\" -pix_fmt yuv420p -profile:v baseline -movflags +faststart /content/final.mp4 -y\n",
        "  else:\n",
        "    with imageio.get_reader('/content/background_%05d'%0, format='mp4') as reader:\n",
        "      fps = reader.get_meta_data()['fps']\n",
        "    if copy_audio:\n",
        "      with open('/content/list.txt','w') as f:\n",
        "        if bg_time_params:\n",
        "          start_seconds = float(bg_time_params.split(' ')[1])\n",
        "          end_seconds = start_seconds + float(bg_time_params.split(' ')[3])\n",
        "        for file in sorted(glob.glob('/content/background_*')):\n",
        "          f.write(\"file '%s'\\n\"%file)\n",
        "          if bg_time_params:\n",
        "            f.write('inpoint %f\\n'%start_seconds)\n",
        "            f.write('outpoint %f\\n'%end_seconds)\n",
        "      !ffmpeg -f concat -safe 0 -i /content/list.txt -c copy /content/bg_audio.mp4 -y\n",
        "      !ffmpeg -framerate $fps -i $result_dir/frame_%05d.png -i /content/bg_audio.mp4 -c:v libx264 -c:a aac -map 0:v -map 1:a? -vf \"crop=trunc(iw/2)*2:trunc(ih/2)*2\" -pix_fmt yuv420p -profile:v baseline -movflags +faststart /content/final.mp4 -y\n",
        "    else:\n",
        "      !ffmpeg -framerate $fps -i $result_dir/frame_%05d.png -c:v libx264 -c:a aac -vf \"crop=trunc(iw/2)*2:trunc(ih/2)*2\" -pix_fmt yuv420p -profile:v baseline -movflags +faststart /content/final.mp4 -y\n",
        "  #video can be downloaded from /content/final.mp4\n",
        "  save_time = time()-start\n",
        "  total_time = time()-grand_start  \n",
        "  clear_output()\n",
        "  with open('/content/final.mp4', 'rb') as f:\n",
        "    data_url = \"data:video/mp4;base64,\" + b64encode(f.read()).decode()\n",
        "  display(HTML(\"\"\"\n",
        "  <video width=600 controls autoplay loop>\n",
        "        <source src=\"%s\" type=\"video/mp4\">\n",
        "  </video>\"\"\" % data_url))\n",
        "else:\n",
        "  shutil.move(out_dir+'/frame_%05d.png'%1, '/content/final.png')\n",
        "  #image can be downloaded from /content/final.png\n",
        "  save_time = time()-start\n",
        "  total_time = time()-grand_start\n",
        "  clear_output()\n",
        "  display(Image('/content/final.png', width=600))\n",
        "if model.startswith('u2net'):\n",
        "    print('frames=%d prepare=%d mask=%d blend=%d save=%d total=%d'%(len(iter_files), prepare_time, mask_time, blend_time, save_time, total_time))\n",
        "else:\n",
        "    print('frames=%d prepare=%d mask+blend=%d save=%d total=%d'%(len(iter_files), prepare_time, mask_time+blend_time, save_time, total_time))\n",
        "'''"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}