Real_ESRGAN_Pytorch_Inference_NO_UI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "6BPxh_VmVVIu"
      ],
      "gpuType": "T4",
      "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/Nick088Official/Real-ESRGAN_Pytorch/blob/main/Real_ESRGAN_Pytorch_Inference_NO_UI.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IRDbDYYMQt_Y"
      },
      "source": [
        "# Real-ESRGAN Pytorch Inference NO UI\n",
        "\n",
        "[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2107.10833)\n",
        "[![GitHub Stars](https://img.shields.io/github/stars/Nick088Official/Real-ESRGAN_Pytorch?style=social)](https://github.com/Nick088Official/Real-ESRGAN_Pytorch)\n",
        "\n",
        "This is a **Practical Image Restoration Demo** of the Pytorch ai-forever custom pretrained Real-ESRGAN models.\n",
        "We extend the powerful ESRGAN to a practical restoration application (namely, Real-ESRGAN), which is trained with pure synthetic data. <br>\n",
        "The following figure shows some real-life examples.\n",
        "\n",
        "Low quality image:\n",
        "\n",
        "![](https://github.com/Nick088Official/Real-ESRGAN-Pytorch/blob/main/inputs/lr_lion.png?raw=1)\n",
        "\n",
        "Real-ESRGAN_Pytorch result:\n",
        "\n",
        "![](https://github.com/Nick088Official/Real-ESRGAN-Pytorch/blob/main/results/sr_lion.png?raw=1)\n",
        "\n",
        "**Note that Real-ESRGAN may still fail in some cases as the real-world degradations are really too complex.**<br>\n",
        "Moreover, it **may not** perform well on **human faces, you could try [GFPGAN](https://colab.research.google.com/github/Nick088Official/GFPGAN-Fix/blob/master/GFPGAN_fix_inference.ipynb) for that\n",
        "<br>\n",
        "\n",
        "This is the Pytorch Implementation of https://github.com/ai-forever/Real-ESRGAN\n",
        "\n",
        "**Credits:** [Nick088](https://linktr.ee/Nick088), Xinntao, Tencent, Geeve George, ai-forever"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GnpnrLfMV2jU",
        "cellView": "form"
      },
      "source": [
        "#@title Installation\n",
        "#@markdown Before running, make sure that you choose\n",
        "#@markdown * Runtime Type = Python 3\n",
        "#@markdown * Hardware Accelerator = GPU (Faster, free daily limit of gpu around 12 hours) or CPU (very slow)\n",
        "\n",
        "#@markdown in the **Runtime** menu -> **Change runtime type**\n",
        "\n",
        "#@markdown By running this, we clone the repository, set up the envrironment.\n",
        "\n",
        "\n",
        "from IPython.display import clear_output\n",
        "import torch\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "    device = \"cuda\"\n",
        "    print(\"Using GPU\")\n",
        "else:\n",
        "    device = \"cpu\"\n",
        "    print(\"Using CPU\")\n",
        "\n",
        "!git clone https://github.com/Nick088Official/Real-ESRGAN_Pytorch/\n",
        "%cd Real-ESRGAN_Pytorch/Scripts\n",
        "!pip install -r requirements_no_ui.txt\n",
        "from RealESRGAN import RealESRGAN\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import ffmpeg\n",
        "%cd ..\n",
        "clear_output()\n",
        "print(f'Installed on {\"GPU\" if torch.cuda.is_available() else \"CPU\"}!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Upload Input Video/Image\n",
        "import os\n",
        "from google.colab import files\n",
        "import shutil\n",
        "from io import BytesIO\n",
        "import io\n",
        "\n",
        "upload_folder = 'inputs'\n",
        "result_folder = 'results'\n",
        "\n",
        "# upload files\n",
        "uploaded = files.upload()\n",
        "for filename in uploaded.keys():\n",
        "    dst_path = os.path.join(upload_folder, filename)\n",
        "    shutil.move(filename, dst_path)\n",
        "    print(f'moved input {filename} to {dst_path}')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "J5FiFN3p9php"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Inference Input Image/Video\n",
        "\n",
        "model_scale = \"2\" #@param [\"2\", \"4\", \"8\"] {allow-input: false}\n",
        "\n",
        "infer_no_ui_command = f\"Scripts/Infer_NO_UI.py --file '{filename}' --size {model_scale}\"\n",
        "\n",
        "!python $infer_no_ui_command"
      ],
      "metadata": {
        "id": "TOrRoYCA9tUh",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7IMD5vhOYp68",
        "cellView": "form"
      },
      "source": [
        "#@title Visualize Input vs Output Image (Image Input ONLY)\n",
        "# utils for visualization\n",
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "def display(img1, img2):\n",
        "  fig = plt.figure(figsize=(25, 10))\n",
        "  ax1 = fig.add_subplot(1, 2, 1)\n",
        "  plt.title('Input image', fontsize=16)\n",
        "  ax1.axis('off')\n",
        "  ax2 = fig.add_subplot(1, 2, 2)\n",
        "  plt.title(f'Real-ESRGAN Pytorch x{model_scale} output', fontsize=16)\n",
        "  ax2.axis('off')\n",
        "  ax1.imshow(img1)\n",
        "  ax2.imshow(img2)\n",
        "def imread(img_path):\n",
        "  img = cv2.imread(img_path)\n",
        "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "  return img\n",
        "\n",
        "# display each image in the upload folder\n",
        "import os\n",
        "import glob\n",
        "\n",
        "input_folder = 'inputs'\n",
        "result_folder = 'results'\n",
        "input_list = sorted(glob.glob(os.path.join(input_folder, '*')))\n",
        "output_list = sorted(glob.glob(os.path.join(result_folder, '*')))\n",
        "for input_path, output_path in zip(input_list, output_list):\n",
        "  img_input = imread(input_path)\n",
        "  img_output = imread(output_path)\n",
        "  display(img_input, img_output)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lHNHoP8PZJQ7",
        "cellView": "form"
      },
      "source": [
        "#@title Download the results\n",
        "zip_filename = f'Real-ESRGAN-Pytorch-x{model_scale}_result.zip'\n",
        "if os.path.exists(zip_filename):\n",
        "  os.remove(zip_filename)\n",
        "os.system(f\"zip -r -j {zip_filename} results/*\")\n",
        "files.download(zip_filename)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Delete Inputs & Outputs\n",
        "#@markdown You need to run this to clear the previous inputs and ouputs after you finished inferencing the images you wanted and downloaded the results, only after this you can inference again.\n",
        "\n",
        "shutil.rmtree(upload_folder)\n",
        "shutil.rmtree(result_folder)\n",
        "os.makedirs(upload_folder, exist_ok=True)\n",
        "os.makedirs(result_folder, exist_ok=True)\n",
        "print(\"Deleted previous Inputs & Outputs Images, now you can inference again.\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "qOEvBZYACZIO"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}