Colab_DAIN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colab_DAIN.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pIo4r_Y8cMo",
        "colab_type": "text"
      },
      "source": [
        "# DAIN Colab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGPHW5SOpPe3",
        "colab_type": "text"
      },
      "source": [
        "*DAIN Colab, v1.5.0*\n",
        "\n",
        "Based on the [original Colab file](https://github.com/baowenbo/DAIN/issues/44) by btahir. \n",
        "\n",
        "Enhancements by [Styler00Dollar](https://github.com/styler00dollar) aka \"sudo rm -rf / --no-preserve-root#8353\" on discord and [Alpha](https://github.com/AlphaGit), (Alpha#6137 on Discord). Please do not run this command in your linux terminal. It's rather meant as a joke.\n",
        "\n",
        "[Styler00Dollar's fork](https://github.com/styler00dollar/DAIN) / [Alpha's fork](https://github.com/AlphaGit/DAIN)\n",
        "\n",
        "A simple guide:\n",
        "- Upload this ` .ipynb`  file to your Google Colab.\n",
        "- Create a folder inside of Google Drive named \"DAIN\"\n",
        "- Change the configurations in the next cell\n",
        "- Run cells one by one\n",
        "\n",
        "Stuff that should be improved:\n",
        "- Alpha channel will be removed automatically and won't be added back. Anything related to alpha will be converted to black.\n",
        "- Adding configuration to select speed\n",
        "- Detect scenes to avoid interpolating scene-changes\n",
        "- Auto-resume\n",
        "- Copy `start_frame` - `end_frame` audio from original input to final output\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "enKoi0TR2fOD",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "################# Required Configurations ############################\n",
        "\n",
        "#@markdown # Required Configuration\n",
        "#@markdown Use the values in here to configure what you'd like DAIN to do.\n",
        "\n",
        "#@markdown ## Input file\n",
        "#@markdown Path (relative to the root of your Google Drive) to the input file. For instance, if you save your `example.mkv` file in your Google Drive, inside a `videos` folder, the path would be: `videos/example.mkv`. Currenly videos and gifs are supported.\n",
        "INPUT_FILEPATH = \"DAIN/input.mp4\" #@param{type:\"string\"}\n",
        "\n",
        "#@markdown ## Output file\n",
        "#@markdown Output file path: path (relative to the root of your Google Drive) for the output file. It will also determine the filetype in the destination. `.mp4` is recommended for video input, `.gif` for gif inputs.\n",
        "OUTPUT_FILE_PATH = \"DAIN/output.mp4\" #@param{type:\"string\"}\n",
        "\n",
        "################# Optional configurations ############################\n",
        "\n",
        "#@markdown # Optional Configuration\n",
        "#@markdown Parameters below can be left with their defaults, but feel free to adapt them to your needs.\n",
        "\n",
        "#@markdown ## Target FPS\n",
        "#@markdown  how many frames per second should the result have. This will determine how many intermediate images are interpolated.\n",
        "TARGET_FPS = 60 #@param{type:\"number\"}\n",
        "\n",
        "#@markdown ## Frame input directory\n",
        "#@markdown A path, relative to your GDrive root, where you already have the list of frames in the format 00001.png, 00002.png, etc.\n",
        "FRAME_INPUT_DIR = '/content/DAIN/input_frames' #@param{type:\"string\"}\n",
        "\n",
        "#@markdown ## Frame output directory\n",
        "#@markdown A path, relative to your GDrive root, where you want the generated frame.\n",
        "FRAME_OUTPUT_DIR = '/content/DAIN/output_frames' #@param{type:\"string\"}\n",
        "\n",
        "#@markdown ## Start Frame\n",
        "#@markdown First frame to consider from the video when processing.\n",
        "START_FRAME = 1 #@param{type:\"number\"}\n",
        "\n",
        "#@markdown ## End Frame\n",
        "#@markdown Last frame to consider from the video when processing. To use the whole video use `-1`.\n",
        "END_FRAME = -1 #@param{type:\"number\"}\n",
        "\n",
        "#@markdown ## Seamless playback\n",
        "#@markdown Creates a seamless loop by using the first frame as last one as well. Set this to True this if loop is intended.\n",
        "SEAMLESS = False #@param{type:\"boolean\"}\n",
        "\n",
        "#@markdown ## Resize hotfix\n",
        "#@markdown DAIN frames are a bit \"shifted / smaller\" compared to original input frames. This can partly be mitigated with resizing DAIN frames to the resolution +2px and cropping the result to the original resoultion with the starting point (1,1). Without this fix, DAIN tends to make \"vibrating\" output and it is pretty noticible with static elements like text.\n",
        "#@markdown\n",
        "#@markdown This hotfix tries to make such effects less visible for a smoother video playback. I do not know what DAINAPP uses as a fix for this problem, but the original does show such behaviour with the default test images. More advanced users can change the interpolation method. The methods cv2.INTER_CUBIC and cv2.INTER_LANCZOS4 are recommended. The current default value is cv2.INTER_LANCZOS4.\n",
        "RESIZE_HOTFIX = True #@param{type:\"boolean\"}\n",
        "\n",
        "#@markdown ## Auto-remove PNG directory\n",
        "#@markdown Auto-delete output PNG dir after ffmpeg video creation. Set this to `False` if you want to keep the PNG files.\n",
        "AUTO_REMOVE = True #@param{type:\"boolean\"}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N9cGwalNeyk9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Connect Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "print('Google Drive connected.')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "irzjv1x4e3S4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Check your current GPU\n",
        "# If you are lucky, you get 16GB VRAM. If you are not lucky, you get less. VRAM is important. The more VRAM, the higher the maximum resolution will go.\n",
        "\n",
        "# 16GB: Can handle 720p. 1080p will procude an out-of-memory error. \n",
        "# 8GB: Can handle 480p. 720p will produce an out-of-memory error.\n",
        "\n",
        "!nvidia-smi --query-gpu=gpu_name,driver_version,memory.total --format=csv"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYHTTP91oMvh",
        "colab_type": "text"
      },
      "source": [
        "# Install dependencies.\n",
        "\n",
        "This next step may take somewhere between 15-20 minutes. Run this only once at startup.\n",
        "\n",
        "Look for the \"Finished installing dependencies\"  message."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XunYRBtyPZf_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Install known used versions of PyTorch and SciPy\n",
        "!pip install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "!pip install scipy==1.1.0\n",
        "\n",
        "!CUDA_VISIBLE_DEVICES=0\n",
        "!sudo apt-get install imagemagick imagemagick-doc\n",
        "print(\"Finished installing dependencies.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-mRVRXIGPk9I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Clone DAIN sources\n",
        "%cd /content\n",
        "!git clone -b master --depth 1 https://github.com/baowenbo/DAIN /content/DAIN\n",
        "%cd /content/DAIN\n",
        "!git log -1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e5AHGetTRacZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# This takes a while. Just wait. ~15 minutes.\n",
        "# Building DAIN.\n",
        "%cd /content/DAIN/my_package/\n",
        "!./build.sh\n",
        "print(\"Building #1 done.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UeaU8um5-2NS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Wait again. ~5 minutes.\n",
        "# Building DAIN PyTorch correlation package.\n",
        "%cd /content/DAIN/PWCNet/correlation_package_pytorch1_0\n",
        "!./build.sh\n",
        "print(\"Building #2 done.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "InjqMYIyXCZs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Downloading pre-trained model\n",
        "%cd /content/DAIN\n",
        "!mkdir model_weights\n",
        "!wget -O model_weights/best.pth http://vllab1.ucmerced.edu/~wenbobao/DAIN/best.pth"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zm5kn6vTncL4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Detecting FPS of input file.\n",
        "%shell yes | cp -f /content/gdrive/My\\ Drive/{INPUT_FILEPATH} /content/DAIN/\n",
        "\n",
        "import os\n",
        "filename = os.path.basename(INPUT_FILEPATH)\n",
        "\n",
        "import cv2\n",
        "cap = cv2.VideoCapture(f'/content/DAIN/{filename}')\n",
        "\n",
        "fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "print(f\"Input file has {fps} fps\")\n",
        "\n",
        "if(fps/TARGET_FPS>0.5):\n",
        "  print(\"Define a higher fps, because there is not enough time for new frames. (Old FPS)/(New FPS) should be lower than 0.5. Interpolation will fail if you try.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9YNva-GuKq4Y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ffmpeg extract - Generating individual frame PNGs from the source file.\n",
        "%shell rm -rf '{FRAME_INPUT_DIR}'\n",
        "%shell mkdir -p '{FRAME_INPUT_DIR}'\n",
        "\n",
        "if (END_FRAME==-1):\n",
        "  %shell ffmpeg -i '/content/DAIN/{filename}' -vf 'select=gte(n\\,{START_FRAME}),setpts=PTS-STARTPTS' '{FRAME_INPUT_DIR}/%05d.png'\n",
        "else:\n",
        "  %shell ffmpeg -i '/content/DAIN/{filename}' -vf 'select=between(n\\,{START_FRAME}\\,{END_FRAME}),setpts=PTS-STARTPTS' '{FRAME_INPUT_DIR}/%05d.png'\n",
        "\n",
        "from IPython.display import clear_output\n",
        "clear_output()\n",
        "\n",
        "png_generated_count_command_result = %shell ls '{FRAME_INPUT_DIR}' | wc -l\n",
        "frame_count = int(png_generated_count_command_result.output.strip())\n",
        "\n",
        "import shutil\n",
        "if SEAMLESS:\n",
        "  frame_count += 1\n",
        "  first_frame = f\"{FRAME_INPUT_DIR}/00001.png\"\n",
        "  new_last_frame = f\"{FRAME_INPUT_DIR}/{frame_count.zfill(5)}.png\"\n",
        "  shutil.copyfile(first_frame, new_last_frame)\n",
        "\n",
        "print(f\"{frame_count} frame PNGs generated.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M74tIdzSne3x",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Checking if PNGs do have alpha\n",
        "import subprocess as sp\n",
        "%cd {FRAME_INPUT_DIR}\n",
        "channels = sp.getoutput('identify -format %[channels] 00001.png')\n",
        "print (f\"{channels} detected\")\n",
        "\n",
        "# Removing alpha if detected\n",
        "if \"a\" in channels:\n",
        "  print(\"Alpha channel detected and will be removed.\")\n",
        "  print(sp.getoutput('find . -name \"*.png\" -exec convert \"{}\" -alpha off PNG24:\"{}\" \\;'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W3rrE7L824gL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Interpolation\n",
        "%shell mkdir -p '{FRAME_OUTPUT_DIR}'\n",
        "%cd /content/DAIN\n",
        "\n",
        "!python -W ignore colab_interpolate.py --netName DAIN_slowmotion --time_step {fps/TARGET_FPS} --start_frame 1 --end_frame {frame_count} --frame_input_dir '{FRAME_INPUT_DIR}' --frame_output_dir '{FRAME_OUTPUT_DIR}'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xPYewi8jsq_B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Finding DAIN Frames, upscaling and cropping to match original\n",
        "%cd {FRAME_OUTPUT_DIR}\n",
        "\n",
        "if (RESIZE_HOTFIX):\n",
        "  images = []\n",
        "  for filename in os.listdir(FRAME_OUTPUT_DIR):\n",
        "    img = cv2.imread(os.path.join(FRAME_OUTPUT_DIR, filename))\n",
        "    filename = os.path.splitext(filename)[0]\n",
        "    if(not filename.endswith('0')):\n",
        "      dimensions = (img.shape[1]+2, img.shape[0]+2)\n",
        "      resized = cv2.resize(img, dimensions, interpolation=cv2.INTER_LANCZOS4)\n",
        "      crop = resized[1:(dimensions[1]-1), 1:(dimensions[0]-1)]\n",
        "      cv2.imwrite(f\"{filename}.png\", crop)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TKREDli2IDMV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create output video\n",
        "%cd {FRAME_OUTPUT_DIR}\n",
        "%shell ffmpeg -y -r {TARGET_FPS} -f image2 -pattern_type glob -i '*.png' '/content/gdrive/My Drive/{OUTPUT_FILE_PATH}'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UF5TEo5N374o",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "if(AUTO_REMOVE):\n",
        "  !rm -rf {FRAME_OUTPUT_DIR}/*"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uCBHIXNN-JYu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# [Experimental] Create video with sound\n",
        "# Only run this, if the original had sound.\n",
        "%cd {FRAME_OUTPUT_DIR}\n",
        "%shell ffmpeg -i '/content/DAIN/{filename}' -acodec copy output-audio.aac\n",
        "%shell ffmpeg -y -r {TARGET_FPS} -f image2 -pattern_type glob -i '*.png' -i output-audio.aac -shortest '/content/gdrive/My Drive/{OUTPUT_FILE_PATH}'\n",
        "\n",
        "if (AUTO_REMOVE):\n",
        "  !rm -rf {FRAME_OUTPUT_DIR}/*\n",
        "  !rm -rf output-audio.aac"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}