Super_SloMo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Super-SloMo.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPhd1ckgW0ufe7ww7UC6dh9",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MSFTserver/AI-Colab-Notebooks/blob/main/Super_SloMo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# [Super-SloMo](https://github.com/avinashpaliwal/Super-SloMo.git) [![MIT Licence](https://badges.frapsoft.com/os/mit/mit.svg?v=103)](https://opensource.org/licenses/mit-license.php)\n",
        "PyTorch implementation of \"Super SloMo: High Quality Estimation of Multiple Intermediate Frames for Video Interpolation\" by Jiang H., Sun D., Jampani V., Yang M., Learned-Miller E. and Kautz J. [[Project]](https://people.cs.umass.edu/~hzjiang/projects/superslomo/) [[Paper]](https://arxiv.org/abs/1712.00080)\n",
        "\n",
        "__Check out [our paper](https://github.com/avinashpaliwal/Deep-SloMo) \"Deep Slow Motion Video Reconstruction with Hybrid Imaging System\" published in TPAMI.__\n",
        "\n",
        "---\n",
        "**Colab by HostsServer aka [MSFTserver](https://github.com/MSFTserver)**"
      ],
      "metadata": {
        "id": "Wz4BaariVdh5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9CwLdiALQUfG"
      },
      "outputs": [],
      "source": [
        "#@markdown #  **1) Download Super-Slomo Repo & Model**\n",
        "!git clone -q --depth 1 https://github.com/avinashpaliwal/Super-SloMo.git\n",
        "from os.path import exists\n",
        "def download_file(file, file_name):\n",
        "  !curl -L {file} -o {file_name}\n",
        "  \n",
        "pretrained_model = 'SuperSloMo.ckpt'\n",
        "if not exists(pretrained_model):\n",
        "  download_file(\"https://www.dropbox.com/s/f2f5pi76z6aaehe/SuperSloMo-Python37%2B1.10.0%2Bcu111-epoch150.ckpt\", pretrained_model)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # **2) Run this block and Upload Video by clicking the Button that pops up below this codeblock! Wait till it loads the video and once it's done run the next block**\n",
        "#@markdown You can drag and drop into choose files button\n",
        "import os\n",
        "from google.colab import files\n",
        "import shutil\n",
        "\n",
        "upload_folder = 'upload'\n",
        "result_folder = 'results'\n",
        "temp_folder = 'tmp'\n",
        "\n",
        "if os.path.isdir(upload_folder):\n",
        "    shutil.rmtree(upload_folder)\n",
        "if os.path.isdir(result_folder):\n",
        "    shutil.rmtree(result_folder)\n",
        "if os.path.isdir(temp_folder):\n",
        "    shutil.rmtree(temp_folder)\n",
        "os.mkdir(upload_folder)\n",
        "os.mkdir(result_folder)\n",
        "\n",
        "# upload images\n",
        "uploaded = files.upload()\n",
        "for filename in uploaded.keys():\n",
        "  dst_path = os.path.join(upload_folder, filename)\n",
        "  print(f'move {filename} to {dst_path}')\n",
        "  shutil.move(filename, dst_path)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ScOvXncDQ6Os"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import subprocess in case this cell is run without the above cells\n",
        "import subprocess\n",
        "\n",
        "SLOW_MOTION_FACTOR = 2#@param {type:\"number\"}\n",
        "TARGET_FPS = 15#@param {type:\"number\"}\n",
        "os.mkdir(temp_folder)\n",
        "cmd1 = [\n",
        "    'python',\n",
        "    'Super-SloMo/video_to_slomo.py',\n",
        "    '--checkpoint',\n",
        "    pretrained_model,\n",
        "    '--video',\n",
        "    f'{os.getcwd()}/{upload_folder}/{filename}',\n",
        "    '--sf',\n",
        "    str(SLOW_MOTION_FACTOR),\n",
        "    '--fps',\n",
        "    str(TARGET_FPS),\n",
        "    '--output',\n",
        "    f'{os.getcwd()}/{temp_folder}/{os.path.splitext(filename)[0]}-slomo.mkv',\n",
        "]\n",
        "print(f\"running Super SlowMo on {filename}...\")\n",
        "process = subprocess.Popen(cmd1, cwd=f'/content', stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "process.stdout.readlines()\n",
        "stdout, stderr = process.communicate()\n",
        "if process.returncode != 0:\n",
        "    raise RuntimeError(stderr)\n",
        "\n",
        "cmd2 = [\n",
        "    'ffmpeg',\n",
        "    '-i',\n",
        "    f'{os.getcwd()}/{temp_folder}/{os.path.splitext(filename)[0]}-slomo.mkv',\n",
        "    '-pix_fmt',\n",
        "    'yuv420p',\n",
        "    '-preset',\n",
        "    'veryslow',\n",
        "    f'{os.getcwd()}/{result_folder}/{os.path.splitext(filename)[0]}-slomo.mp4',\n",
        "]\n",
        "print(f\"converting file back to .mp4...\")\n",
        "process = subprocess.Popen(cmd2, cwd=f'/content', stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "process.stdout.readlines()\n",
        "stdout, stderr = process.communicate()\n",
        "if process.returncode != 0:\n",
        "    raise RuntimeError(stderr)\n",
        "    print(stderr)\n",
        "print(\"All Finished!\")\n",
        "shutil.rmtree(temp_folder)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "w0vXIE9sQz44"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title **See video in the browser**\n",
        "# @markdown This process may take a little longer. If you don't want to wait, download it by executing the next cell instead of using this cell.\n",
        "from base64 import b64encode\n",
        "from IPython import display\n",
        "mp4 = open(f'{os.getcwd()}/{result_folder}/{os.path.splitext(filename)[0]}-slomo.mp4','rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "display.HTML(\"\"\"\n",
        "<video width=400 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "hUMBmaKlWeue"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title **Download video**\n",
        "from google.colab import files\n",
        "files.download(f'{os.getcwd()}/{result_folder}/{os.path.splitext(filename)[0]}-slomo.mp4')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "t6_xKJLLWtbM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}