yarok.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "yarok.ipynb",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyOcZp8f845+1bQ0Oe7X4cs1",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "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/yarok.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_ETOKoCh2ur"
      },
      "source": [
        "# Video Green Screen and Cutter-Mixer-Upper ([j.mp/vid2green](https://j.mp/vid2green))\n",
        "\n",
        "## Qin et al., U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection, https://arxiv.org/abs/2005.09007, https://github.com/NathanUA/U-2-Net\n",
        "\n",
        "## Ke et al., MODNet: Is a Green Screen Really Necessary for Real-Time Portrait Matting?, https://arxiv.org/abs/2011.11961, https://github.com/ZHKKKe/MODNet\n",
        "\n",
        "### Made just a little bit more accessible by Eyal Gruss ([@eyaler](twitter.com/eyaler) / [eyalgruss.com](https://eyalgruss.com) / eyalgruss@gmail.com)\n",
        "\n",
        "#### Foreground options:\n",
        "*   Image from web or upload\n",
        "*   Video from web or upload\n",
        "*   Override frame dimensions and rate\n",
        "*   Concatenate multiple videos from folder or list\n",
        "*   Trim videos start times and end times or durations\n",
        "*   Resizing to match override or previous video height keeping aspect ratio\n",
        "*   Cut to short clips and mix up with options for random effects\n",
        "*   A second (\"double\") remixed layer of short clips\n",
        "*   Determine order of double layer by mask area\n",
        "*   Mirror\n",
        "*   Any [G'MIC](https://gmic.eu) filter combinations! (note: can be extermely slow)\n",
        "\n",
        "#### Model options:\n",
        "*   U^2-Net\n",
        "*   U^2-Net p (small)\n",
        "*   U^2-Net human segmentation\n",
        "*   MODNet photographic model\n",
        "*   MODNet \"webcam\" model\n",
        "*   U^2-Net portrait generation (sketch)\n",
        "*   U^2-Net portrait generation + U^2-Net blending\n",
        "*   U^2-Net preprocessing + U^2-Net portrait generation\n",
        "*   U^2-Net preprocessing + U^2-Net portrait generation + U^2-Net blending\n",
        "*   similar combinations as above with any of the other U^2-Net models\n",
        "\n",
        "#### Blending options:\n",
        "*   Continuous blending\n",
        "*   One frame delay smoothing with custom threshold\n",
        "*   Grab sketch colors from original foreground\n",
        "*   Minimal and Maximal mask area thresholds\n",
        "*   Foreground blend refinement with pymatting's [Fast Multi-Level Foreground Estimation](https://pymatting.github.io/pymatting.foreground.html#module-pymatting.foreground.estimate_foreground_ml) (https://arxiv.org/abs/2006.14970)\n",
        "\n",
        "#### Background options:\n",
        "*   Solid white, black, chroma green, chroma blue or any hex value\n",
        "*   Transparent (for images)\n",
        "*   Image from web or upload\n",
        "*   Video from web or upload\n",
        "*   Concatenate multiple videos from folder or list\n",
        "*   Trim videos start times and end times or durations\n",
        "*   Cut to short clips and mix up with options for random effects\n",
        "*   Loop video from start or in reverse\n",
        "*   Resize to match foreground, optionally keeping aspect ratio\n",
        "*   Match frame rate to foreground or override\n",
        "*   Use foreground image/video with the optional following effects\n",
        "*   Mirror\n",
        "*   Any [G'MIC](https://gmic.eu) filter combination! (note: can be extermely slow)\n",
        "*   Grayscale\n",
        "*   Bokeh (gamma-corrected blurred)\n",
        "\n",
        "#### Examples:\n",
        "*   U^2-Net/MODNet model comparison: https://twitter.com/eyaler/status/1342853322127110146\n",
        "*   U^2-Net portrait sketching \"Take On Me\": https://twitter.com/eyaler/status/1346395761236443138, https://vimeo.com/showcase/8491391\n",
        "*   Sexy vaporwave lady series: https://youtu.be/74vgrfIVE9A&list=PLZV8sKAO3LzZwHGWbmRxLuQQq2lXZsBaN\n",
        "*   Adam Cutmix videoart: https://vimeo.com/561436016\n",
        "\n",
        "##### A curated list of online generative tools: [j.mp/generativetools](https://j.mp/generativetools)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6K5a29Wbgi0L",
        "cellView": "form"
      },
      "source": [
        "#@title Setup\n",
        "%cd /content\n",
        "!git clone --depth 1 https://github.com/eyaler/U-2-Net\n",
        "import os\n",
        "!pip install -U gdown\n",
        "\n",
        "!mkdir -p /content/U-2-Net/saved_models/u2net\n",
        "%cd /content/U-2-Net/saved_models/u2net\n",
        "if not os.path.exists('/content/U-2-Net/saved_models/u2net/u2net.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ -O /content/U-2-Net/saved_models/u2net/u2net.pth\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/u2net.pth\n",
        "\n",
        "!mkdir -p /content/U-2-Net/saved_models/u2netp\n",
        "%cd /content/U-2-Net/saved_models/u2netp\n",
        "if not os.path.exists('/content/U-2-Net/saved_models/u2netp/u2netp.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy -O /content/U-2-Net/saved_models/u2netp/u2netp.pth\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/u2netp.pth\n",
        "\n",
        "!mkdir -p /content/U-2-Net/saved_models/u2net_human_seg\n",
        "%cd /content/U-2-Net/saved_models/u2net_human_seg\n",
        "if not os.path.exists('/content/U-2-Net/saved_models/u2net_human_seg/u2net_human_seg.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1-Yg0cxgrNhHP-016FPdp902BR-kSsA4P -O /content/U-2-Net/saved_models/u2net_human_seg/u2net_human_seg.pth\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/u2net_human_seg.pth\n",
        "\n",
        "!mkdir -p /content/U-2-Net/saved_models/u2net_portrait\n",
        "%cd /content/U-2-Net/saved_models/u2net_portrait\n",
        "if not os.path.exists('/content/U-2-Net/saved_models/u2net_portrait/u2net_portrait.pth'):\n",
        "  !gdown https://drive.google.com/file/d/1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy -O /content/U-2-Net/saved_models/u2net_portrait/u2net_portrait.pth\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/u2net_portrait.pth\n",
        "\n",
        "%cd /content\n",
        "!git clone --depth 1 https://github.com/eyaler/MODNet\n",
        "%cd MODNet/pretrained\n",
        "\n",
        "if not os.path.exists('/content/MODNet/pretrained/modnet_photographic_portrait_matting.ckpt'):\n",
        "  !gdown https://drive.google.com/uc?id=1mcr7ALciuAsHCpLnrtG_eop5-EYhbCmz -O /content/MODNet/pretrained/modnet_photographic_portrait_matting.ckpt\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/modnet_photographic_portrait_matting.ckpt\n",
        "\n",
        "if not os.path.exists('/content/MODNet/pretrained/modnet_webcam_portrait_matting.ckpt'):\n",
        "  !gdown https://drive.google.com/uc?id=1Nf1ZxeJZJL8Qx9KadcYYyEmmlKhTADxX -O /content/MODNet/pretrained/modnet_webcam_portrait_matting.ckpt\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/modnet_webcam_portrait_matting.ckpt\n",
        "\n",
        "%cd /content\n",
        "#!pip install youtube-dl\n",
        "!pip install git+https://github.com/ytdl-org/youtube-dl\n",
        "!pip install imageio==2.9.0\n",
        "!pip install imageio-ffmpeg==0.4.5\n",
        "!pip install pymatting\n",
        "!pip install gmic==2.9.4a1\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3dExwJXtQN2j",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally mount Google Drive { run: \"auto\" }\n",
        "#@markdown This will also save output to your Google drive in vid2green_output\n",
        "mount_google_drive = False #@param {type:\"boolean\"}\n",
        "if mount_google_drive:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IT3ERnOxl-jz",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally enter lists of videos\n",
        "#@markdown 1. **Run** cell to allow input. \n",
        "#@markdown 2. Do **not** rerun after input. \n",
        "#@markdown 3. You need to **select** \"List 1\" or \"List 2\" in the **next** cell\n",
        "\n",
        "from ipywidgets import Textarea, HBox\n",
        "try:\n",
        "  value=list1.value\n",
        "except:\n",
        "  value=None\n",
        "list1 = Textarea(\n",
        "    placeholder='''https://www.youtube.com/watch?v=R8WzG0Ts6L0, 1:07, 2:02\n",
        "https://www.youtube.com/watch?v=Kls_hbMml_Y, 30, 50\n",
        "https://www.youtube.com/watch?v=Gqf_cuDf9qI, 2:35, 3:30\n",
        "https://www.youtube.com/watch?v=SMDkvJRHaNM, 1:06, 1:43\n",
        "https://www.youtube.com/watch?v=5e0MWOOqxxQ, , 22\n",
        "''',\n",
        "    value=value,\n",
        "    layout={'width': '45%', 'height':'95%'},\n",
        "    description='List 1')\n",
        "\n",
        "try:\n",
        "  value=list2.value\n",
        "except:\n",
        "  value=None\n",
        "list2 = Textarea(\n",
        "    placeholder='''https://www.youtube.com/watch?v=rQLA8ZIEM3A, , 7\n",
        "https://www.youtube.com/watch?v=5BHoDW9f7vY, 1:39, 2:29\n",
        "https://www.youtube.com/watch?v=kyd1j3mRsw4, 3, 2:16\n",
        "https://www.youtube.com/watch?v=yNHRXNvFmZ8, , 3:10\n",
        "https://www.youtube.com/watch?v=Q93VrYOXSe8, 1:05, 2:10\n",
        "''',\n",
        "    value=value,\n",
        "    layout={'width': '45%', 'height':'95%'},\n",
        "    description='List 2')\n",
        "display(HBox([list1,list2], layout={'height':'250px'}))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vGtTGhKUh1fW",
        "cellView": "form"
      },
      "source": [
        "#@title Get the foreground image/video and background image/video\n",
        "#@markdown 1. You can change the URLs to your **own** stuff! (e.g. youtube, mp4)\n",
        "#@markdown 2. You can add comma separated **start and end times** or use &t=XXXs or ?t=XXXs notations\n",
        "#@markdown 3. For the background you can use the drop-down menu to alternatively choose a **solid color**, **transparent** (for images only) or the **foreground** iteslf (with optional effects). You can also feed a **hex** color value (e.g. de0000 or #de0000)\n",
        "#@markdown 4. Advanced: You can use the **above lists** (select \"List 1\" or \"List 2\") or point to a txt/csv file with **a list** urls/paths and optional **start and end times**.\n",
        "#@markdown 5. Advanced: You can add files mannualy (e.g. from a mounted Google drive) and put in the path the location on Colab. This path can also be to a folder containing **several videos**.\n",
        "#@markdown 6. Alternatively, you select Upload to **manually upload** image/video/videos from your machine. After pressing play, a button will appear bellow to select files.\n",
        "#@markdown 7. You can **reuse** a previously fetched resources by selecting \"Reuse\" or deleting the entry.\n",
        "\n",
        "#foreground_url = 'https://www.youtube.com/watch?v=HzpzvAPj1kw' #@param ['List 1', 'List 2', 'Upload', 'Reuse'] {allow-input: true}\n",
        "#background_url = 'https://www.youtube.com/watch?v=pXpvh6eIFBk' #@param ['White', 'Black', 'Chroma Green', 'Chroma Blue', 'Transparent', 'Foreground', 'List 1', 'List 2', 'Upload', 'Reuse'] {allow-input: true}\n",
        "\n",
        "#foreground_url = 'https://www.youtube.com/watch?v=kMpnwIGDQvU' #@param ['List 1', 'List 2', 'Upload', 'Reuse'] {allow-input: true}\n",
        "#background_url = 'https://www.youtube.com/watch?v=dMvnCyznteU' #@param ['White', 'Black', 'Chroma Green', 'Chroma Blue', 'Transparent', 'Foreground', 'List 1', 'List 2', 'Upload', 'Reuse'] {allow-input: true}\n",
        "\n",
        "foreground_url = 'https://www.youtube.com/watch?v=kMpnwIGDQvU' #@param ['List 1', 'List 2', 'Upload', 'Reuse'] {allow-input: true}\n",
        "background_url = 'https://www.youtube.com/watch?v=dMvnCyznteU' #@param ['White', 'Black', 'Chroma Green', 'Chroma Blue', 'Transparent', 'Foreground', 'List 1', 'List 2', 'Upload', 'Reuse'] {allow-input: true}\n",
        "\n",
        "import os\n",
        "import csv\n",
        "import re\n",
        "import youtube_dl\n",
        "from google.colab import files\n",
        "\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",
        "    if '://' not in url:\n",
        "      url = 'https://'+url\n",
        "    for e in extractors:\n",
        "      if e.suitable(url) and e.IE_NAME != 'generic':\n",
        "        return True\n",
        "    return False\n",
        "\n",
        "def get_media(url, name):\n",
        "  !rm -rf /content/$name\n",
        "  !mkdir -p /content/$name\n",
        "  %cd /content\n",
        "  if os.path.isdir(url):\n",
        "    return sorted(os.path.join(os.path.abspath(url),file) for file in os.listdir(url))\n",
        "  if url=='Upload':   \n",
        "    %cd /content/sample_data\n",
        "    print('Upload %s:'%name)\n",
        "    try:\n",
        "      uploaded = files.upload()\n",
        "    except Exception:\n",
        "      %cd /content\n",
        "      raise\n",
        "    urls = sorted(os.path.abspath(file) for file in uploaded)\n",
        "    %cd /content\n",
        "    return urls\n",
        "  lines = [url.split(',')]\n",
        "  urls = []\n",
        "  if url=='List 1':\n",
        "    lines = csv.reader(list1.value.splitlines())\n",
        "    assert lines, url+' is empty'\n",
        "  elif url=='List 2':\n",
        "    lines = csv.reader(list2.value.splitlines())\n",
        "    assert lines, url+' is empty'\n",
        "  elif url.endswith(('.txt','.csv')):\n",
        "    if not os.path.exists(url):\n",
        "      !wget \"$url\" -O /content/list.csv\n",
        "      url = '/content/list.csv'\n",
        "    with open(url,encoding='utf8') as f: \n",
        "      lines = list(csv.reader(f))\n",
        "  for line in lines:\n",
        "    try:\n",
        "      if not line or not line[0].strip():\n",
        "        continue\n",
        "      url = line[0]\n",
        "      start_time = '0'\n",
        "      end_time = '0'\n",
        "      if len(line)>1 and line[1].strip():\n",
        "        start_time = line[1]\n",
        "      if not re.search('[^0:. ]', start_time):\n",
        "        stamp = re.search(r'[&?#]t=(\\d+.?\\d+)',url)\n",
        "        if stamp:\n",
        "          start_time = stamp.group(1)\n",
        "      url = re.split('[&?#]t=', url, maxsplit=1)[0]\n",
        "      if len(line)>2 and line[2].strip():\n",
        "        end_time = line[2]\n",
        "      outfile = os.path.join('/content',name,'%06d'%len(urls))\n",
        "      if os.path.isfile(url):\n",
        "        infile = url\n",
        "        if not re.search('[^0:. ]', start_time+end_time):\n",
        "          outfile = os.path.abspath(url)\n",
        "      else:\n",
        "        infile = '/content/temp'\n",
        "        if not re.search('[^0:. ]', start_time+end_time):\n",
        "          infile = outfile\n",
        "        !rm -f $infile\n",
        "        if is_supported(url):\n",
        "          !rm -f /content/temp.mp4\n",
        "          !youtube-dl --no-playlist -f \"bestvideo[ext=mp4][vcodec!*=av01]+bestaudio[ext=m4a]/mp4\" \"$url\" --merge-output-format mp4 -o /content/temp\n",
        "          !mv /content/temp.mp4 $infile\n",
        "        if not os.path.exists(infile):\n",
        "          !wget \"$url\" -O $infile\n",
        "      urls.append(outfile)\n",
        "      if re.search('[^0:. ]', start_time+end_time):\n",
        "        end_time_arg = '-copyts -to %s'%end_time if re.search('[^0:. ]', end_time) else ''\n",
        "        !ffmpeg -ss $start_time -i \"$infile\" $end_time_arg -f mp4 \"$outfile\" -y\n",
        "      assert os.path.exists(outfile), 'Cannot find file: %s'%outfile\n",
        "    except Exception:\n",
        "      print('Error for',line)\n",
        "      raise\n",
        "  return urls\n",
        "\n",
        "got_bg = False\n",
        "if foreground_url != 'Upload' and background_url == 'Upload':\n",
        "  got_bg = True\n",
        "  background_src = get_media(background_url, 'background')\n",
        "\n",
        "if foreground_url not in ('', 'Reuse'):\n",
        "  foreground_src = get_media(foreground_url, 'foreground')\n",
        "\n",
        "if not got_bg:\n",
        "  if '/' in background_url or '.' in background_url or '\\\\' in background_url or background_url in ['List 1','List 2','Upload']:\n",
        "    background_src = get_media(background_url, 'background')\n",
        "  elif background_url not in ('', 'Reuse'):\n",
        "    background_src = background_url.lower()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X8otpdxD88FH",
        "cellView": "form"
      },
      "source": [
        "#@title Green that screen!\n",
        "#@markdown Model notes:\n",
        "#@markdown 1. u2net tends to remove more unwanted parts, but may also remove desired parts of the foreground objects.\n",
        "#@markdown 2. modnet tends to keep more of the desired parts and also gives a finer boundary, but may leave in more unwanted parts (which is the more useful option if you further post edit the video).\n",
        "#@markdown 3. u2net_portrait will generate a sketch + if background is image/video/\"Foreground\" then use u2net mask of the original foreground to blend with that background.\n",
        "#@markdown 4. u2net + u2net_portrait adds a preprocessing stage to remove the background before generating a sketch of the foreground.\n",
        "#@markdown\n",
        "#@markdown Sketching notes:\n",
        "#@markdown 1. background=\"White\" + model=\"u2net_portrait\" -> everything becomes a sketch\n",
        "#@markdown 2. background=\"Foreground\" + model=\"u2net_portrait\" -> sketched foreground on top of original background\n",
        "#@markdown 3. background=\"White\" + model=\"u2net + u2net_portrait\" -> sketched foreground with white background\n",
        "#@markdown 4. background=\"Foreground\" + model=\"u2net + u2net_portrait\" -> variation similar to (2)\n",
        "#@markdown 5. background=upload result of (1) + model=\"u2net\" -> original foreground on top of sketched background  \n",
        "\n",
        "model = 'u2net' #@param ['u2net', 'modnet_photographic', 'modnet_webcam', 'u2net_portrait', 'u2net + u2net_portrait']\n",
        "u2net_variant = 'u2net' #@param ['u2net', 'u2netp', 'u2net_human_seg']\n",
        "sketch_color = 'Gray' #@param ['Gray','Foreground','Tint outline','Tint fill']\n",
        "one_frame_delay = True #@param {type:\"boolean\"} \n",
        "one_frame_delay_threshold = 0.1 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "min_mask_area_threshold = 0.0 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "max_mask_area_threshold = 1.0 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "refine_foreground_blend = False #@param {type:\"boolean\"}  \n",
        "start_secs_foreground = 0#@param {type:\"number\"}\n",
        "duration_secs_foreground = 0#@param {type:\"number\"}\n",
        "start_secs_background = 0#@param {type:\"number\"}\n",
        "duration_secs_background = 0#@param {type:\"number\"}\n",
        "mirror_foreground = False #@param {type:\"boolean\"} \n",
        "mirror_background = False #@param {type:\"boolean\"} \n",
        "gmic_filters_foreground = ' ' #@param [' ', 'choose_randomly_for_cut_mix', 'fx_boxfitting 3,0,0.1,0', 'cartoon 3,200,20,0.25,1.5,8,0', 'fx_diffusiontensors 10,5,3,1,0.15,1,0,3,0', 'fx_engrave_preview 0.5,50,0,8,40,0,0,0,10,1,0,0,0,1,0', 'fx_engrave_preview 0.5,50,0,8,40,0,0,1,10,1,0,0,0,1,0', 'fx_pixelsort 1,0,0,1,0,100,0,0,0', 'fx_polygonize 300,10,10,10,10,0,0,0,255,0', 'rodilius 10,10,300,5,30,1,1,0,0,0', 'fx_stained_glass 20,0.1,1,1,1,0,0,0,0', 'fx_8bits 25,800,16,0'] {allow-input: true}\n",
        "gmic_filters_background = ' ' #@param [' ', 'choose_randomly_for_cut_mix', 'fx_boxfitting 3,0,0.1,0', 'cartoon 3,200,20,0.25,1.5,8,0', 'fx_diffusiontensors 10,5,3,1,0.15,1,0,3,0', 'fx_engrave_preview 0.5,50,0,8,40,0,0,0,10,1,0,0,0,1,0', 'fx_engrave_preview 0.5,50,0,8,40,0,0,1,10,1,0,0,0,1,0', 'fx_pixelsort 1,0,0,1,0,100,0,0,0', 'fx_polygonize 300,10,10,10,10,0,0,0,255,0', 'rodilius 10,10,300,5,30,1,1,0,0,0', 'fx_stained_glass 20,0.1,1,1,1,0,0,0,0', 'fx_8bits 25,800,16,0'] {allow-input: true}\n",
        "gray_background = False #@param {type:\"boolean\"} \n",
        "bokeh_background = False #@param {type:\"boolean\"} \n",
        "bokeh_prcnt = 5 #@param {type:\"slider\", min:1, max:50, step:1}\n",
        "bokeh_gamma = 5 #@param {type:\"slider\", min:1, max:50, step:1}\n",
        "keep_aspect_background = True #@param {type:\"boolean\"}\n",
        "loop_reverse_background = False #@param {type:\"boolean\"}\n",
        "override_width = 0#@param {type:\"integer\"}\n",
        "override_height = 0#@param {type:\"integer\"}\n",
        "override_fps = 0#@param {type:\"number\"}\n",
        "#cut_mix_total_minutes = 0#@param {type:\"number\"}\n",
        "cut_mix_min_secs_foreground = 0#@param {type:\"number\"}\n",
        "cut_mix_max_secs_foreground = 0#@param {type:\"number\"}\n",
        "cut_mix_tied_background = False #@param {type:\"boolean\"}\n",
        "cut_mix_min_secs_background = 0#@param {type:\"number\"}\n",
        "cut_mix_max_secs_background = 0#@param {type:\"number\"}\n",
        "cut_mix_phase_secs_background = 0#@param {type:\"number\"}\n",
        "cut_mix_double_foreground = False #@param {type:\"boolean\"}\n",
        "cut_mix_double_sketch_color = 'Gray' #@param ['Gray','Foreground','Tint outline','Tint fill']\n",
        "cut_mix_mirror_double = False #@param {type:\"boolean\"}\n",
        "cut_mix_gmic_filters_double = ' ' #@param [' ', 'choose_randomly_for_cut_mix', 'fx_boxfitting 3,0,0.1,0', 'cartoon 3,200,20,0.25,1.5,8,0', 'fx_diffusiontensors 10,5,3,1,0.15,1,0,3,0', 'fx_engrave_preview 0.5,50,0,8,40,0,0,0,10,1,0,0,0,1,0', 'fx_engrave_preview 0.5,50,0,8,40,0,0,1,10,1,0,0,0,1,0', 'fx_pixelsort 1,0,0,1,0,100,0,0,0', 'fx_polygonize 300,10,10,10,10,0,0,0,255,0', 'rodilius 10,10,300,5,30,1,1,0,0,0', 'fx_stained_glass 20,0.1,1,1,1,0,0,0,0', 'fx_8bits 25,800,16,0'] {allow-input: true}\n",
        "cut_mix_min_secs_double = 0#@param {type:\"number\"}\n",
        "cut_mix_max_secs_double = 0#@param {type:\"number\"}\n",
        "cut_mix_phase_secs_double = 0#@param {type:\"number\"}\n",
        "cut_mix_layer_order_double = 'Front if smaller mask' #@param ['Middle', 'Front', 'Front if smaller mask', 'Front if larger mask']\n",
        "cut_mix_horizontal_flip_prob = 0#@param {type:\"slider\", min:0, max:0.5, step:0.01}\n",
        "cut_mix_vertical_flip_prob = 0#@param {type:\"slider\", min:0, max:0.5, step:0.01}\n",
        "cut_mix_temporal_flip_prob = 0#@param {type:\"slider\", min:0, max:0.5, step:0.01}\n",
        "cut_mix_random_gmic_filters_prob = 0#@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "random_seed = 42#@param {type:\"integer\"}\n",
        "copy_audio = True #@param {type:\"boolean\"}\n",
        "override_defaults_for_adam_cutmix = False #@param {type:\"boolean\"}\n",
        "\n",
        "if override_defaults_for_adam_cutmix:\n",
        "  if override_width==0:\n",
        "    override_width = 854\n",
        "  if override_height==0:\n",
        "    override_height = 480\n",
        "  if override_fps==0:\n",
        "    override_fps = 25\n",
        "  if cut_mix_min_secs_foreground==0:\n",
        "    cut_mix_min_secs_foreground = 4\n",
        "  if cut_mix_max_secs_foreground==0:  \n",
        "    cut_mix_max_secs_foreground = 4\n",
        "  if cut_mix_min_secs_background==0:\n",
        "    cut_mix_min_secs_background = 4\n",
        "  if cut_mix_max_secs_background==0:\n",
        "    cut_mix_max_secs_background = 4\n",
        "  if cut_mix_phase_secs_background==0:\n",
        "    cut_mix_phase_secs_background = -1.33\n",
        "  cut_mix_double_foreground = True\n",
        "  if cut_mix_min_secs_double==0:\n",
        "    cut_mix_min_secs_double = 4\n",
        "  if cut_mix_max_secs_double==0:\n",
        "    cut_mix_max_secs_double = 4\n",
        "  if cut_mix_phase_secs_double==0:\n",
        "    cut_mix_phase_secs_double = 1.33\n",
        "  if cut_mix_horizontal_flip_prob==0:\n",
        "    cut_mix_horizontal_flip_prob = 0.4\n",
        "  if cut_mix_vertical_flip_prob==0:\n",
        "    cut_mix_vertical_flip_prob = 0.04\n",
        "  if cut_mix_temporal_flip_prob==0:\n",
        "    cut_mix_temporal_flip_prob = 0.2\n",
        "\n",
        "calc_area = min_mask_area_threshold or max_mask_area_threshold<1  \n",
        "frame_png_compression = 1\n",
        "frame_jpg_quality = 75\n",
        "frame_jpg_qscale = 2\n",
        "frame_file_format = 'jpg'\n",
        "read_imglib = 'imageio-fi'\n",
        "write_imglib = 'imageio-fi'\n",
        "quality_arg = '-compression_level %d'%frame_png_compression if frame_file_format=='png' else '-q:v %d'%frame_jpg_qscale\n",
        "bg_mode_max_w = 1920\n",
        "chroma_thresholds = [0.5,0.25]\n",
        "sketch_color = sketch_color.lower()\n",
        "u2net_size = 320\n",
        "u2net_portrait_size = 512\n",
        "modnet_size = 512\n",
        "\n",
        "start_secs_foreground = max(start_secs_foreground,0)\n",
        "duration_secs_foreground = max(duration_secs_foreground,0)\n",
        "fg_time_params = ''\n",
        "if duration_secs_foreground: \n",
        "  fg_time_params = '-ss %f -t %f'%(start_secs_foreground, duration_secs_foreground)\n",
        "\n",
        "start_secs_background = max(start_secs_background,0)\n",
        "duration_secs_background = max(duration_secs_background,0)\n",
        "bg_time_params = ''\n",
        "if duration_secs_background: \n",
        "  bg_time_params = '-ss %f -t %f'%(start_secs_background, duration_secs_background)\n",
        "\n",
        "override_width = max(override_width, 0)\n",
        "override_height = max(override_height, 0)\n",
        "override_fps = max(override_fps, 0)\n",
        "\n",
        "fast_filters = ['cartoon 3,200,20,0.25,1.5,8,0', 'fx_diffusiontensors 10,5,3,1,0.15,1,0,3,0', 'fx_pixelsort 1,0,0,1,0,100,0,0,0', 'fx_polygonize 300,10,10,10,10,0,0,0,255,0', 'fx_8bits 25,800,16,0']\n",
        "cut_mix_random_gmic_filters_foreground = [gmic_filters_foreground]\n",
        "if gmic_filters_foreground == 'choose_randomly_for_cut_mix':\n",
        "  gmic_filters_foreground = ''\n",
        "  cut_mix_random_gmic_filters_foreground = fast_filters\n",
        "cut_mix_random_gmic_filters_background = [gmic_filters_background]\n",
        "if gmic_filters_background == 'choose_randomly_for_cut_mix':\n",
        "  gmic_filters_background = ''\n",
        "  cut_mix_random_gmic_filters_background = fast_filters\n",
        "cut_mix_random_gmic_filters_double = [cut_mix_gmic_filters_double]\n",
        "if cut_mix_gmic_filters_double == 'choose_randomly_for_cut_mix':\n",
        "  cut_mix_gmic_filters_double = ''\n",
        "  cut_mix_random_gmic_filters_double = fast_filters\n",
        "\n",
        "#cut_mix_total_minutes = max(cut_mix_total_minutes, 0)\n",
        "cut_mix_min_secs_foreground = max(cut_mix_min_secs_foreground, 0)\n",
        "cut_mix_max_secs_foreground = max(cut_mix_min_secs_foreground, cut_mix_max_secs_foreground)\n",
        "cut_mix_min_secs_background = max(cut_mix_min_secs_background, 0)\n",
        "cut_mix_max_secs_background = max(cut_mix_min_secs_background, cut_mix_max_secs_background)\n",
        "cut_mix_min_secs_double = max(cut_mix_min_secs_double, 0)\n",
        "cut_mix_max_secs_double = max(cut_mix_min_secs_double, cut_mix_max_secs_double)\n",
        "cut_mix_layer_order_double = cut_mix_layer_order_double.lower()\n",
        "if background_src=='foreground' and cut_mix_tied_background:\n",
        "  cut_mix_min_secs_background = 0\n",
        "  cut_mix_max_secs_background = 0\n",
        "  cut_mix_phase_secs_background = 0\n",
        "\n",
        "%cd /content/U-2-Net\n",
        "from data_loader import RescaleT, ToTensorLab, SalObjDataset\n",
        "from model import U2NET, U2NETP\n",
        "%cd /content\n",
        "\n",
        "import imageio\n",
        "imageio.plugins.freeimage.download()\n",
        "import cv2\n",
        "import numpy as np\n",
        "import os\n",
        "import glob\n",
        "from time import time\n",
        "from datetime import datetime\n",
        "import io\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as transforms\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import DataLoader\n",
        "from MODNet.src.models.modnet import MODNet\n",
        "from pymatting import estimate_foreground_ml\n",
        "import gmic\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "if random_seed:\n",
        "  np.random.seed(random_seed)\n",
        "\n",
        "fg_dir = '/content/fg_frames'\n",
        "mask_dir = '/content/mask_frames'\n",
        "bg_dir = '/content/bg_frames'\n",
        "result_dir = '/content/out_frames'\n",
        "portrait_in_dir = '/content/portrait_in'\n",
        "portrait_out_dir = '/content/portrait_out'\n",
        "!rm -rf $fg_dir\n",
        "!mkdir -p $fg_dir\n",
        "!rm -rf $mask_dir\n",
        "!mkdir -p $mask_dir\n",
        "!rm -rf $bg_dir\n",
        "!mkdir -p $bg_dir\n",
        "!rm -rf $result_dir\n",
        "!mkdir -p $result_dir\n",
        "!rm -rf $portrait_in_dir\n",
        "!mkdir -p $portrait_in_dir\n",
        "!rm -rf $portrait_out_dir\n",
        "!mkdir -p $portrait_out_dir\n",
        "\n",
        "torch_transforms = transforms.Compose(\n",
        "  [\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
        "  ]\n",
        ")\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 imread(path, file=None, file_format=None, imglib=read_imglib):\n",
        "  if file is not None:\n",
        "    if type(file)==int:\n",
        "      file = 'frame_%06d.jpg'%file\n",
        "    path = os.path.join(path,file)\n",
        "  try:\n",
        "    im = imageio.imread(path, format='JPEG-FI')\n",
        "  except Exception:\n",
        "    im = imageio.imread(path)\n",
        "  im = fix_dims(im)/255\n",
        "  return im\n",
        "\n",
        "def imwrite(im, path, file=None, file_format=None, png_compression=9, jpg_quality=95, imglib=write_imglib):\n",
        "  if file is not None:\n",
        "    if type(file)==int:\n",
        "      file = 'frame_%06d.jpg'%file\n",
        "    path = os.path.join(path,file)\n",
        "  imageio.imwrite(path, np.uint8(im*255), quality=jpg_quality, format='JPEG-FI')\n",
        "\n",
        "def safe_resize(im, w_h_size, interp):\n",
        "  w, h = w_h_size\n",
        "  assert max(w,h)>0\n",
        "  aspect = im.shape[1]/im.shape[0]\n",
        "  if w<=0:\n",
        "    w = int(np.ceil(h*aspect))\n",
        "  elif h<=0:\n",
        "    h = int(np.ceil(w/aspect))\n",
        "  return cv2.resize(im, (w,h), interpolation=interp)\n",
        "\n",
        "def resize_crop_pad(im, size, crop_by_max_ratio=False, crop_pad_to_height=False):\n",
        "  assert not crop_by_max_ratio or not crop_pad_to_height\n",
        "  size = tuple(size)\n",
        "  if not size[0] and not size[1] or im.shape[:2] == size:\n",
        "    return im\n",
        "  if not size[0] or size[0]==-1:\n",
        "    size = (int(round(im.shape[0]/im.shape[1]*size[1])), size[1])\n",
        "  elif not size[1] or size[1]==-1:\n",
        "    size = (size[0], int(round(im.shape[1]/im.shape[0]*size[0])))\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 crop_by_max_ratio:\n",
        "    max_ratio = max(size[0]/im.shape[0], size[1]/im.shape[1])\n",
        "    im = safe_resize(im, (int(np.ceil(im.shape[1]*max_ratio)), int(np.ceil(im.shape[0]*max_ratio))), interp=interp)\n",
        "    diff0 = im.shape[0]-size[0]\n",
        "    diff1 = im.shape[1]-size[1]\n",
        "    if diff0>0 or diff1>0:\n",
        "      im = im[diff0//2:size[0]+diff0//2, diff1//2:size[1]+diff1//2]\n",
        "  elif crop_pad_to_height:\n",
        "    im = safe_resize(im, (-1, size[0]), interp=interp)\n",
        "    diff1 = im.shape[1]-size[1]\n",
        "    if diff1>0:\n",
        "      im = im[:, diff1//2:size[1]+diff1//2]\n",
        "    elif diff1<0:\n",
        "      im = np.pad(im, pad_width=[(0,0),(-(diff1//2),diff1//2-diff1),(0,0)])\n",
        "  else:\n",
        "    im = safe_resize(im, size[::-1], interp=interp)\n",
        "  assert im.shape[:2] == size, (im.shape[:2],size)\n",
        "  return np.clip(im,0,1)\n",
        "\n",
        "g = gmic.Gmic()\n",
        "def do_gmic(im, filters):\n",
        "  if not filters.strip():\n",
        "    return im\n",
        "  \n",
        "  # see https://github.com/myselfhimself/gmic-py/issues/61\n",
        "  #im = [gmic.GmicImage.from_numpy_helper(im, deinterleave=True)]\n",
        "  #g.run(filters, im)\n",
        "  #return fix_dims(im[0].to_numpy_helper(interleave=True, permute='yxzc', squeeze_shape=True))\n",
        "\n",
        "  imwrite(im, 'gmic.jpg')\n",
        "  g.run('input gmic.jpg '+filters+' output gmic.jpg')    \n",
        "  return imread('gmic.jpg') \n",
        "\n",
        "def preproc_fg(im, orig_size=None):\n",
        "  if orig_size:\n",
        "    im = resize_crop_pad(im, orig_size)\n",
        "  return resize_crop_pad(im, (standard_size[0], standard_size[1]), crop_pad_to_height=True)\n",
        "\n",
        "def preproc_bg(im, flip_x=False, flip_y=False, filters='', gray=False, bokeh=False):\n",
        "  im = do_gmic(im, filters)\n",
        "  \n",
        "  if gray:\n",
        "    im = np.dot(im, [0.2989, 0.5870, 0.1140])[...,None]\n",
        "  \n",
        "  if standard_size[0]*standard_size[1]<im.shape[0]*im.shape[1]:\n",
        "    im = resize_crop_pad(im, standard_size, crop_by_max_ratio=keep_aspect_background)\n",
        "  \n",
        "  if bokeh:\n",
        "    if bokeh_gamma>1:\n",
        "      im = im**bokeh_gamma\n",
        "    radius = int(bokeh_prcnt/100*np.sqrt(im.shape[0]*im.shape[1])//2*2+1)\n",
        "    if radius>1:\n",
        "      im = cv2.GaussianBlur(im,(radius,radius),0)\n",
        "    if bokeh_gamma>1:\n",
        "      im **= 1/bokeh_gamma\n",
        "  \n",
        "  if flip_x:\n",
        "    im = np.fliplr(im)\n",
        "  if flip_y:\n",
        "    im = np.flipud(im)\n",
        "\n",
        "  if standard_size[0]*standard_size[1]>=im.shape[0]*im.shape[1]:\n",
        "    im = resize_crop_pad(im, standard_size, crop_by_max_ratio=keep_aspect_background)\n",
        "  \n",
        "  return fix_dims(im)\n",
        "\n",
        "def u2net_matting(model_name, image_dir, prediction_dir):\n",
        "    def normPRED(d):\n",
        "        ma = torch.max(d)\n",
        "        mi = torch.min(d)\n",
        "        dn = (d-mi)/(ma-mi)\n",
        "        return dn\n",
        "\n",
        "    model_dir = '/content/U-2-Net/saved_models/%s/%s.pth'%(model_name,model_name)\n",
        "\n",
        "    img_name_list = glob.glob(image_dir + os.sep + '*')\n",
        "    #print(img_name_list)\n",
        "\n",
        "    size = u2net_portrait_size if model_name=='u2net_portrait' else u2net_size\n",
        "    ops = [RescaleT(size), ToTensorLab(flag=0)]\n",
        "\n",
        "    test_salobj_dataset = SalObjDataset(img_name_list = img_name_list,\n",
        "                                        lbl_name_list = [],\n",
        "                                        transform=transforms.Compose(ops)\n",
        "                                        )\n",
        "    test_salobj_dataloader = DataLoader(test_salobj_dataset,\n",
        "                                        batch_size=1,\n",
        "                                        shuffle=False,\n",
        "                                        num_workers=1)\n",
        "\n",
        "    if(model_name=='u2netp'):\n",
        "        net = U2NETP(3,1)\n",
        "    else:\n",
        "        net = U2NET(3,1)\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "        net.load_state_dict(torch.load(model_dir))\n",
        "        net.cuda()\n",
        "    else:\n",
        "        net.load_state_dict(torch.load(model_dir, map_location='cpu'))\n",
        "    net.eval()\n",
        "\n",
        "    for i_test, data_test in enumerate(test_salobj_dataloader):\n",
        "        #print(\"inferencing:\",img_name_list[i_test].split(os.sep)[-1])\n",
        "\n",
        "        inputs_test = data_test['image']\n",
        "        inputs_test = inputs_test.type(torch.FloatTensor)\n",
        "\n",
        "        if torch.cuda.is_available():\n",
        "            inputs_test = Variable(inputs_test.cuda())\n",
        "        else:\n",
        "            inputs_test = Variable(inputs_test)\n",
        "\n",
        "        d1,d2,d3,d4,d5,d6,d7 = net(inputs_test)\n",
        "\n",
        "        pred = d1[:,0,:,:]\n",
        "        if model_name=='u2net_portrait':\n",
        "          pred = 1.0 - pred\n",
        "        pred = normPRED(pred).squeeze().cpu().data.numpy()\n",
        "        \n",
        "        file_format='png'\n",
        "        png_compression=9\n",
        "        jpg_quality=95\n",
        "        if len(img_name_list)>1:\n",
        "          file_format=frame_file_format\n",
        "          png_compression=frame_png_compression\n",
        "          jpg_quality=frame_jpg_quality\n",
        "        imwrite(pred, prediction_dir, os.path.basename(img_name_list[i_test]), file_format=file_format, png_compression=png_compression, jpg_quality=jpg_quality)\n",
        "\n",
        "        del d1,d2,d3,d4,d5,d6,d7\n",
        "\n",
        "def modnet_matting(modnet, im):\n",
        "  im_h, im_w = im.shape[:2]\n",
        "  im_tensor = torch_transforms(im).float()\n",
        "  im_tensor = im_tensor[None, :, :, :].cuda()\n",
        "  \n",
        "  if max(im_h, im_w) < modnet_size or min(im_h, im_w) > modnet_size:\n",
        "    if im_w >= im_h:\n",
        "      im_rh = modnet_size\n",
        "      im_rw = int(im_w / im_h * modnet_size)\n",
        "    elif im_w < im_h:\n",
        "      im_rw = modnet_size\n",
        "      im_rh = int(im_h / im_w * modnet_size)\n",
        "  else:\n",
        "    im_rh = im_h\n",
        "    im_rw = im_w\n",
        "        \n",
        "  im_rw = im_rw - im_rw % 32\n",
        "  im_rh = im_rh - im_rh % 32\n",
        "  \n",
        "  if (im_h,im_w)!=(im_rh,im_rw):\n",
        "    im_tensor = F.interpolate(im_tensor, size=(im_rh, im_rw), mode='area')\n",
        "  \n",
        "  with torch.no_grad():\n",
        "    _, _, matte_tensor = modnet(im_tensor, True)\n",
        "  \n",
        "  if (im_h,im_w)!=(im_rh,im_rw):\n",
        "    matte_tensor = F.interpolate(matte_tensor, size=(im_h, im_w), mode='area')\n",
        "  matte_tensor = matte_tensor.repeat(1, 3, 1, 1)\n",
        "  return matte_tensor[0].data.cpu().numpy().transpose(1, 2, 0)\n",
        "\n",
        "def make_file_list(folder, fps, cut_mix_min_secs=0, cut_mix_max_secs=0, cut_mix_phase_secs=0, start_frames=None):\n",
        "  files = [x for x in sorted(os.listdir(folder)) if x.endswith('.'+frame_file_format)]\n",
        "  if not fps or not cut_mix_max_secs or len(files)<=1:\n",
        "    return files\n",
        "  start_frames = start_frames + [len(files)]\n",
        "  i = 1\n",
        "  cuts = []\n",
        "  phase = round(cut_mix_phase_secs*fps)\n",
        "  frames = phase\n",
        "  total = 0\n",
        "  minlen = round(cut_mix_min_secs*fps*2/3)\n",
        "  while len(files):\n",
        "    if not phase or cuts:\n",
        "      frames = np.random.randint(round(cut_mix_min_secs*fps),round(cut_mix_max_secs*fps)+1)\n",
        "    frames = min(frames,len(files))\n",
        "    while total+frames>start_frames[i]:\n",
        "      i += 1\n",
        "    diff1 = total+frames-start_frames[i-1]\n",
        "    diff2 = start_frames[i]-total-frames\n",
        "    if diff1<minlen or 0<diff2<minlen:\n",
        "      if diff1<diff2 and start_frames[i-1]-total>=minlen:\n",
        "        frames = start_frames[i-1]-total\n",
        "      elif start_frames[i-1]+minlen <= start_frames[i]-minlen and start_frames[i-1]>=total:\n",
        "        frames = start_frames[i-1]+minlen-total\n",
        "      else:\n",
        "        frames = start_frames[i]-total\n",
        "    assert frames>0, (len(cuts), len(files), total, frames, total+frames, start_frames[i-1], start_frames[i], diff1, diff2, minlen)\n",
        "    cuts.append(files[:frames])\n",
        "    if np.random.random()<cut_mix_temporal_flip_prob and (not phase or cuts):\n",
        "      cuts[-1] = cuts[-1][::-1]\n",
        "    files = files[frames:]\n",
        "    total += frames\n",
        "  cuts = np.array(cuts)\n",
        "  np.random.shuffle(cuts[bool(phase):])\n",
        "  return [file for sublist in cuts for file in sublist]\n",
        "\n",
        "grand_start = time()\n",
        "start = time()\n",
        "fg_start_frames = []\n",
        "bad_files = []\n",
        "try:\n",
        "  fg_now = imread(foreground_src[0])\n",
        "  imwrite(fg_now,fg_dir,0,file_format='png')\n",
        "except Exception:\n",
        "  for fg_file in foreground_src:\n",
        "    vf_arg = '-vf \"fps=%f\"'%override_fps if override_fps else ''\n",
        "    j = len(glob.glob('%s/frame_*'%fg_dir))\n",
        "    fg_start_frames.append(j)\n",
        "    try:\n",
        "      outfile = os.path.join(fg_dir,'frame_%06d.'+frame_file_format)\n",
        "      !ffmpeg $fg_time_params -i \"$fg_file\" $vf_arg -start_number $j $quality_arg $outfile\n",
        "      if not override_fps:\n",
        "        with imageio.get_reader(fg_file, format='mp4') as reader:\n",
        "          override_fps = reader.get_meta_data()['fps']\n",
        "      assert os.path.exists(outfile%0), 'Cannot find file: %s'%(outfile%0)\n",
        "    except Exception:\n",
        "      bad_files.append(fg_file)\n",
        "  fg_now = imread(fg_dir,0,file_format=frame_file_format)\n",
        "standard_size = resize_crop_pad(fg_now, (override_height, override_width), crop_pad_to_height=True).shape[:2]  \n",
        "\n",
        "bg_phase = max(cut_mix_phase_secs_background,0)\n",
        "fg_phase = max(-cut_mix_phase_secs_background,0)\n",
        "fg_files = make_file_list(fg_dir, override_fps, cut_mix_min_secs_foreground, cut_mix_max_secs_foreground, fg_phase, fg_start_frames)\n",
        "if cut_mix_double_foreground and len(fg_files)>1:\n",
        "  dbl_phase = max(fg_phase+cut_mix_phase_secs_double,0)\n",
        "  dbl_files = make_file_list(fg_dir, override_fps, cut_mix_min_secs_double, cut_mix_max_secs_double, dbl_phase, fg_start_frames)\n",
        "else:\n",
        "  cut_mix_double_foreground = False\n",
        "\n",
        "prepare_time = time()-start\n",
        "\n",
        "start = time()\n",
        "have_u2_mask = False\n",
        "if model.startswith('u2net') and (model!='u2net_portrait' or type(background_src)==list or background_src=='foreground'):\n",
        "  u2net_matting(u2net_variant, fg_dir, mask_dir)\n",
        "  have_u2_mask = True\n",
        "elif model.startswith('modnet'):\n",
        "  modnet = MODNet(backbone_pretrained=False)\n",
        "  modnet = nn.DataParallel(modnet).cuda()\n",
        "  modnet.load_state_dict(torch.load(os.path.join('/content/MODNet/pretrained',model+'_portrait_matting.ckpt')))\n",
        "  modnet.eval()\n",
        "mask_time = time()-start\n",
        "\n",
        "blend_time = 0\n",
        "is_fg = True\n",
        "bg_files = []\n",
        "iter_files = fg_files\n",
        "rounds = 0\n",
        "\n",
        "def extract_file_index(file):\n",
        "  return int(file.rsplit('_',1)[1].split('.')[0])\n",
        "\n",
        "def mat(in_dir, out_dir, bg_mode, orig_dir=None, is_pre=False):\n",
        "  global prepare_time, blend_time, fg_now, is_fg, bg_files, iter_files, rounds, override_fps\n",
        "  rounds += 1\n",
        "  start = time()\n",
        "  bg_start_frames = []\n",
        "  if bg_mode == 'white':\n",
        "    bg = np.ones((*standard_size, 3))\n",
        "  elif bg_mode == 'black':\n",
        "    bg = np.zeros((*standard_size, 3))\n",
        "  elif bg_mode == 'chroma green':\n",
        "    bg = np.full((*standard_size, 3), [0,177/255,64/255])\n",
        "  elif bg_mode == 'chroma blue':\n",
        "    bg = np.full((*standard_size, 3), [0,71/255,187/255])\n",
        "  elif type(bg_mode)==list:\n",
        "    try:\n",
        "      bg = imread(bg_mode[0])\n",
        "      bg = preproc_bg(bg, flip_x=mirror_background, flip_y=False, filters=gmic_filters_background, gray=gray_background, bokeh=bokeh_background)\n",
        "    except Exception:\n",
        "      for bg_file in bg_mode:\n",
        "        vf_arg = '-vf \"fps=%f\"'%override_fps if override_fps else ''\n",
        "        j = len(glob.glob('%s/frame_*'%bg_dir))\n",
        "        bg_start_frames.append(j)\n",
        "        try:\n",
        "          outfile = os.path.join(bg_dir,'frame_%06d.'+frame_file_format)\n",
        "          !ffmpeg $bg_time_params -i \"$bg_file\" $vf_arg -start_number $j $quality_arg $outfile\n",
        "          if not override_fps:\n",
        "            with imageio.get_reader(bg_file, format='mp4') as reader:\n",
        "              override_fps = reader.get_meta_data()['fps']\n",
        "          assert os.path.exists(outfile%0), 'Cannot find file: %s'%(outfile%0)\n",
        "        except Exception:\n",
        "          bad_files.append(bg_file)\n",
        "  elif bg_mode not in ['foreground', 'transparent']:\n",
        "    bg = np.full((*standard_size, 3), [int(bg_mode.lstrip('#')[i:i+2], 16)/255 for i in [0, 2, 4]])\n",
        "  use_dir = bg_dir\n",
        "  use_start_frames = bg_start_frames\n",
        "  if bg_mode=='foreground' and len(fg_files)>1 and not cut_mix_tied_background and (cut_mix_max_secs_foreground or cut_mix_max_secs_background):\n",
        "    use_dir = fg_dir\n",
        "    use_start_frames = fg_start_frames\n",
        "  bg_files = make_file_list(use_dir, override_fps, cut_mix_min_secs_background, cut_mix_max_secs_background, bg_phase, use_start_frames)\n",
        "  prepare_time += time()-start\n",
        "\n",
        "  start = time()\n",
        "  fg_plus = None\n",
        "  mask_plus = None\n",
        "  orig = None\n",
        "  dbl_plus = None\n",
        "  dbl_mask_plus = None\n",
        "  dbl_orig = None\n",
        "  is_fg = len(fg_files)>1 or not bg_files\n",
        "  if is_fg:\n",
        "    iter_files = fg_files\n",
        "  else:\n",
        "    iter_files = bg_files\n",
        "    fg_now = imread(in_dir,0)\n",
        "    if have_u2_mask:\n",
        "      mask = imread(mask_dir,0)\n",
        "    elif model.startswith('modnet'):\n",
        "      mask = modnet_matting(modnet, fg_now)\n",
        "    else:\n",
        "      mask = np.ones_like(fg_now)\n",
        "    if orig_dir is not None:\n",
        "      orig = imread(orig_dir,0)\n",
        "      orig = preproc_fg(orig)\n",
        "    if mirror_foreground: #this will only happen on last pass\n",
        "      fg_now = np.fliplr(fg_now)\n",
        "      mask = np.fliplr(mask)\n",
        "      if orig is not None:\n",
        "        orig = np.fliplr(orig)\n",
        "\n",
        "  j = -1\n",
        "  j_direction = 1\n",
        "  for i,file in enumerate(iter_files):\n",
        "    iter_file_index = extract_file_index(file)\n",
        "    first = i==0 or abs(extract_file_index(iter_files[i-1])-iter_file_index)>1\n",
        "    last = i==len(iter_files)-1 or abs(extract_file_index(iter_files[i+1])-iter_file_index)>1\n",
        "    first_bg = first\n",
        "    if is_fg:\n",
        "      if one_frame_delay and not first and not last:\n",
        "        fg_now = fg_plus\n",
        "        mask_minus = mask_now\n",
        "        mask_now = mask_plus\n",
        "        if fg_now is None:\n",
        "          fg_now = imread(in_dir,file)\n",
        "        if have_u2_mask:\n",
        "          if mask_now is None:\n",
        "            mask_now = imread(mask_dir,file)\n",
        "          mask_plus = imread(mask_dir,fg_files[i+1])\n",
        "        elif model.startswith('modnet'):\n",
        "          if mask_now is None:\n",
        "            mask_now = modnet_matting(modnet, fg_now)\n",
        "          fg_plus = imread(in_dir,fg_files[i+1])\n",
        "          mask_plus = modnet_matting(modnet, fg_plus)\n",
        "        else:\n",
        "          if mask_now is None:\n",
        "            mask_now = np.ones_like(fg_now)\n",
        "          mask_plus = np.ones_like(fg_now)          \n",
        "        mask = mask_now\n",
        "        if mask_plus.shape==mask_minus.shape==mask_now.shape and (not calc_area or (max_mask_area_threshold>=mask_plus.sum()/mask_plus.shape[0]/mask_plus.shape[1]/mask_plus.shape[2]>=min_mask_area_threshold and max_mask_area_threshold>=mask_minus.sum()/mask_minus.shape[0]/mask_minus.shape[1]/mask_minus.shape[2]>=min_mask_area_threshold)):\n",
        "          cond = (np.abs(mask_plus-mask_minus)<=one_frame_delay_threshold) & (np.abs(mask_now-mask_minus)>one_frame_delay_threshold) & (np.abs(mask_now-mask_plus)>one_frame_delay_threshold)\n",
        "          mask = mask_now*(1-cond) + (mask_minus+mask_plus)/2*cond\n",
        "      else:\n",
        "        fg_now = imread(in_dir,file)\n",
        "        if have_u2_mask:\n",
        "          mask_now = imread(mask_dir,file)\n",
        "        elif model.startswith('modnet'):\n",
        "          mask_now = modnet_matting(modnet, fg_now)\n",
        "        else:\n",
        "          mask_now = np.ones_like(fg_now)\n",
        "        mask = mask_now\n",
        "\n",
        "      if orig_dir is not None:  \n",
        "        orig = imread(orig_dir,file)\n",
        "        orig = preproc_fg(orig)\n",
        "\n",
        "      if first:\n",
        "        flip_x_fg = mirror_foreground\n",
        "        flip_y_fg = False\n",
        "        gmic_fg = gmic_filters_foreground\n",
        "        if cut_mix_max_secs_foreground and len(iter_files)>cut_mix_min_secs_foreground*override_fps and (not fg_phase or i):\n",
        "          if np.random.random()<cut_mix_horizontal_flip_prob:\n",
        "            flip_x_fg = not flip_x_fg\n",
        "          if np.random.random()<cut_mix_vertical_flip_prob:\n",
        "            flip_y_fg = not flip_y_fg\n",
        "          gmic_fg = ''\n",
        "          if np.random.random()<cut_mix_random_gmic_filters_prob:\n",
        "            gmic_fg = np.random.choice(cut_mix_random_gmic_filters_foreground)\n",
        "\n",
        "      if not is_pre:\n",
        "        fg_now = do_gmic(fg_now, gmic_fg)\n",
        "\n",
        "        if flip_x_fg:\n",
        "          fg_now = np.fliplr(fg_now)\n",
        "          mask = np.fliplr(mask)\n",
        "          if orig is not None:\n",
        "            orig = np.fliplr(orig)\n",
        "\n",
        "        if flip_y_fg:\n",
        "          fg_now = np.flipud(fg_now)\n",
        "          mask = np.flipud(mask)\n",
        "          if orig is not None:\n",
        "            orig = np.flipud(orig)\n",
        "\n",
        "      if cut_mix_double_foreground and not is_pre:\n",
        "        dbl_file = dbl_files[i]\n",
        "        dbl_file_index = extract_file_index(dbl_file)\n",
        "        first_dbl = i==0 or abs(extract_file_index(dbl_files[i-1])-dbl_file_index)>1\n",
        "        last_dbl = i==len(dbl_files)-1 or abs(extract_file_index(dbl_files[i+1])-dbl_file_index)>1\n",
        "        if one_frame_delay and not first_dbl and not last_dbl:\n",
        "          dbl_now = dbl_plus\n",
        "          dbl_mask_minus = dbl_mask_now\n",
        "          dbl_mask_now = dbl_mask_plus\n",
        "          if dbl_now is None:\n",
        "            dbl_now = imread(in_dir,dbl_file)\n",
        "          if have_u2_mask:\n",
        "            if dbl_mask_now is None:\n",
        "              dbl_mask_now = imread(mask_dir,dbl_file)\n",
        "            dbl_mask_plus = imread(mask_dir,dbl_files[i+1])\n",
        "          elif model.startswith('modnet'):\n",
        "            if dbl_mask_now is None:\n",
        "              dbl_mask_now = modnet_matting(modnet, dbl_now)\n",
        "            dbl_plus = imread(in_dir,dbl_files[i+1])\n",
        "            dbl_mask_plus = modnet_matting(modnet, dbl_plus)\n",
        "          else:\n",
        "            if dbl_mask_now is None:\n",
        "              dbl_mask_now = np.ones_like(dbl_now)\n",
        "            dbl_mask_plus = np.ones_like(dbl_now)          \n",
        "          dbl_mask = dbl_mask_now\n",
        "          if dbl_mask_plus.shape==dbl_mask_minus.shape==dbl_mask_now.shape and (not calc_area or (max_mask_area_threshold>=dbl_mask_plus.sum()/dbl_mask_plus.shape[0]/dbl_mask_plus.shape[1]/dbl_mask_plus.shape[2]>=min_mask_area_threshold and min_mask_area_threshold>=dbl_mask_minus.sum()/dbl_mask_minus.shape[0]/dbl_mask_minus.shape[1]/dbl_mask_minus.shape[2]>=min_mask_area_threshold)):\n",
        "            cond = (np.abs(dbl_mask_plus-dbl_mask_minus)<=one_frame_delay_threshold) & (np.abs(dbl_mask_now-dbl_mask_minus)>one_frame_delay_threshold) & (np.abs(dbl_mask_now-dbl_mask_plus)>one_frame_delay_threshold)\n",
        "            dbl_mask = dbl_mask_now*(1-cond) + (dbl_mask_minus+dbl_mask_plus)/2*cond\n",
        "        else:\n",
        "          dbl_now = imread(in_dir,dbl_file)\n",
        "          if have_u2_mask:\n",
        "            dbl_mask_now = imread(mask_dir,dbl_file)\n",
        "          elif model.startswith('modnet'):\n",
        "            dbl_mask_now = modnet_matting(modnet, dbl_now)\n",
        "          else:\n",
        "            dbl_mask_now = np.ones_like(dbl_now)\n",
        "          dbl_mask = dbl_mask_now\n",
        "\n",
        "        if orig_dir is not None:  \n",
        "          dbl_orig = imread(orig_dir,dbl_file)\n",
        "          dbl_orig = preproc_fg(dbl_orig)\n",
        "\n",
        "        if first_dbl:\n",
        "          flip_x_dbl = cut_mix_mirror_double\n",
        "          flip_y_dbl = False\n",
        "          gmic_dbl = cut_mix_gmic_filters_double\n",
        "          if cut_mix_max_secs_double and len(dbl_files)>cut_mix_min_secs_double*override_fps and (not dbl_phase or i):\n",
        "            if np.random.random()<cut_mix_horizontal_flip_prob:\n",
        "              flip_x_dbl = not flip_x_dbl\n",
        "            if np.random.random()<cut_mix_vertical_flip_prob:\n",
        "              flip_y_dbl = not flip_y_dbl\n",
        "            gmic_dbl = ''\n",
        "            if np.random.random()<cut_mix_random_gmic_filters_prob:\n",
        "              gmic_dbl = np.random.choice(cut_mix_random_gmic_filters_double)\n",
        "          \n",
        "        if not is_pre:\n",
        "          dbl_now = do_gmic(dbl_now, gmic_dbl)\n",
        "          \n",
        "          if flip_x_dbl:\n",
        "            dbl_now = np.fliplr(dbl_now)\n",
        "            dbl_mask = np.fliplr(dbl_mask)\n",
        "            if dbl_orig is not None:\n",
        "              dbl_orig = np.fliplr(dbl_orig)\n",
        "\n",
        "          if flip_y_dbl:\n",
        "            dbl_now = np.flipud(dbl_now)\n",
        "            dbl_mask = np.flipud(dbl_mask)\n",
        "            if dbl_orig is not None:\n",
        "              dbl_orig = np.flipud(dbl_orig)\n",
        "\n",
        "      if bg_mode=='foreground' and (len(fg_files)==1 or cut_mix_tied_background or not cut_mix_max_secs_foreground and not cut_mix_max_secs_background):\n",
        "        if orig is not None:\n",
        "          bg = orig\n",
        "        else:\n",
        "          bg = preproc_fg(fg_now)\n",
        "      elif bg_files:\n",
        "        if loop_reverse_background: \n",
        "          j += j_direction\n",
        "          if j>=len(bg_files):\n",
        "            j = 2*len(bg_files)-j-1\n",
        "            j_direction = -1\n",
        "          elif j<0:\n",
        "            j = 0\n",
        "            j_direction = 1\n",
        "        else:\n",
        "          j = i%len(bg_files)\n",
        "        bg = imread(use_dir,bg_files[j])\n",
        "        bg_file_index = extract_file_index(bg_files[j])\n",
        "        if j_direction>0:\n",
        "          first_bg = j==0 or abs(extract_file_index(bg_files[j-1])-bg_file_index)>1\n",
        "        else:\n",
        "          first_bg = j==len(bg_files)-1 or abs(extract_file_index(bg_files[j+1])-bg_file_index)>1\n",
        "    \n",
        "    else:\n",
        "      bg = imread(bg_dir,file)\n",
        "\n",
        "    if bg_files:\n",
        "      if first_bg:\n",
        "        flip_x_bg = mirror_background\n",
        "        flip_y_bg = False\n",
        "        gmic_bg = gmic_filters_background\n",
        "        if cut_mix_max_secs_background and len(bg_files)>cut_mix_min_secs_background*override_fps:\n",
        "          if np.random.random()<cut_mix_horizontal_flip_prob:\n",
        "            flip_x_bg = not flip_x_bg\n",
        "          if np.random.random()<cut_mix_vertical_flip_prob:\n",
        "            flip_y_bg = not flip_y_bg\n",
        "          gmic_bg = ''\n",
        "          if np.random.random()<cut_mix_random_gmic_filters_prob:\n",
        "            gmic_bg = np.random.choice(cut_mix_random_gmic_filters_background)\n",
        "      bg = preproc_bg(bg, flip_x=flip_x_bg, flip_y=flip_y_bg, filters=gmic_bg, gray=gray_background, bokeh=bokeh_background)\n",
        "    if is_fg or i==0:\n",
        "      fg_area_ok = True\n",
        "      if calc_area or cut_mix_double_foreground and not is_pre and cut_mix_layer_order_double.startswith('front if'):\n",
        "        fg_area = mask.sum()/mask.shape[0]/mask.shape[1]/mask.shape[2]\n",
        "        fg_area_ok = max_mask_area_threshold>=fg_area>=min_mask_area_threshold\n",
        "\n",
        "      fg = fg_now\n",
        "\n",
        "      if refine_foreground_blend and fg_area_ok:\n",
        "        try:\n",
        "          fg = estimate_foreground_ml(fg, mask[...,0])\n",
        "        except Exception as e:\n",
        "          pass\n",
        "          #print(e)\n",
        "          #imwrite(fg, '/content/debug_fg', file_format=frame_file_format, png_compression=frame_png_compression, jpg_quality=frame_jpg_quality)\n",
        "          #imwrite(mask, '/content/debug_mask', file_format=frame_file_format, png_compression=frame_png_compression, jpg_quality=frame_jpg_quality)\n",
        "\n",
        "      orig_size = fg.shape[:2]\n",
        "      fg = preproc_fg(fg)\n",
        "      mask = preproc_fg(mask, orig_size)\n",
        "\n",
        "      if orig is not None and fg_area_ok:\n",
        "        if sketch_color=='foreground':\n",
        "          fg = 1-(1-fg)*(1-orig)\n",
        "        elif 'tint' in sketch_color:\n",
        "          non_black_mask = np.any(mask != [0, 0, 0], axis=-1)\n",
        "          colors = (orig*mask)[non_black_mask]\n",
        "          chroma = colors.max(axis=-1)-colors.min(axis=-1)\n",
        "          chroma_thresholds.sort(reverse=True)\n",
        "          if sketch_color=='tint fill':\n",
        "            chroma_thresholds.append(0)\n",
        "          else:\n",
        "            color = np.array([0,0,0])\n",
        "          for threshold in chroma_thresholds:\n",
        "            cond = chroma>=threshold\n",
        "            if np.any(cond):\n",
        "              unique, counts = np.unique(colors[cond], axis=0, return_counts=True)\n",
        "              color = unique[np.argmax(counts)]\n",
        "              break\n",
        "          if sketch_color=='tint outline':\n",
        "            fg = 1-(1-fg)*(1-color)\n",
        "          elif sketch_color=='tint fill':\n",
        "            fg = fg*color\n",
        "\n",
        "      if cut_mix_double_foreground and not is_pre:\n",
        "          \n",
        "          if calc_area or cut_mix_layer_order_double.startswith('front if'):\n",
        "            dbl_area = dbl_mask.sum()/dbl_mask.shape[0]/dbl_mask.shape[1]/dbl_mask.shape[2]\n",
        "\n",
        "          if not calc_area or max_mask_area_threshold>=dbl_area>=min_mask_area_threshold:\n",
        "            \n",
        "            dbl = dbl_now\n",
        "\n",
        "            if refine_foreground_blend:\n",
        "              try:\n",
        "                dbl = estimate_foreground_ml(dbl, dbl_mask[...,0])\n",
        "              except Exception:\n",
        "                pass\n",
        "\n",
        "            orig_size = dbl.shape[:2]\n",
        "            dbl = preproc_fg(dbl)\n",
        "            dbl_mask = preproc_fg(dbl_mask, orig_size)\n",
        "\n",
        "            if dbl_orig is not None:\n",
        "              if cut_mix_double_sketch_color=='foreground':\n",
        "                dbl = 1-(1-dbl)*(1-dbl)\n",
        "              elif 'tint' in cut_mix_double_sketch_color:\n",
        "                non_black_mask = np.any(dbl_mask != [0, 0, 0], axis=-1)\n",
        "                colors = (dbl_orig*dbl_mask)[non_black_mask]\n",
        "                chroma = colors.max(axis=-1)-colors.min(axis=-1)\n",
        "                chroma_thresholds.sort(reverse=True)\n",
        "                if cut_mix_double_sketch_color=='tint fill':\n",
        "                  chroma_thresholds.append(0)\n",
        "                else:\n",
        "                  color = np.array([0,0,0])\n",
        "                for threshold in chroma_thresholds:\n",
        "                  cond = chroma>=threshold\n",
        "                  if np.any(cond):\n",
        "                    unique, counts = np.unique(colors[cond], axis=0, return_counts=True)\n",
        "                    color = unique[np.argmax(counts)]\n",
        "                    break\n",
        "                if cut_mix_double_sketch_color=='tint outline':\n",
        "                  dbl = 1-(1-dbl)*(1-color)\n",
        "                elif cut_mix_double_sketch_color=='tint fill':\n",
        "                  dbl = dbl*color\n",
        "            \n",
        "            if cut_mix_layer_order_double.startswith('front'):\n",
        "              if fg_area_ok and (cut_mix_layer_order_double=='front' or cut_mix_layer_order_double=='front if smaller mask' and dbl_area<fg_area or cut_mix_layer_order_double=='front if larger mask' and dbl_area>fg_area):\n",
        "                fg, dbl = dbl, fg\n",
        "                mask, dbl_mask = dbl_mask, mask\n",
        "\n",
        "            if bg_mode == 'transparent':\n",
        "              bg = np.dstack([dbl,dbl_mask[...,0]])\n",
        "            else:\n",
        "              bg = bg*(1-dbl_mask)+dbl*dbl_mask\n",
        "\n",
        "    if bg_mode == 'transparent':\n",
        "      im = fg\n",
        "      if bg is not None:\n",
        "        im = bg\n",
        "        if fg_area_ok:\n",
        "          fg = fg*mask + dbl*dbl_mask*(1-mask)\n",
        "          mask = mask + dbl_mask*(1-mask)\n",
        "          fg = np.divide(fg,mask,out=mask,where=mask>0)  \n",
        "      if bg is None or bg is not None and fg_area_ok:\n",
        "        im = np.dstack([fg,mask[...,0]*fg_area_ok])\n",
        "    else:\n",
        "      im = bg\n",
        "      if fg_area_ok:\n",
        "        im = bg*(1-mask)+fg*mask\n",
        "    \n",
        "    if not override_width and not override_height and is_fg and bg_mode_max_w and fg.shape[1]>bg_mode_max_w: # prevent creation of too high resolutions videos due to a high resolution foreground image\n",
        "      im = resize_crop_pad(im, (-1,bg_mode_max_w))\n",
        "    if len(fg_files)>1 or bg_files:\n",
        "      imwrite(im, out_dir, file if is_pre else i, file_format=frame_file_format, png_compression=frame_png_compression, jpg_quality=frame_jpg_quality)\n",
        "    else:\n",
        "      imwrite(im, out_dir, file if is_pre else i, file_format='png', png_compression=9, jpg_quality=95)\n",
        "    print('%d/%d (%d/%d)'%(i+1,len(iter_files),rounds,2 if is_pre else rounds))\n",
        "  blend_time += time()-start\n",
        "\n",
        "out_dir = result_dir\n",
        "bg_mode = background_src\n",
        "if 'u2net_portrait' in model:\n",
        "  out_dir = portrait_in_dir\n",
        "  bg_mode = 'white' \n",
        "if model!='u2net_portrait':\n",
        "  mat(fg_dir, out_dir, bg_mode, is_pre='+' in model)\n",
        "elif 'u2net_portrait' in model:\n",
        "  if background_src=='foreground' or sketch_color!='gray':\n",
        "    !cp $fg_dir/* $portrait_in_dir\n",
        "  else:\n",
        "    !mv $fg_dir/* $portrait_in_dir\n",
        "if 'u2net_portrait' in model:\n",
        "  start = time()\n",
        "  u2net_matting('u2net_portru2ait', portrait_in, portrait_out)\n",
        "  blend_time += time()-start\n",
        "  if type(background_src)==list or background_src=='foreground' or sketch_color!='gray':\n",
        "    mat(portrait_out_dir, result_dir, background_src, orig_dir=fg_dir if background_src=='foreground' or sketch_color!='gray' else None)\n",
        "  else:\n",
        "    !mv $portrait_out_dir/* $result_dir\n",
        "\n",
        "start = time()\n",
        "from IPython.display import HTML, clear_output, Image\n",
        "import shutil\n",
        "!rm -f /content/final.mp4\n",
        "!rm -f /content/final.png\n",
        "date = datetime.utcnow().strftime('%Y%m%d%H%M%S')\n",
        "if len(fg_files)>1 or len(bg_files)>1:\n",
        "  if not copy_audio:\n",
        "    !ffmpeg -framerate $override_fps -i $result_dir/frame_%06d.$frame_file_format -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",
        "  elif is_fg:\n",
        "      if len(foreground_src)==1:\n",
        "        !ffmpeg -framerate $override_fps -i $result_dir/frame_%06d.$frame_file_format $fg_time_params -i '{foreground_src[0]}' -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",
        "        with open('/content/list.txt','w',encoding='utf8') as f:\n",
        "          if fg_time_params:\n",
        "            start_seconds = float(fg_time_params.split(' ')[1])\n",
        "            end_seconds = start_seconds + float(fg_time_params.split(' ')[3])\n",
        "          for file in foreground_src:\n",
        "            f.write(\"file '%s'\\n\"%file)\n",
        "            if fg_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/audio.mp4 -y\n",
        "        !ffmpeg -framerate $override_fps -i $result_dir/frame_%06d.$frame_file_format -i /content/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",
        "  elif type(background_src)!=list or len(background_src)==1:\n",
        "        !ffmpeg -framerate $override_fps -i $result_dir/frame_%06d.$frame_file_format $bg_time_params -i '{background_src[0]}' -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",
        "    with open('/content/list.txt','w',encoding='utf8') 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 background_src:\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/audio.mp4 -y\n",
        "    !ffmpeg -framerate $override_fps -i $result_dir/frame_%06d.$frame_file_format -i /content/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",
        "  #video can be downloaded from /content/final.mp4\n",
        "  save_time = time()-start\n",
        "  total_time = time()-grand_start  \n",
        "  if os.path.exists('/content/drive/MyDrive'):\n",
        "    !mkdir -p /content/drive/MyDrive/vid2green_output\n",
        "    drive_output = '/content/drive/MyDrive/vid2green_output/%s.mp4'%date\n",
        "    !cp /content/final.mp4 $drive_output\n",
        "  clear_output()\n",
        "  # robust display for large videos (https://stackoverflow.com/questions/67591072/displaying-large-video-files-in-google-colab): \n",
        "  !pkill -f \"python3 -m http.server\"\n",
        "  !nohup python3 -m http.server -d /content 8000 &>/dev/null &\n",
        "  display(HTML(\"\"\"\n",
        "  <video width=600 controls autoplay loop>\n",
        "    <source src=\"https://localhost:8000/final.mp4\" type=\"video/mp4\">\n",
        "  </video>\"\"\"))\n",
        "else:\n",
        "  shutil.move(os.path.join(out_dir,'frame_%06d.jpg'%1), '/content/final.jpg')\n",
        "  #image can be downloaded from /content/final.png\n",
        "  save_time = time()-start\n",
        "  total_time = time()-grand_start\n",
        "  if os.path.exists('/content/drive/MyDrive'):\n",
        "    !mkdir -p /content/drive/MyDrive/vid2green_output\n",
        "    drive_output = '/content/drive/MyDrive/vid2green_output/%s.jpg'%date\n",
        "    !cp /content/final.png $drive_output\n",
        "  clear_output()\n",
        "  display(Image('/content/final.jpg', width=600))\n",
        "if bad_files:\n",
        "  print('bad files:',bad_files)\n",
        "if os.path.exists('/content/drive/MyDrive'):\n",
        "  print('saved output to: %s'%drive_output)\n",
        "if model.startswith('u2net'):\n",
        "  print('%d frames took: prepare=%.1f mask=%d blend=%.1f save=%.1f total=%.1f minutes'%(len(iter_files), prepare_time/60, mask_time/60, blend_time/60, save_time/60, total_time/60))\n",
        "else:\n",
        "  print('%d frames took: prepare=%.1f mask+blend=%.1f save=%.1f total=%.1f minutes'%(len(iter_files), prepare_time/60, (mask_time+blend_time)/60, save_time/60, total_time/60))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zxy4hscknBvt",
        "cellView": "form"
      },
      "source": [
        "#@title Download\n",
        "from google.colab import files\n",
        "if os.path.exists('/content/final.mp4'):\n",
        "  files.download('/content/final.mp4')\n",
        "else:\n",
        "  files.download('/content/final.png')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}