4k_Video_Upscaler_Colab_(Real_ESRGAN).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNo7Sv82VzH+t3sj4q6tf29",
      "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/yuvraj108c/4k-video-upscaler-colab/blob/main/4k_Video_Upscaler_Colab_(Real_ESRGAN).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 4k Video Upscaler Colab (Real-ESRGAN)\n",
        "\n",
        "Adapted from: [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN)\n",
        "\n",
        "Made with ❤️ by: [yuvraj108c](https://github.com/yuvraj108c)\n",
        "\n",
        "Github repository: https://github.com/yuvraj108c/4k-video-upscaler-colab"
      ],
      "metadata": {
        "id": "zEksijsOSulF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. Setup (~1 minute)"
      ],
      "metadata": {
        "id": "Jt28DQ93QK5Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "assert torch.cuda.is_available(), \"GPU not detected.. Please change runtime to GPU\"\n",
        "\n",
        "from PIL import Image\n",
        "import cv2, os, subprocess\n",
        "from google.colab import drive\n",
        "from tqdm import tqdm\n",
        "\n",
        "!git clone https://github.com/xinntao/Real-ESRGAN.git\n",
        "%cd Real-ESRGAN\n",
        "\n",
        "!pip install -q torch==2.0.1 torchvision==0.15.2 --extra-index-url https://download.pytorch.org/whl/cu118\n",
        "!pip install -q basicsr facexlib gfpgan ffmpeg ffmpeg-python\n",
        "!pip install -q -r requirements.txt\n",
        "!python setup.py develop\n",
        "\n",
        "mount_drive=False #@param{type:\"boolean\"}\n",
        "\n",
        "if mount_drive:\n",
        "  drive.mount('/content/gdrive/')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "kzcD1CMyQtbv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. Upscale video\n",
        "\n",
        "- The upscaled video will be saved to `output_dir`\n",
        "- If google drive is mounted, it will be also saved at `MyDrive/Upscaled Videos (REAL-ESRGAN)`\n"
      ],
      "metadata": {
        "id": "mZZPG007SB4A"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NX2v5PH9O7Sw"
      },
      "outputs": [],
      "source": [
        "video_path=\"/content/input.mp4\" #@param{type:\"string\"}\n",
        "output_dir=\"/content/\" #@param{type:\"string\"}\n",
        "resolution = \"FHD (1920 x 1080)\" # @param [\"FHD (1920 x 1080)\", \"2k (2560 x 1440)\", \"4k (3840 x 2160)\",\"2 x original\", \"3 x original\", \"4 x original\"] {type:\"string\"}\n",
        "model = \"RealESRGAN_x4plus\" #@param [\"RealESRGAN_x4plus\" , \"RealESRGAN_x4plus_anime_6B\", \"realesr-animevideov3\"]\n",
        "\n",
        "assert os.path.exists(video_path), \"Video file does not exist\"\n",
        "\n",
        "video_capture = cv2.VideoCapture(video_path)\n",
        "video_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "video_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "\n",
        "final_width = None\n",
        "final_height = None\n",
        "aspect_ratio = float(video_width/video_height)\n",
        "\n",
        "# Get output resolutions\n",
        "match resolution:\n",
        "  case \"FHD (1920 x 1080)\":\n",
        "    final_width=1920\n",
        "    final_height=1080\n",
        "  case \"2k (2560 x 1440)\":\n",
        "    final_width=2560\n",
        "    final_height=1440\n",
        "  case \"4k (3840 x 2160)\":\n",
        "    final_width=3840\n",
        "    final_height=2160\n",
        "  case \"2 x original\":\n",
        "    final_width=2*video_width\n",
        "    final_height=2*video_height\n",
        "  case \"3 x original\":\n",
        "    final_width=3*video_width\n",
        "    final_height=3*video_height\n",
        "  case \"4 x original\":\n",
        "    final_width=4*video_width\n",
        "    final_height=4*video_height\n",
        "\n",
        "if aspect_ratio == 1.0 and \"original\" not in resolution:\n",
        "  final_height = final_width\n",
        "\n",
        "if aspect_ratio < 1.0 and \"original\" not in resolution:\n",
        "  temp = final_width\n",
        "  final_width = final_height\n",
        "  final_height = temp\n",
        "\n",
        "scale_factor = max(final_width/video_width, final_height/video_height)\n",
        "isEven = int(video_width * scale_factor) % 2 == 0 and int(video_height * scale_factor) % 2 == 0\n",
        "\n",
        "# scale_factor needs to be even\n",
        "while isEven == False:\n",
        "  scale_factor += 0.01\n",
        "  isEven = int(video_width * scale_factor) % 2 == 0 and int(video_height * scale_factor) % 2 == 0\n",
        "\n",
        "print(f\"Upscaling from {video_width}x{video_height} to {final_width}x{final_height}, scale_factor={scale_factor}\")\n",
        "\n",
        "!python inference_realesrgan_video.py -n {model} -i '{video_path}' -o '{output_dir}' --outscale {scale_factor}\n",
        "\n",
        "video_name_with_ext = os.path.basename(video_path)\n",
        "video_name = video_name_with_ext.replace(\".mp4\",\"\")\n",
        "upscaled_video_path = f\"{output_dir}{video_name}_out.mp4\"\n",
        "final_video_name = f\"{video_name}_upscaled_{final_width}_{final_height}.mp4\"\n",
        "final_video_path = os.path.join(output_dir, final_video_name)\n",
        "\n",
        "# crop to fit\n",
        "if \"original\" not in resolution:\n",
        "  print(\"Cropping to fit...\")\n",
        "  command = f\"ffmpeg -loglevel error -hwaccel cuda -y -i '{upscaled_video_path}' -c:v h264_nvenc -filter:v  'crop={final_width}:{final_height}:(in_w-{final_width})/2:(in_h-{final_height})/2' -c:v libx264 -pix_fmt yuv420p '{final_video_path}'\"\n",
        "  subprocess.run(command,shell=True)\n",
        "\n",
        "print(f\"Upscaled video saved to: {final_video_path}\")\n",
        "\n",
        "# save to drive\n",
        "if mount_drive:\n",
        "  drive_folder = \"MyDrive/Upscaled Videos (REAL-ESRGAN)\"\n",
        "  save_directory_drive = f\"/content/gdrive/{drive_folder}\"\n",
        "  os.makedirs(save_directory_drive, exist_ok=True)\n",
        "\n",
        "  command = f\"'{final_video_path}' '{save_directory_drive}'/'{final_video_name}'\"\n",
        "  subprocess.run(command,shell=True)\n",
        "  print(f\"Saved to drive: /{drive_folder}/{final_video_name}\" )\n",
        "\n",
        "!rm \"{upscaled_video_path}\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3. Disconnect runtime"
      ],
      "metadata": {
        "id": "XydAN-xZT0AS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import runtime\n",
        "\n",
        "disconnect_when_finish = False  #@param{type:\"boolean\"}\n",
        "\n",
        "if disconnect_when_finish:\n",
        "  runtime.unassign()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "rIfUT_e9TqLp"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}