ESRGAN_4_Video.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ESRGAN-4-Video.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "private_outputs": true,
      "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/ESRGAN_4_Video.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D6HDAdt3tE15"
      },
      "source": [
        "# **ESRGAN for increasing resolution ON VIDEOS!!!**\n",
        "by HostsServer aka [MSFTserver](https://github.com/MSFTserver)\n",
        "####**-SEEMS TO ONLY WORK CONSISTENTLY ON CHROME!!**\n",
        "#### **-DOUBLE CHECK YOUR RUNTIME TYPE IS SET TO GPU**\n",
        "example: an image with height 500px by width 500px scaled to an image with height 1762px by width 1762px\n",
        "\n",
        "**You might need to go to Runtime>Restart Runtime**\n",
        "\n",
        "If you open the folder icon (in the left menu) you will see all your files, the good thing about doing **Runtime>Restart Runtime** is that it keeps all the files you downloaded but Restarts the Memory on the Cloud Computer so you can put the ESRGAN model in memory, compared to **Runtime>Factory Reset Runtime** which Restarts the whole Cloud Computer to a default clear state.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GnpnrLfMV2jU",
        "cellView": "form"
      },
      "source": [
        "#@markdown # **1) Download required Libraries**\n",
        "# # Clone Real-ESRGAN and enter the Real-ESRGAN\n",
        "from google.colab import drive, files\n",
        "import os, shutil, subprocess\n",
        "drive_mounted = False\n",
        "temp_folder = 'tmp_frames'\n",
        "result_folder = 'results'\n",
        "!git clone https://github.com/xinntao/Real-ESRGAN.git\n",
        "%cd Real-ESRGAN\n",
        "# Set up the environment\n",
        "!pip install basicsr facexlib gfpgan\n",
        "!pip install -r requirements.txt\n",
        "!python setup.py develop\n",
        "# Download the pre-trained model\n",
        "!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models\n",
        "%cd .."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # **2.A) Optional:** Mount Google Drive\n",
        "#@markdown uploading videos is really slow it will be faster to connect your drive and spcify the filename or path the video you want to use here\n",
        "\n",
        "#@markdown Step 2.B can be Skipped if using google drive!\n",
        "if os.path.exists('/content/drive') is not True:\n",
        "  drive.mount('/content/drive')\n",
        "\n",
        "drive_mounted = True\n",
        "drive_path = 'test.mp4'#@param{type: \"string\"}\n",
        "file_name = drive_path.split('/')[-1]\n",
        "input_path = f'/content/drive/MyDrive/{drive_path}'\n",
        "print(\"video path set:\\n\",input_path)"
      ],
      "metadata": {
        "id": "S8epjA_OEgHn",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u1tiyMZJW5td",
        "cellView": "form"
      },
      "source": [
        "#@markdown # **2.B) Run this block and Upload Video using the DIALOG box!**\n",
        "#@markdown Wait till it loads the video and once it's done run the next block\n",
        "\n",
        "#@markdown You can drag and drop into choose files button\n",
        "#@markdown ## SKIP this if you mounted your Google Drive!\n",
        "if drive_mounted is not True:\n",
        "\n",
        "  upload_folder = 'upload'\n",
        "\n",
        "  if os.path.isdir(upload_folder):\n",
        "      shutil.rmtree(upload_folder)\n",
        "\n",
        "  os.mkdir(upload_folder)\n",
        "  \n",
        "  # upload images\n",
        "  uploaded = files.upload()\n",
        "  file_name = next(iter(uploaded))\n",
        "  input_path = f'/content/{upload_folder}/{file_name}'\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)\n",
        "else:\n",
        "  print(\"drive mounted nothing to do here move along\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # **3) Extract Frames from Video**\n",
        "\n",
        "if os.path.isdir(temp_folder):\n",
        "  shutil.rmtree(temp_folder)\n",
        "\n",
        "os.mkdir(temp_folder)\n",
        "print(f'Extracting Frames to: {temp_folder}')\n",
        "cmd = [\n",
        "       'ffmpeg',\n",
        "       '-i',\n",
        "       input_path,\n",
        "       '-qscale:v',\n",
        "       '1',\n",
        "       '-qmin',\n",
        "       '1',\n",
        "       '-qmax',\n",
        "       '1',\n",
        "       '-vsync',\n",
        "       '0',\n",
        "       f'{temp_folder}/frame_%08d.png'\n",
        "]\n",
        "process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "stdout, stderr = process.communicate()\n",
        "if process.returncode != 0:\n",
        "    print(stderr)\n",
        "    raise RuntimeError(stderr)\n",
        "else:\n",
        "    frame_count = len(os.listdir(f'/content/{temp_folder}'))\n",
        "    print(f\"Done, Extracted {frame_count} Frames\")\n"
      ],
      "metadata": {
        "id": "F5RCrfMkBEbQ",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XTVL4m_zXqBU",
        "cellView": "form"
      },
      "source": [
        "#@markdown # **4) Run ERSGAN on Extracted Frames!**\n",
        "%cd /content/Real-ESRGAN\n",
        "frame_count = len(os.listdir(f'/content/{temp_folder}'))\n",
        "print(f\"Enhancing {frame_count} Frames with ESRGAN...\")\n",
        "cmd = [\n",
        "       'python',\n",
        "       'inference_realesrgan.py',\n",
        "       '-n',\n",
        "       'RealESRGAN_x4plus',\n",
        "       '-i',\n",
        "       f'../{temp_folder}',\n",
        "       '--outscale',\n",
        "       '4',\n",
        "       '--half',\n",
        "       '--face_enhance'\n",
        "]\n",
        "#process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)\n",
        "# stdout, stderr = process.communicate()\n",
        "# if process.returncode != 0:\n",
        "#     print(stderr)\n",
        "#     raise RuntimeError(stderr)\n",
        "# else:\n",
        "#     print(\"Done Extracting Frames\")\n",
        "!{' '.join(cmd)}\n",
        "print(\"Done Enhancing Frames\")\n",
        "%cd .."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # **5) Recreate Video From Enhanced Frames!**\n",
        "import subprocess\n",
        "frame_count = len(os.listdir(f'/content/{temp_folder}'))\n",
        "if os.path.isdir(result_folder):\n",
        "  shutil.rmtree(result_folder)\n",
        "os.mkdir(result_folder)\n",
        "\n",
        "fps = 15 #@param {type: \"number\"}\n",
        "print(f\"Recompiling {frame_count} Frames into Video...\")\n",
        "cmd = [\n",
        "       'ffmpeg',\n",
        "       '-i',\n",
        "       f'/content/Real-ESRGAN/results/frame_%08d_out.png',\n",
        "       '-c:a',\n",
        "       'copy',\n",
        "       '-c:v',\n",
        "       'libx264',\n",
        "       '-r',\n",
        "       str(fps),\n",
        "       '-pix_fmt',\n",
        "       'yuv420p',\n",
        "       f'{result_folder}/enhanced_{file_name}'\n",
        "]\n",
        "process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "stdout, stderr = process.communicate()\n",
        "if process.returncode != 0:\n",
        "    print(stderr)\n",
        "    raise RuntimeError(stderr)\n",
        "else:\n",
        "    print(\"Done Recreating Video\")"
      ],
      "metadata": {
        "id": "Q9V7MOOgBTY7",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lHNHoP8PZJQ7",
        "cellView": "form"
      },
      "source": [
        "#@markdown # **6) Download Super Enhanced Video**\n",
        "#@markdown if drive is mounted this will save the file to the root of your google drive\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "filename= os.listdir(result_folder)\n",
        "if drive_mounted:\n",
        "  if os.path.exists('/content/drive/MyDrive/Hosts-AI-Edits/ESRGAN-4-Vid') is not True:\n",
        "    os.makedirs('/content/drive/MyDrive/Hosts-AI-Edits/ESRGAN-4-Vid')\n",
        "  shutil.move(f'{result_folder}/{filename[0]}',f'/content/drive/MyDrive/Hosts-AI-Edits/ESRGAN-4-Vid/{filename[0]}')\n",
        "  print(f\"moved {filename[0]} to the folder Hosts-AI-Edits/ESRGAN-4-Vid on your Google Drive\")\n",
        "else:\n",
        "  files.download(f'{result_folder}/{filename[0]}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # **Optional**: Unmount Google Drive\n",
        "#@markdown might be useful if you decide you want to upload on here instead\n",
        "drive_mounted = False\n",
        "drive.flush_and_unmount()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "kshS1oE3lLls"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}