ESRGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ESRGAN.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "ChddBZ3XQhtH",
        "cellView": "form"
      },
      "source": [
        "#@title ## <center>1. Download rich package & Mount Gdrive</center>\n",
        "\n",
        "gdrive_mode = 'MOUNT' #@param ['MOUNT', 'UNMOUNT']\n",
        "gdrive_type = 'SHARED DRIVE' #@param ['MY DRIVE', 'SHARED DRIVE']\n",
        "gdrive_shared_drive = 'my_shared_drive' #@param {type:\"string\"}\n",
        "gdrive_subpath = \"work\" #@param {type:\"string\"}\n",
        "\n",
        "if gdrive_type == \"MY DRIVE\":\n",
        "    gdrive_path = f'/content/drive/MyDrive'\n",
        "else:\n",
        "    gdrive_path = f'/content/drive/Shareddrives/{gdrive_shared_drive}/{gdrive_subpath}'\n",
        "\n",
        "!pip install rich rich[jupyter] > /dev/null\n",
        "\n",
        "from rich import print\n",
        "from rich.console import Console\n",
        "from pathlib import Path\n",
        "\n",
        "console = Console()\n",
        "print(f'Installed rich package')\n",
        "\n",
        "from google.colab import drive\n",
        "\n",
        "drive.mount._DEBUG = False\n",
        "if gdrive_mode == 'MOUNT':\n",
        "    drive.mount('/content/drive', force_remount=True)\n",
        "    # print(f'Google Drive mounted on \"{gdrive_path}\"')\n",
        "elif gdrive_mode == 'UNMOUNT':\n",
        "    try:\n",
        "        drive.flush_and_unmount()\n",
        "    except ValueError:\n",
        "        pass\n",
        "    get_ipython().system_raw('rm -rf /root/.config/Google/DriveFS')\n",
        "    # print(f'Google Drive unmounted')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5bnghRm9QuLm",
        "cellView": "form"
      },
      "source": [
        "#@title ##  <center>2. Download stuff</center>\n",
        "#@markdown ### Download/Update ESRGAN:\n",
        "# #@markdown ```Run this once. Running this while you already have a folder named Colab-ESRGAN in your Google Drive will give you an error.```\n",
        "#@markdown The first time you have to select `CLONE`, the other times select `PULL`.\n",
        "git_mode = 'PULL' #@param ['CLONE', 'PULL']\n",
        "esrgan_dirname = \"Colab-ESRGAN\" #@param {type:\"string\"}\n",
        "\n",
        "%cd \"{gdrive_path}\"\n",
        "esrgan_path = f'{gdrive_path}/{esrgan_dirname}'\n",
        "if git_mode == 'CLONE':\n",
        "    !git clone https://github.com/efonte/ESRGAN \"{esrgan_path}\"\n",
        "    %cd \"{esrgan_path}\"\n",
        "else:\n",
        "    %cd \"{esrgan_path}\"\n",
        "    !git pull\n",
        "with console.status(\"Installing packages...\"):\n",
        "    !pip install -q -r requirements.txt > /dev/null\n",
        "    !apt-get -qq update > /dev/null\n",
        "    !apt-get -qq install imagemagick > /dev/null\n",
        "    # !apt-get -qq install parallel > /dev/null\n",
        "    !pip install -U imageio \n",
        "    !pip install -U imageio-ffmpeg\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DhYF9Qi_QwYH",
        "cellView": "form"
      },
      "source": [
        "#@title ##  <center>Get available models</center>\n",
        "#@markdown #### Get the models located in the `{esrgan_path}/models` directory\n",
        "models_path= Path(f'{esrgan_path}/models')\n",
        "for model_path in models_path.glob(\"**/*.pth\"):\n",
        "    print(model_path.relative_to(models_path).as_posix())\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ja3SZv8EQxuH",
        "cellView": "form"
      },
      "source": [
        "#@title ## <center>Upscale - Image Mode</center>\n",
        "#@markdown #### ESRGAN image mode.\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ### Models:\n",
        "#@markdown ##### Examples:<br>• `1x_model1.pth>4x_model2.pth`<br>• `1x_model1.pth@50|1x_model2.pth@50>2x_model3.pth`<br>• `1x_model1.pth:50&1x_model2.pth:50>2x_model3.pth`\n",
        "models = \"1x_NoiseToner_Uniform_100k.pth>4x_DigitalFrames_2.1_Aggressive.pth\" #@param {type:\"string\"}\n",
        "#@markdown ---\n",
        "#@markdown ### Parameters:\n",
        "#@markdown ##### Input folder:\n",
        "input_path = \"input/test\" #@param {type:\"string\"}\n",
        "#@markdown ##### Output folder:\n",
        "output_path = \"output/test\" #@param {type:\"string\"}\n",
        "#@markdown ##### Reverse Order:\n",
        "reverse = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Skip existing output files:\n",
        "skip_existing = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Delete input files after upscaling:\n",
        "delete_input = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Use FloatingPoint16/Halftensor type:\n",
        "floating_point_16 = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Helps seamlessly upscale an image:<br>`None` = disable seamless.<br>`tile` = repeating along edges.<br>`mirror` = reflected along edges.<br>`replicate` = extended pixels along edges.<br>`alpha_pad` = extended alpha border.\n",
        "seamless = 'None' #@param ['None', 'tile', 'mirror', 'replicate', 'alpha_pad']\n",
        "#@markdown ##### Caches the maximum recursion depth used by the split/merge function. Useful only when upscaling images of the same size:\n",
        "cache_max_split_depth = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Whether to use a 1 bit alpha transparency channel, Useful for PSX upscaling:\n",
        "binary_alpha = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Whether to use a 2 bit alpha transparency channel, Useful for PSX upscaling \n",
        "ternary_alpha = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Only used when binary_alpha is supplied. Defines the alpha threshold for binary transparency:\n",
        "alpha_threshold = 0.5 #@param {type:\"number\"}\n",
        "#@markdown ##### Only used when binary_alpha is supplied. Determines the offset boundary from the alpha threshold for half transparency:\n",
        "alpha_boundary_offset = 0.2 #@param {type:\"number\"}\n",
        "#@markdown ##### Type of alpha processing to use:<br>`no_alpha` = no alpha processing.<br>`bas` = BA's difference method.<br>`alpha_separately` = upscaling the alpha channel separately (like IEU).<br>`swapping` = swapping an existing channel with the alpha channel.\n",
        "alpha_mode ='alpha_separately' #@param ['no_alpha', 'bas', 'alpha_separately', 'swapping']\n",
        "#@markdown ##### Use ImageMagick to save the upscaled image (higher quality but slower):\n",
        "imagemagick = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Convert the image to jpg:\n",
        "jpg = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Resize percent (`100` = no resize, `200` = x2, `50` = x0.5):\n",
        "resize = 100 #@param {type:\"number\"}\n",
        "#@markdown ##### Compress the output to zip file:\n",
        "zip = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Verbose mode:\n",
        "verbose = False #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "\n",
        "esrgan_args = f'\"{models}\" --input \"{input_path}\" --output \"{output_path}\"'\n",
        "if reverse:\n",
        "    esrgan_args += ' --reverse'\n",
        "if skip_existing:\n",
        "    esrgan_args += ' --skip-existing'\n",
        "if delete_input:\n",
        "    esrgan_args += ' --delete-input'\n",
        "if floating_point_16:\n",
        "    esrgan_args += ' --floating-point-16'\n",
        "if seamless != 'None':\n",
        "    esrgan_args += f' --seamless {seamless}'\n",
        "if cache_max_split_depth:\n",
        "    esrgan_args += ' --cache-max-split-depth'\n",
        "if binary_alpha:\n",
        "    esrgan_args += ' --binary-alpha'\n",
        "if ternary_alpha:\n",
        "    esrgan_args += ' --ternary-alpha'\n",
        "esrgan_args += f' --alpha-threshold {alpha_threshold} --alpha-boundary-offset {alpha_boundary_offset} --alpha-mode {alpha_mode}'\n",
        "if jpg:\n",
        "    esrgan_args += ' --jpg'\n",
        "esrgan_args += f' --resize {resize}'\n",
        "if imagemagick:\n",
        "    esrgan_args += ' --imagemagick'\n",
        "if zip:\n",
        "    esrgan_args += ' --zip'\n",
        "if verbose:\n",
        "    esrgan_args += ' -v' \n",
        "%cd {esrgan_path}\n",
        "# %run -i run.py image $esrgan_args > /dev/null\n",
        "%run -i run.py image $esrgan_args\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "nrzZENrvQz0m"
      },
      "source": [
        "#@title ## <center>ImageMagick - Utils</center>\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ##### Input folder:\n",
        "input_path = \"output/test\" #@param {type:\"string\"}\n",
        "#@markdown ##### Output folder (*must not exist*):\n",
        "output_path = \"output/test_resized\" #@param {type:\"string\"}\n",
        "#@markdown ##### Convert to jpg:\n",
        "jpg = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Resize percent (`100` = no resize, `200` = x2, `50` = x0.5):\n",
        "resize = 80 #@param {type:\"number\"}\n",
        "#@markdown ##### Show progress:\n",
        "monitor = True #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "\n",
        "# !rm -r \"{output_path}\"\n",
        "!mkdir \"{output_path}\"\n",
        "\n",
        "mogrify_args = f'-resize {resize}%'\n",
        "if monitor:\n",
        "    mogrify_args += ' -monitor' \n",
        "if jpg:\n",
        "    mogrify_args += ' -format jpg -sampling-factor 4:2:0 -strip -interlace JPEG -colorspace sRGB'\n",
        "else:\n",
        "    mogrify_args += ' -format png'\n",
        "\n",
        "!mogrify {mogrify_args} -path \"{output_path}\" \"{input_path}/*.png\"\n",
        "# !find \"{input_path}\" -type f -name '*.png' -print0 | xargs -0 -n3 mogrify {mogrify_args} -path \"{output_path}\"\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CPbVooh7Q1V1",
        "cellView": "form"
      },
      "source": [
        "#@title ## <center>Upscale - Video Mode</center>\n",
        "#@markdown #### ESRGAN video mode.\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ### Models:\n",
        "#@markdown ##### Examples:<br>• `1x_model1.pth>4x_model2.pth`<br>• `1x_model1.pth@50|1x_model2.pth@50>2x_model3.pth`<br>• `1x_model1.pth:50&1x_model2.pth:50>2x_model3.pth`\n",
        "models = \"1x_DeJpeg_Fatality_PlusULTRA_200k.pth\" #@param {type:\"string\"}\n",
        "#@markdown ---\n",
        "#@markdown ### Parameters:\n",
        "#@markdown ##### Input vide:\n",
        "input_path = \"input/video.mp4\" #@param {type:\"string\"}\n",
        "#@markdown ##### Output video:\n",
        "output_path = \"output/video.mp4\" #@param {type:\"string\"}\n",
        "#@markdown ##### Use FloatingPoint16/Halftensor type:\n",
        "floating_point_16 = True #@param {type:\"boolean\"}\n",
        "#@markdown ##### Helps seamlessly upscale an image:<br>`None` = disable seamless.<br>`tile` = repeating along edges.<br>`mirror` = reflected along edges.<br>`replicate` = extended pixels along edges.<br>`alpha_pad` = extended alpha border.\n",
        "seamless = 'None' #@param ['None', 'tile', 'mirror', 'replicate', 'alpha_pad']\n",
        "#@markdown ##### Caches the maximum recursion depth used by the split/merge function. Useful only when upscaling images of the same size:\n",
        "cache_max_split_depth = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Min SSIM value:\n",
        "min_ssim = 0.998 #@param {type:\"number\"}\n",
        "#@markdown ##### Enable diff mode (*beta*):\n",
        "diff = False #@param {type:\"boolean\"}\n",
        "#@markdown ##### Only used with diff mode. Chunk size to be able to generate the frame difference (*beta*):\n",
        "chunk_size = 2 #@param {type:\"number\"}\n",
        "#@markdown ##### Only used with diff mode. Padding size between each chunk (*beta*):\n",
        "padding_size = 2 #@param {type:\"number\"}\n",
        "#@markdown ##### Video quality:\n",
        "quality = 10 #@param {type:\"slider\", min:0, max:10, step:0.5}\n",
        "#@markdown ##### Verbose mode:\n",
        "verbose = False #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "\n",
        "esrgan_args = f'\"{models}\" --input \"{input_path}\" --output \"{output_path}\"'\n",
        "if floating_point_16:\n",
        "    esrgan_args += ' --floating-point-16'\n",
        "if seamless != 'None':\n",
        "    esrgan_args += f' --seamless {seamless}'\n",
        "if cache_max_split_depth:\n",
        "    esrgan_args += ' --cache-max-split-depth'\n",
        "esrgan_args += f' --chunk-size {chunk_size} --padding-size {padding_size} --min-ssim {min_ssim} --quality {quality}'\n",
        "if diff:\n",
        "    esrgan_args += ' --diff'\n",
        "if verbose:\n",
        "    esrgan_args += ' -v' \n",
        "%cd {esrgan_path}\n",
        "# %run -i run.py video $esrgan_args > /dev/null\n",
        "%run -i run.py video $esrgan_args\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GN2I2U9SvDCN"
      },
      "source": [
        "# concatenate scenes\n",
        "# %cd \"/content/drive/Shareddrives/my_shared_drive/work/Colab-ESRGAN/output/video\"\n",
        "# !ffmpeg -f concat -safe 0 -i ffmpeg_list.txt -c copy video.mp4"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}