Image_Enhancement_with_Real_ESRGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "I-nRIXJdJNFH",
        "GM2B1l3SKaf3",
        "99LlExMcKi1Q"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **Test it Yourself!!!** 😉"
      ],
      "metadata": {
        "id": "Jea9yNNrAiEF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Getting Started**\n",
        "\n",
        "- Click on profile picture on top right. Login to Google Colab with your gmail account.\n",
        "\n",
        "- Click *File > Save copy in Drive*. This will open a new copy of this notebook in a new tab.\n",
        "\n",
        "- In the notebook copy opened up in the new tab, click on *Connect* button on top right. This will connect you to a cloud runtime.\n",
        "\n",
        "- Select *T4 GPU* option under *Runtime > Change runtime type*. Or choose any other GPU as you prefer if you have Google Colab Pro.\n"
      ],
      "metadata": {
        "id": "Tztw9_ZLY4yE"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download either *RealESRGAN_x2plus.pth* or *RealESRGAN_x4plus.pth* (or both of them) from [this model zoo]( https://github.com/NightmareAI/Real-ESRGAN/blob/master/docs/model_zoo.md).\n",
        "\n",
        "If you would like to try out the face enhancement model too, download *GFPGANv1.4.pth* from [here](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth).\n",
        "\n",
        "Upload the ***.pth PyTorch model files*** to this cloud runtime. You can click on the Folder icon on the left side panel, and then click on the Upload icon on the top of that left panel to upload your files."
      ],
      "metadata": {
        "id": "RHK1KjIFbQsn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The gray box with code below is called a 'code cell'. You can run the code in the code cell by clicking the play button on the left of the code cell. Run the pip install code cell below, it should download all the necessary packages to your cloud runtime."
      ],
      "metadata": {
        "id": "cJXd2lnlE0x4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torch torchvision torchaudio basicsr gfpgan --quiet"
      ],
      "metadata": {
        "id": "mIznOMeCrCSq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run this next cell to import all necessary packages:"
      ],
      "metadata": {
        "id": "kV4XpOJcJqmQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import math\n",
        "import numpy as np\n",
        "import os\n",
        "import queue\n",
        "import threading\n",
        "import torch\n",
        "\n",
        "from torch.nn import functional as F\n",
        "from torch import nn as nn\n",
        "from PIL import Image\n",
        "from basicsr.archs.rrdbnet_arch import RRDBNet\n",
        "from math import ceil, floor, sqrt\n",
        "from PIL import Image, ImageFilter\n",
        "from IPython.display import Image as display_image"
      ],
      "metadata": {
        "id": "LEAQj5yghv6H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Codes [**Click the play button below to run all the hidden cells. No need to open this tab unless you want to inspect all the long long codes...**]"
      ],
      "metadata": {
        "id": "I-nRIXJdJNFH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Utilities"
      ],
      "metadata": {
        "id": "GM2B1l3SKaf3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "These are the Utils functions and the classes containing the actual model architecture:"
      ],
      "metadata": {
        "id": "BWKDOTk2dHeV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Utils\n",
        "# ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\n",
        "\n",
        "def convert_to_jpg(input_path, output_path):\n",
        "    # Open the PNG file\n",
        "    with Image.open(input_path) as img:\n",
        "        # Save the image in JPEG format\n",
        "        img.convert(\"RGB\").save(output_path, 'JPEG')\n",
        "\n",
        "\n",
        "class RealESRGANer():\n",
        "    \"\"\"A helper class for upsampling images with RealESRGAN.\n",
        "\n",
        "    Args:\n",
        "        scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.\n",
        "        model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).\n",
        "        model (nn.Module): The defined network. Default: None.\n",
        "        tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop\n",
        "            input images into tiles, and then process each of them. Finally, they will be merged into one image.\n",
        "            0 denotes for do not use tile. Default: 0.\n",
        "        tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.\n",
        "        pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.\n",
        "        half (float): Whether to use half precision during inference. Default: False.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 scale,\n",
        "                 model_path,\n",
        "                 model=None,\n",
        "                 tile=0,\n",
        "                 tile_pad=10,\n",
        "                 pre_pad=10,\n",
        "                 half=False,\n",
        "                 device=None,\n",
        "                 gpu_id=None):\n",
        "        self.scale = scale\n",
        "        self.tile_size = tile\n",
        "        self.tile_pad = tile_pad\n",
        "        self.pre_pad = pre_pad\n",
        "        self.mod_scale = None\n",
        "        self.half = half\n",
        "\n",
        "        # initialize model\n",
        "        if gpu_id:\n",
        "            self.device = torch.device(\n",
        "                f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device\n",
        "        else:\n",
        "            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device\n",
        "        # if the model_path starts with https, it will first download models to the folder: realesrgan/weights\n",
        "        # if model_path.startswith('https://'):\n",
        "        #     model_path = load_file_from_url(\n",
        "        #         url=model_path, model_dir='realesrgan/weights')\n",
        "        loadnet = torch.load(model_path, map_location=torch.device('cpu'))\n",
        "        # prefer to use params_ema\n",
        "        if 'params_ema' in loadnet:\n",
        "            keyname = 'params_ema'\n",
        "        else:\n",
        "            keyname = 'params'\n",
        "        model.load_state_dict(loadnet[keyname], strict=True)\n",
        "        model.eval()\n",
        "        self.model = model.to(self.device)\n",
        "        if self.half:\n",
        "            self.model = self.model.half()\n",
        "\n",
        "    def pre_process(self, img):\n",
        "        \"\"\"Pre-process, such as pre-pad and mod pad, so that the images can be divisible\n",
        "        \"\"\"\n",
        "        img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()\n",
        "        self.img = img.unsqueeze(0).to(self.device)\n",
        "        if self.half:\n",
        "            self.img = self.img.half()\n",
        "\n",
        "        # pre_pad\n",
        "        if self.pre_pad != 0:\n",
        "            self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')\n",
        "        # mod pad for divisible borders\n",
        "        if self.scale == 2:\n",
        "            self.mod_scale = 2\n",
        "        elif self.scale == 1:\n",
        "            self.mod_scale = 4\n",
        "        if self.mod_scale is not None:\n",
        "            self.mod_pad_h, self.mod_pad_w = 0, 0\n",
        "            _, _, h, w = self.img.size()\n",
        "            if (h % self.mod_scale != 0):\n",
        "                self.mod_pad_h = (self.mod_scale - h % self.mod_scale)\n",
        "            if (w % self.mod_scale != 0):\n",
        "                self.mod_pad_w = (self.mod_scale - w % self.mod_scale)\n",
        "            self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')\n",
        "\n",
        "    def process(self):\n",
        "        # model inference\n",
        "        self.output = self.model(self.img)\n",
        "\n",
        "    def tile_process(self):\n",
        "        \"\"\"It will first crop input images to tiles, and then process each tile.\n",
        "        Finally, all the processed tiles are merged into one images.\n",
        "\n",
        "        Modified from: https://github.com/ata4/esrgan-launcher\n",
        "        \"\"\"\n",
        "        batch, channel, height, width = self.img.shape\n",
        "        output_height = height * self.scale\n",
        "        output_width = width * self.scale\n",
        "        output_shape = (batch, channel, output_height, output_width)\n",
        "\n",
        "        # start with black image\n",
        "        self.output = self.img.new_zeros(output_shape)\n",
        "        tiles_x = math.ceil(width / self.tile_size)\n",
        "        tiles_y = math.ceil(height / self.tile_size)\n",
        "\n",
        "        # loop over all tiles\n",
        "        for y in range(tiles_y):\n",
        "            for x in range(tiles_x):\n",
        "                # extract tile from input image\n",
        "                ofs_x = x * self.tile_size\n",
        "                ofs_y = y * self.tile_size\n",
        "                # input tile area on total image\n",
        "                input_start_x = ofs_x\n",
        "                input_end_x = min(ofs_x + self.tile_size, width)\n",
        "                input_start_y = ofs_y\n",
        "                input_end_y = min(ofs_y + self.tile_size, height)\n",
        "\n",
        "                # input tile area on total image with padding\n",
        "                input_start_x_pad = max(input_start_x - self.tile_pad, 0)\n",
        "                input_end_x_pad = min(input_end_x + self.tile_pad, width)\n",
        "                input_start_y_pad = max(input_start_y - self.tile_pad, 0)\n",
        "                input_end_y_pad = min(input_end_y + self.tile_pad, height)\n",
        "\n",
        "                # input tile dimensions\n",
        "                input_tile_width = input_end_x - input_start_x\n",
        "                input_tile_height = input_end_y - input_start_y\n",
        "                tile_idx = y * tiles_x + x + 1\n",
        "                input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]\n",
        "\n",
        "                # upscale tile\n",
        "                try:\n",
        "                    with torch.no_grad():\n",
        "                        output_tile = self.model(input_tile)\n",
        "                except RuntimeError as error:\n",
        "                    print('Error', error)\n",
        "                print(f'\\tTile {tile_idx}/{tiles_x * tiles_y}')\n",
        "\n",
        "                # output tile area on total image\n",
        "                output_start_x = input_start_x * self.scale\n",
        "                output_end_x = input_end_x * self.scale\n",
        "                output_start_y = input_start_y * self.scale\n",
        "                output_end_y = input_end_y * self.scale\n",
        "\n",
        "                # output tile area without padding\n",
        "                output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale\n",
        "                output_end_x_tile = output_start_x_tile + input_tile_width * self.scale\n",
        "                output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale\n",
        "                output_end_y_tile = output_start_y_tile + input_tile_height * self.scale\n",
        "\n",
        "                # put tile into output image\n",
        "                self.output[:, :, output_start_y:output_end_y,\n",
        "                            output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,\n",
        "                                                                       output_start_x_tile:output_end_x_tile]\n",
        "\n",
        "    def post_process(self):\n",
        "        # remove extra pad\n",
        "        if self.mod_scale is not None:\n",
        "            _, _, h, w = self.output.size()\n",
        "            self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]\n",
        "        # remove prepad\n",
        "        if self.pre_pad != 0:\n",
        "            _, _, h, w = self.output.size()\n",
        "            self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]\n",
        "        return self.output\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):\n",
        "        h_input, w_input = img.shape[0:2]\n",
        "        # img: numpy\n",
        "        img = img.astype(np.float32)\n",
        "        if np.max(img) > 256:  # 16-bit image\n",
        "            max_range = 65535\n",
        "            print('\\tInput is a 16-bit image')\n",
        "        else:\n",
        "            max_range = 255\n",
        "        img = img / max_range\n",
        "        if len(img.shape) == 2:  # gray image\n",
        "            img_mode = 'L'\n",
        "            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n",
        "        elif img.shape[2] == 4:  # RGBA image with alpha channel\n",
        "            img_mode = 'RGBA'\n",
        "            alpha = img[:, :, 3]\n",
        "            img = img[:, :, 0:3]\n",
        "            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "            if alpha_upsampler == 'realesrgan':\n",
        "                alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)\n",
        "        else:\n",
        "            img_mode = 'RGB'\n",
        "            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "        # ------------------- process image (without the alpha channel) ------------------- #\n",
        "        self.pre_process(img)\n",
        "        if self.tile_size > 0:\n",
        "            self.tile_process()\n",
        "        else:\n",
        "            self.process()\n",
        "        output_img = self.post_process()\n",
        "        output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()\n",
        "        output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))\n",
        "        if img_mode == 'L':\n",
        "            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)\n",
        "\n",
        "        # ------------------- process the alpha channel if necessary ------------------- #\n",
        "        if img_mode == 'RGBA':\n",
        "            if alpha_upsampler == 'realesrgan':\n",
        "                self.pre_process(alpha)\n",
        "                if self.tile_size > 0:\n",
        "                    self.tile_process()\n",
        "                else:\n",
        "                    self.process()\n",
        "                output_alpha = self.post_process()\n",
        "                output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()\n",
        "                output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))\n",
        "                output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)\n",
        "            else:  # use the cv2 resize for alpha channel\n",
        "                h, w = alpha.shape[0:2]\n",
        "                output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)\n",
        "\n",
        "            # merge the alpha channel\n",
        "            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)\n",
        "            output_img[:, :, 3] = output_alpha\n",
        "\n",
        "        # ------------------------------ return ------------------------------ #\n",
        "        if max_range == 65535:  # 16-bit image\n",
        "            output = (output_img * 65535.0).round().astype(np.uint16)\n",
        "        else:\n",
        "            output = (output_img * 255.0).round().astype(np.uint8)\n",
        "\n",
        "        if outscale is not None and outscale != float(self.scale):\n",
        "            output = cv2.resize(\n",
        "                output, (\n",
        "                    int(w_input * outscale),\n",
        "                    int(h_input * outscale),\n",
        "                ), interpolation=cv2.INTER_LANCZOS4)\n",
        "\n",
        "        return output, img_mode\n",
        "\n",
        "\n",
        "class PrefetchReader(threading.Thread):\n",
        "    \"\"\"Prefetch images.\n",
        "\n",
        "    Args:\n",
        "        img_list (list[str]): A image list of image paths to be read.\n",
        "        num_prefetch_queue (int): Number of prefetch queue.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, img_list, num_prefetch_queue):\n",
        "        super().__init__()\n",
        "        self.que = queue.Queue(num_prefetch_queue)\n",
        "        self.img_list = img_list\n",
        "\n",
        "    def run(self):\n",
        "        for img_path in self.img_list:\n",
        "            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)\n",
        "            self.que.put(img)\n",
        "\n",
        "        self.que.put(None)\n",
        "\n",
        "    def __next__(self):\n",
        "        next_item = self.que.get()\n",
        "        if next_item is None:\n",
        "            raise StopIteration\n",
        "        return next_item\n",
        "\n",
        "    def __iter__(self):\n",
        "        return self\n",
        "\n",
        "\n",
        "class IOConsumer(threading.Thread):\n",
        "\n",
        "    def __init__(self, opt, que, qid):\n",
        "        super().__init__()\n",
        "        self._queue = que\n",
        "        self.qid = qid\n",
        "        self.opt = opt\n",
        "\n",
        "    def run(self):\n",
        "        while True:\n",
        "            msg = self._queue.get()\n",
        "            if isinstance(msg, str) and msg == 'quit':\n",
        "                break\n",
        "\n",
        "            output = msg['output']\n",
        "            save_path = msg['save_path']\n",
        "            cv2.imwrite(save_path, output)\n",
        "        print(f'IO worker {self.qid} is done.')"
      ],
      "metadata": {
        "id": "0hvEpPJLJQKp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Main Code"
      ],
      "metadata": {
        "id": "0OFGjn1AKfYj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "And this is the main method to use to enhance our images:"
      ],
      "metadata": {
        "id": "gZMc6F4AM-R6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def enhance_image(input_file, layers=2, upscale=2, final_filename=\"\", enhance_faces=False):\n",
        "  ## Set models to use\n",
        "  if layers == 4:\n",
        "    # 4 layers\n",
        "    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)\n",
        "    netscale = 4\n",
        "    model_file = 'RealESRGAN_x4plus.pth'\n",
        "  elif layers == 2:\n",
        "    # 2 layers\n",
        "    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)\n",
        "    netscale = 2\n",
        "    model_file = 'RealESRGAN_x2plus.pth'\n",
        "  else:\n",
        "    print(\"Layers parameter must be either 2 or 4.\")\n",
        "    return\n",
        "\n",
        "  # Final enhanced image will be upscaled by this factor using LANCZOS4 resampling\n",
        "\n",
        "  # Input image\n",
        "  imgname, org_extension = input_file.split('.')\n",
        "  image = cv2.imread(input_file)\n",
        "  org_width, org_height = image.shape[:2]\n",
        "\n",
        "  # Convert image to JPG if need be\n",
        "  if org_extension not in [\"jpeg\", \"jpg\"]:\n",
        "      \"\"\"JPG file format reduces the file size and makes it feasable for\n",
        "      faster enhancement using the model.\n",
        "      \"\"\"\n",
        "      convert_to_jpg(input_file, f\"{imgname}.jpg\")\n",
        "      input_file = f\"{imgname}.jpg\"\n",
        "\n",
        "  # Compute tile size\n",
        "  if min(org_width, org_height) <= 800:\n",
        "    tile_size = 0\n",
        "    print(f\"Small image so batching is not necessary.\")\n",
        "  else:\n",
        "    tile_size = ceil(sqrt(min(org_width, org_height))) * 10\n",
        "  if tile_size > 500:\n",
        "    tile_size = 350\n",
        "  print(f\"Tile size being used: {tile_size}\")\n",
        "\n",
        "  # restorer\n",
        "  upsampler = RealESRGANer(\n",
        "      scale=netscale,\n",
        "      model_path=model_file,\n",
        "      model=model,\n",
        "      tile=tile_size,\n",
        "      tile_pad=2,\n",
        "      half=False)\n",
        "\n",
        "  # Use GFPGAN for face enhancement\n",
        "  if enhance_faces:\n",
        "    from gfpgan import GFPGANer\n",
        "    face_enhancer = GFPGANer(\n",
        "        model_path='GFPGANv1.4.pth',\n",
        "        upscale=upscale,\n",
        "        arch='clean',\n",
        "        channel_multiplier=2,\n",
        "        bg_upsampler=upsampler)\n",
        "\n",
        "  img = cv2.imread(input_file, cv2.IMREAD_UNCHANGED)\n",
        "\n",
        "  try:\n",
        "    if enhance_faces:\n",
        "      _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)\n",
        "    else:\n",
        "      output, _ = upsampler.enhance(img, outscale=upscale)\n",
        "  except RuntimeError as error:\n",
        "      print('Error', error)\n",
        "      print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')\n",
        "      print('Else, the file you are using may be too large.')\n",
        "  else:\n",
        "    if final_filename != \"\":\n",
        "      if not (final_filename.endswith(\".jpg\") or final_filename.endswith(\".jpeg\")):\n",
        "        print(\n",
        "          \"Your preferred final filename for the output image does not or has a wrong have a file extenstion.\"\\\n",
        "          \"Append .jpg or .jpg to your preferred filename.\"\n",
        "        )\n",
        "        return\n",
        "      save_path = final_filename\n",
        "    else:\n",
        "      save_path = f'{imgname}_out.jpg'\n",
        "\n",
        "    cv2.imwrite(save_path, output)\n",
        "    print(f\"Enhanced image has been saved to {save_path}.\\nClick refresh button on the left panel to get latest version of {save_path}\")\n",
        "    return save_path"
      ],
      "metadata": {
        "id": "jAALQsiSJEvz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Metrics Code"
      ],
      "metadata": {
        "id": "99LlExMcKi1Q"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "These are some helper fucntions to compute the quality of images based on certain metrics like resolution, sharpness, contrast, and noise:"
      ],
      "metadata": {
        "id": "YXCf3jRTTkUx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Metrics\n",
        "\n",
        "def get_resolution(image):\n",
        "    return image.shape[:2]\n",
        "\n",
        "def get_noise_level(image):\n",
        "    # Convert the image to grayscale\n",
        "    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
        "\n",
        "    # Compute the Discrete Fourier Transform (DFT)\n",
        "    f_transform = np.fft.fft2(gray)\n",
        "    f_transform_shifted = np.fft.fftshift(f_transform)\n",
        "\n",
        "    # Compute the magnitude spectrum\n",
        "    magnitude_spectrum = np.abs(f_transform_shifted)\n",
        "\n",
        "    # Calculate the noise level using the standard deviation of the magnitude spectrum\n",
        "    noise_level = np.std(np.log1p(magnitude_spectrum))\n",
        "\n",
        "    return round(noise_level, 2)\n",
        "\n",
        "def get_sharpness(image):\n",
        "    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
        "    # Apply an edge-enhancing filter (Laplacian) and compute variance as a measure of sharpness\n",
        "    laplacian = cv2.Laplacian(cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2GRAY), cv2.CV_64F)\n",
        "    return round(laplacian.var(), 2)\n",
        "\n",
        "def get_contrast(image):\n",
        "    # Using Michelson contrast measure\n",
        "    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
        "    I_max = np.max(gray)\n",
        "    I_min = np.min(gray)\n",
        "    print(I_min, I_max)\n",
        "\n",
        "    contrast = (I_max - I_min) / (I_max + I_min)\n",
        "    return round(contrast, 5)\n",
        "\n",
        "\n",
        "def get_filesize(image_file):\n",
        "  file_size = os.path.getsize(image_file)\n",
        "  return round(file_size / 1_000_000, 2)\n",
        "\n",
        "\n",
        "import time\n",
        "class Timer:\n",
        "    def __init__(self) -> None:\n",
        "        self.start = 0\n",
        "        self.end = 0\n",
        "\n",
        "    def start(self):\n",
        "        self.start = time.time()\n",
        "\n",
        "    def end(self):\n",
        "        self.end = time.time()\n",
        "        elapsed_time = self.end - self.start\n",
        "        print(f\"Elapsed Time: {elapsed_time} seconds\")\n",
        "\n",
        "\n",
        "def print_quality(image_file):\n",
        "    # Load the image\n",
        "    image = cv2.imread(image_file)\n",
        "\n",
        "    # Get image metrics\n",
        "    resolution = get_resolution(image)\n",
        "    noise_level = get_noise_level(image)\n",
        "    sharpness = get_sharpness(image)\n",
        "    try:\n",
        "      contrast = get_contrast(image)\n",
        "    except:\n",
        "      contrast = \"unknown\"\n",
        "    image_size = get_filesize(image_file)\n",
        "\n",
        "    # Output the results\n",
        "    print(f\"Resolution: {resolution} pixels\")\n",
        "    print(f\"Noise Level: {noise_level} dB\")\n",
        "    print(f\"Sharpness: {sharpness}\")\n",
        "    print(f\"Contrast: {contrast}\")\n",
        "    print(f\"Size of image: {image_size} MB\")"
      ],
      "metadata": {
        "id": "esWmYaaSROAt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Here we go woohoo :)"
      ],
      "metadata": {
        "id": "iUyFFEyoFqeM"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ok so now you can run the code cell below and scroll a little down to below the code cell. It will prompt you to upload an image file. After it is uploaded, it will enhance your image and display it for you.\n",
        "\n",
        "You can change the values of \"layers\" and \"enhance_faces\" in the actual method below which has been cordoned off for you! ;)"
      ],
      "metadata": {
        "id": "RWI0OZ7HMJDF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import files\n",
        "\n",
        "Image.MAX_IMAGE_PIXELS = 933120000\n",
        "\n",
        "# Upload a file\n",
        "uploaded = files.upload()\n",
        "\n",
        "# Get the file name\n",
        "filename = list(uploaded.keys())[0]\n",
        "\n",
        "print(f\"File uploaded: {filename}\")\n",
        "\n",
        "# Start timing\n",
        "t = Timer()\n",
        "t.start\n",
        "\n",
        "########################################################################################################################################################\n",
        "\n",
        "# You can edit the layers and enhance_faces parameters if you want.\n",
        "result = enhance_image(\n",
        "    input_file=filename, # DO NOT CHANGE THIS LINE\n",
        "    layers=2, # Choose either 2 or 4 as the value here.\n",
        "    upscale=1.5, # This value indicates the no of times the output image's resolution needs to enlarged from the original.\n",
        "    enhance_faces=True # Choose between 'True' or 'False' [First letter is capital!]\n",
        "  )\n",
        "\n",
        "########################################################################################################################################################\n",
        "\n",
        "t.end # End timing\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "import numpy as np\n",
        "\n",
        "# Image paths (replace these with your image paths)\n",
        "image_paths = [filename, f\"{filename.split('.')[0]}_out.jpg\"]\n",
        "\n",
        "# Load images\n",
        "images = [mpimg.imread(path) for path in image_paths]\n",
        "\n",
        "# Display images in the same row\n",
        "fig, axs = plt.subplots(1, len(images), figsize=(10, 5))  # Adjust the figsize as needed\n",
        "\n",
        "for i, (img, path) in enumerate(zip(images, image_paths)):\n",
        "    axs[i].imshow(img)\n",
        "    axs[i].axis('off')\n",
        "\n",
        "axs[0].set_title(f\"Original Image\")\n",
        "axs[1].set_title(f\"Enhanced Image\")\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "uTPGERzujrMB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can download the enhanced image to your local computer by opening the Folder icon on the left side panel, and there you will find both your original file and an _out file.\n",
        "\n",
        "The _out file is the enhanced image. You can right click on it and click Download.\n",
        "\n",
        "You can also check out the metrics of your original and enhanced images by running the code cells below:"
      ],
      "metadata": {
        "id": "yR86F23jArrq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_quality(filename)"
      ],
      "metadata": {
        "id": "bvg5Rf7OrtYs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print_quality(f\"{filename.split('.')[0]}_out.jpg\")"
      ],
      "metadata": {
        "id": "lv84RVebNJzg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e590afb5-d0e8-4c9c-dd54-ac653b356eb3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 255\n",
            "Resolution: (5250, 10500) pixels\n",
            "Noise Level: 1.57 dB\n",
            "Sharpness: 91.58\n",
            "Contrast: 1.0\n",
            "Size of image: 14.79 MB\n"
          ]
        }
      ]
    }
  ]
}