Skip to content

Commit

Permalink
Merge pull request #56 from styler00dollar/master
Browse files Browse the repository at this point in the history 
Colab end-to-end runnable version
  • Loading branch information
@baowenbo
baowenbo committed Apr 18, 2020
2 parents 37281cc + 610d812 commit 05685d8
Show file tree
Hide file tree
Showing 4 changed files with 544 additions and 1 deletion.
366 changes: 366 additions & 0 deletions Colab_DAIN.ipynb
View file
Original file line number Diff line number Diff line change
@@ -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": []
}
]
}
11 changes: 10 additions & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit 05685d8

Please sign in to comment.