Colab_DAIN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colab_DAIN_new.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pIo4r_Y8cMo"
      },
      "source": [
        "# DAIN Colab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGPHW5SOpPe3"
      },
      "source": [
        "*DAIN Colab, v1.6.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",
        "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 ## 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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N9cGwalNeyk9",
        "cellView": "form"
      },
      "source": [
        "#@title Connect Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "print('Google Drive connected.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "irzjv1x4e3S4",
        "cellView": "form"
      },
      "source": [
        "#@title 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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYHTTP91oMvh"
      },
      "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": "e5AHGetTRacZ",
        "cellView": "form"
      },
      "source": [
        "#@title Setup everything. This takes a while. Just wait ~20 minutes in total.\n",
        "\n",
        "# Install old pytorch to avoid faulty output\n",
        "%cd /content/\n",
        "!wget -c https://repo.anaconda.com/miniconda/Miniconda3-4.5.4-Linux-x86_64.sh\n",
        "!chmod +x Miniconda3-4.5.4-Linux-x86_64.sh\n",
        "!bash ./Miniconda3-4.5.4-Linux-x86_64.sh -b -f -p /usr/local\n",
        "!conda install pytorch==1.1 cudatoolkit torchvision -c pytorch -y\n",
        "!conda install ipykernel -y\n",
        "\n",
        "!pip install scipy==1.1.0\n",
        "!pip install imageio\n",
        "!CUDA_VISIBLE_DEVICES=0\n",
        "!sudo apt-get install imagemagick imagemagick-doc\n",
        "print(\"Finished installing dependencies.\")\n",
        "\n",
        "# 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\n",
        "\n",
        "# Building DAIN\n",
        "%cd /content/DAIN/my_package/\n",
        "!./build.sh\n",
        "print(\"Building #1 done.\")\n",
        "\n",
        "# Building DAIN PyTorch correlation package.\n",
        "%cd /content/DAIN/PWCNet/correlation_package_pytorch1_0\n",
        "!./build.sh\n",
        "print(\"Building #2 done.\")\n",
        "\n",
        "# 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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zm5kn6vTncL4",
        "cellView": "form"
      },
      "source": [
        "#@title 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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9YNva-GuKq4Y",
        "cellView": "form"
      },
      "source": [
        "#@title 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.\")\n",
        "\n",
        "#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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W3rrE7L824gL",
        "cellView": "form"
      },
      "source": [
        "#@title 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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TKREDli2IDMV",
        "cellView": "form"
      },
      "source": [
        "#@title 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}'\n",
        "\n",
        "if(AUTO_REMOVE):\n",
        "  !rm -rf {FRAME_OUTPUT_DIR}/*\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UF5TEo5N374o",
        "cellView": "form"
      },
      "source": [
        "#@title [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": null,
      "outputs": []
    }
  ]
}