Colab end-to-end runnable version

Colab_DAIN.ipynb

@@ -0,0 +1,366 @@
+{
+  "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": [
+        "Credits to original Colab file: \n",
+        "https://github.com/baowenbo/DAIN/issues/44\n",
+        "\n",
+        "My fork:\n",
+        "https://github.com/styler00dollar/DAIN\n",
+        "\n",
+        "Enhancement by [Styler00Dollar](https://github.com/styler00dollar) aka \"sudo rm -rf / --no-preserve-root#8353\" on discord and [Alpha](https://github.com/AlphaGit). Please do not run this command in your linux terminal. It's rather meant as a joke.\n",
+        "\n",
+        "A simple guide:\n",
+        "- Copy the .ipynb-file to your drive.\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"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "enKoi0TR2fOD",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "################# Configurations ############################\n",
+        "\n",
+        "# Use the values in here to configure what you'd like DAIN to do.\n",
+        "\n",
+        "# Input file: Path (relative to the root of your Google Drive) to the input file.\n",
+        "# For instance, if you save your \"example.mkv\" file in your Google Drive, inside a \"videos\" folder, the path would be:\n",
+        "# videos/example.mkv. Currenly videos and gifs are supported.\n",
+        "INPUT_FILEPATH = \"DAIN/input.mp4\"\n",
+        "\n",
+        "# Output file path: path (relative to the root of your Google Drive) for the output file. It will also determine the\n",
+        "# filetype in the destination. MP4 is recommended.\n",
+        "OUTPUT_FILE_PATH = \"DAIN/output.mp4\"\n",
+        "\n",
+        "################# Optional configurations ############################\n",
+        "\n",
+        "# Target FPS = how many frames per second should the result have. This will determine how many intermediate images are\n",
+        "# interpolated.\n",
+        "TARGET_FPS = 60\n",
+        "\n",
+        "# Frame input directly\n",
+        "# Use a path that is in your GDrive if you already have the list of frames in the format 00001.png, 00002.png, etc.\n",
+        "# Your GDrive is located at `/content/gdrive/My Drive/`\n",
+        "FRAME_INPUT_DIR = '/content/DAIN/input_frames'\n",
+        "\n",
+        "# Frame output directory\n",
+        "# Use a location in your GDrive if you want the generated frames stored to your Google Drive.\n",
+        "# Your GDrive is located at `/content/gdrive/My Drive/DAIN/tmp`\n",
+        "FRAME_OUTPUT_DIR = '/content/DAIN/output_frames'\n",
+        "\n",
+        "# Seamless playback\n",
+        "# 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\n",
+        "\n",
+        "# Resize hotfix\n",
+        "# DAIN frames are a bit \"shifted / smaller\" compared to original input frames. This can partly be mitigated with resizing\n",
+        "# DAIN frames to the resolution +2px and cropping the result to the original resoultion with the starting point (1,1).\n",
+        "# Without this fix, DAIN tends to make \"vibrating\" output and it is pretty noticible with static elements like text.\n",
+        "# This hotfix tries to make such effects less visible for a smoother video playback. I do not know what DAINAPP\n",
+        "# uses as a fix for this problem, but the original does show such behaviour with the default test images. More advanced\n",
+        "# users can change the interpolation method. The methods cv2.INTER_CUBIC and cv2.INTER_LANCZOS4 are recommended.\n",
+        "# The current default value is cv2.INTER_LANCZOS4.\n",
+        "RESIZE_HOTFIX = True\n",
+        "\n",
+        "# Auto-delete output PNG dir after ffmpeg video creation.\n",
+        "AUTO_REMOVE = True"
+      ],
+      "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",
+        "# P100: 16GB (Works)\n",
+        "# T4: 16GB [RuntimeError: CUDA call failed]\n",
+        "# P4: 8GB (Works)\n",
+        "# K80: 8GB (Not tested)\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"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "UeaU8um5-2NS",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "from IPython.display import clear_output\n",
+        "!git clone https://github.com/styler00dollar/Colab-DAIN /content/DAIN\n",
+        "\n",
+        "# 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.\")\n",
+        "\n",
+        "# 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.\")\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\n",
+        "\n",
+        "!CUDA_VISIBLE_DEVICES=0\n",
+        "\n",
+        "!sudo apt-get install imagemagick imagemagick-doc"
+      ],
+      "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",
+        "\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",
+        "%shell ffmpeg -i '/content/DAIN/{filename}' '{FRAME_INPUT_DIR}/%05d.png'\n",
+        "\n",
+        "png_generated_count_command_result = %shell ls '{FRAME_INPUT_DIR}' | wc -l\n",
+        "clear_output()\n",
+        "\n",
+        "pngs_generated_count = int(png_generated_count_command_result.output.strip())\n",
+        "\n",
+        "import shutil\n",
+        "if SEAMLESS==True:\n",
+        "  pngs_generated_count += 1\n",
+        "  original = str(FRAME_INPUT_DIR)+\"/00001.png\"\n",
+        "  target = str(FRAME_INPUT_DIR)+\"/\"+str(pngs_generated_count).zfill(5)+\".png\"\n",
+        "  shutil.copyfile(original, target)\n",
+        "\n",
+        "print(f\"Input FPS: {fps}\")\n",
+        "print(f\"{pngs_generated_count} frame PNGs generated.\")\n",
+        "\n",
+        "# Checking if PNG 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 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 colab_interpolate.py --netName DAIN_slowmotion --time_step {fps/TARGET_FPS} --start_frame 1 --end_frame {pngs_generated_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",
+        "import numpy as np\n",
+        "%cd {FRAME_OUTPUT_DIR}\n",
+        "if(RESIZE_HOTFIX==True):\n",
+        "  images = []\n",
+        "  for filename in os.listdir(f'{FRAME_OUTPUT_DIR}'):\n",
+        "    img = cv2.imread(os.path.join(f'{FRAME_OUTPUT_DIR}',filename))\n",
+        "    part_filename = os.path.splitext(filename)\n",
+        "    if(part_filename[0].endswith('0')==False):\n",
+        "      dimension = (img.shape[1]+2, img.shape[0]+2)\n",
+        "      resized = cv2.resize(img, dimension, interpolation=cv2.INTER_LANCZOS4)\n",
+        "      crop = resized[1:(dimension[1]-1), 1:(dimension[0]-1)]\n",
+        "      cv2.imwrite(part_filename[0]+\".png\", crop)\n",
+        "\n",
+        "%cd /content/DAIN"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "TKREDli2IDMV",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Create 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",
+        "if(AUTO_REMOVE==True):\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",
+        "if(AUTO_REMOVE==True):\n",
+        "  !rm -rf {FRAME_OUTPUT_DIR}/*\n",
+        "  !rm -rf output-audio.aac"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "8tTXpWkRBm8T",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Manually removing output PNG dir (before working with another file)\n",
+        "!rm -rf {FRAME_OUTPUT_DIR}/*"
+      ],
+      "execution_count": 0,
+      "outputs": []
+    }
+  ]
+}

README.md

@@ -20,7 +20,8 @@ This work is developed based on our TPAMI work [MEMC-Net](https://github.com/bao
 1. [Testing Pre-trained Models](#testing-pre-trained-models)
 1. [Downloading Results](#downloading-results)
 1. [Slow-motion Generation](#slow-motion-generation)
-1. [Training New Models](#training-new-models) 
+1. [Training New Models](#training-new-models)
+1. [Google Colab Demo](#google-colab-demo)
 
 ### Introduction
 We propose the **D**epth-**A**ware video frame **IN**terpolation (**DAIN**) model to explicitly detect the occlusion by exploring the depth cue.
@@ -224,6 +225,14 @@ Then test the new model by executing:
 
     $ CUDA_VISIBLE_DEVICES=0 python demo_MiddleBury.py
 
+### Google Colab Demo
+This is a modification of DAIN that allows the usage of Google Colab and is able to do a full demo interpolation from a source video to a target video.
+
+Original Notebook File by btahir can be found [here](https://github.com/baowenbo/DAIN/issues/44).
+This is a modification by [Styler00Dollar](https://github.com/styler00dollar) and [Alpha](https://github.com/AlphaGit).
+
+Simply upload the `Colab_DAIN.ipynb` file to your Google Drive or use this [link](https://colab.research.google.com/drive/1gzsfDV_MIdehr7Y8ZzWjTuW-mMZRP4Vy).
+
 ### Contact
 [Wenbo Bao](mailto:bwb0813@gmail.com); [Wei-Sheng (Jason) Lai](mailto:phoenix104104@gmail.com)