File: //opt/DAIN/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": []
}
]
}