colab_inference_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cSLfR2K1fRgR"
      },
      "source": [
        "# Inference Notebook\n",
        "\n",
        "[MAXIM: Multi-Axis MLP for Image Processing (CVPR 2022 Oral)](https://github.com/google-research/maxim)\n",
        "\n",
        "**This is just the inference code. Maximum you can do is to come in with your images and get results using trained models**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXxVwI1SfJc-"
      },
      "source": [
        "# Clone repo and install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RbskPSxADHXM",
        "outputId": "35c15104-7636-4ef6-cfda-f8ffe0e7e3fe"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ""
          ]
        }
      ],
      "source": [
        "! git clone https://github.com/google-research/maxim/\n",
        "%cd ./maxim\n",
        "\n",
        "!pip install -r requirements.txt\n",
        "!pip install --upgrade jax\n",
        "! pip install gdown\n",
        "\n",
        "!python setup.py build\n",
        "! python setup.py install\n",
        "\n",
        "# https://console.cloud.google.com/storage/browser/gresearch/maxim/ckpt/Enhancement/FiveK;tab=objects?prefix=&forceOnObjectsSortingFiltering=false"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l06nrsoVdFRA"
      },
      "source": [
        "# Imports and Defaults\n",
        "Imports from libraries and from the modules written by authors of the repo\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LEbaNP5IdNOQ"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive # works only for colab\n",
        "from PIL import Image\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import collections\n",
        "import importlib\n",
        "import io\n",
        "import os\n",
        "import math\n",
        "import requests\n",
        "from tqdm import tqdm\n",
        "import gdown # to download weights from Drive\n",
        "\n",
        "import flax\n",
        "import jax.numpy as jnp\n",
        "import ml_collections\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from jax.experimental import jax2tf\n",
        "\n",
        "\n",
        "# below code lines are from run_eval.py\n",
        "_MODEL_FILENAME = 'maxim'\n",
        "\n",
        "_MODEL_VARIANT_DICT = {\n",
        "    'Denoising': 'S-3',\n",
        "    'Deblurring': 'S-3',\n",
        "    'Deraining': 'S-2',\n",
        "    'Dehazing': 'S-2',\n",
        "    'Enhancement': 'S-2',\n",
        "}\n",
        "\n",
        "_MODEL_CONFIGS = {\n",
        "    'variant': '',\n",
        "    'dropout_rate': 0.0,\n",
        "    'num_outputs': 3,\n",
        "    'use_bias': True,\n",
        "    'num_supervision_scales': 3,\n",
        "}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EILubkkjc1P5"
      },
      "source": [
        "# Link Google Drive for data input and output \n",
        "Not necessary but ease of use for Data input / Output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w4qz8CMTJonN"
      },
      "outputs": [],
      "source": [
        "# drive.mount('/content/gdrive/',)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "54_eNSRJdUHz"
      },
      "source": [
        "# Helpers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "36QlK8Rfk5ai"
      },
      "outputs": [],
      "source": [
        "def sizeof_fmt(size, suffix='B'):\n",
        "    \"\"\"Get human readable file size.\n",
        "    Args:\n",
        "        size (int): File size.\n",
        "        suffix (str): Suffix. Default: 'B'.\n",
        "    Return:\n",
        "        str: Formated file siz.\n",
        "    \"\"\"\n",
        "    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:\n",
        "        if abs(size) < 1024.0:\n",
        "            return f'{size:3.1f} {unit}{suffix}'\n",
        "        size /= 1024.0\n",
        "    return f'{size:3.1f} Y{suffix}'\n",
        "\n",
        "\n",
        "def download_file_from_google_drive(file_id, save_path):\n",
        "    \"\"\"Download files from google drive.\n",
        "\n",
        "    Ref:\n",
        "    https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive  # noqa E501\n",
        "\n",
        "    Args:\n",
        "        file_id (str): File id.\n",
        "        save_path (str): Save path.\n",
        "    \"\"\"\n",
        "\n",
        "    session = requests.Session()\n",
        "    URL = 'https://docs.google.com/uc?export=download'\n",
        "    params = {'id': file_id}\n",
        "\n",
        "    response = session.get(URL, params=params, stream=True)\n",
        "    token = get_confirm_token(response)\n",
        "    if token:\n",
        "        params['confirm'] = token\n",
        "        response = session.get(URL, params=params, stream=True)\n",
        "\n",
        "    # get file size\n",
        "    response_file_size = session.get(\n",
        "        URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})\n",
        "    if 'Content-Range' in response_file_size.headers:\n",
        "        file_size = int(\n",
        "            response_file_size.headers['Content-Range'].split('/')[1])\n",
        "    else:\n",
        "        file_size = None\n",
        "\n",
        "    save_response_content(response, save_path, file_size)\n",
        "\n",
        "\n",
        "def get_confirm_token(response):\n",
        "    for key, value in response.cookies.items():\n",
        "        if key.startswith('download_warning'):\n",
        "            return value\n",
        "    return None\n",
        "\n",
        "\n",
        "def save_response_content(response,\n",
        "                          destination,\n",
        "                          file_size=None,\n",
        "                          chunk_size=32768):\n",
        "    if file_size is not None:\n",
        "        pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')\n",
        "\n",
        "        readable_file_size = sizeof_fmt(file_size)\n",
        "    else:\n",
        "        pbar = None\n",
        "\n",
        "    with open(destination, 'wb') as f:\n",
        "        downloaded_size = 0\n",
        "        for chunk in response.iter_content(chunk_size):\n",
        "            downloaded_size += chunk_size\n",
        "            if pbar is not None:\n",
        "                pbar.update(1)\n",
        "                pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} '\n",
        "                                     f'/ {readable_file_size}')\n",
        "            if chunk:  # filter out keep-alive new chunks\n",
        "                f.write(chunk)\n",
        "        if pbar is not None:\n",
        "            pbar.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TY2GyIbn6j85"
      },
      "outputs": [],
      "source": [
        "\n",
        "def resize(path, new_width_height = 1280, save_image = False, convert_RGB = True, clip_full_hd = False, quality = 100):\n",
        "  '''\n",
        "  Resize and return Given Image\n",
        "  args:\n",
        "    path: Image Path\n",
        "    new_width_height = Reshaped image's width and height. # If integer is given, it'll keep the aspect ratio as it is by shrinking the Bigger dimension (width or height) to the max of new_width_height  and then shring the smaller dimension accordingly \n",
        "    save_image = Whether to save the image or not\n",
        "    convert_RGB: Whether to Convert the RGBA image to RGB (by default backgroud is white)\n",
        "  '''\n",
        "  image = Image.open(path)\n",
        "  w, h = image.size\n",
        "\n",
        "  fixed_size = new_width_height if isinstance(new_width_height, int) else False\n",
        "\n",
        "  if fixed_size:\n",
        "    if h > w:\n",
        "      fixed_height = fixed_size\n",
        "      height_percent = (fixed_height / float(h))\n",
        "      width_size = int((float(w) * float(height_percent)))\n",
        "      image = image.resize((width_size, fixed_height), Image.NEAREST)\n",
        "\n",
        "    else:\n",
        "      fixed_width = fixed_size\n",
        "      width_percent = (fixed_width / float(w))\n",
        "      height_size = int((float(h) * float(width_percent)))\n",
        "      image = image.resize((fixed_width, height_size), Image.NEAREST) # Try Image.ANTIALIAS inplace of Image.NEAREST\n",
        "\n",
        "  else:\n",
        "    image = image.resize(new_width_height)\n",
        "\n",
        "  if image.mode == \"RGBA\" and convert_RGB:\n",
        "    # image.load() # required for png.split()\n",
        "    # new = Image.new(\"RGB\", image.size, (255, 255, 255)) # White Background\n",
        "    # image = new.paste(image, mask=image.split()[3]) # 3 is the alpha channel\n",
        "\n",
        "    new = Image.new(\"RGBA\", image.size, \"WHITE\") # Create a white rgba background\n",
        "    new.paste(image, (0, 0), image) # Paste the image on the background.\n",
        "    image = new.convert('RGB')\n",
        "\n",
        "  if save_image:\n",
        "    image.save(path, quality = quality)\n",
        "\n",
        "  return image\n",
        "\n",
        "\n",
        "class DummyFlags():\n",
        "  def __init__(self, ckpt_path:str, task:str, input_dir: str = \"./maxim/images/Enhancement\", output_dir:str = \"./maxim/images/Results\", has_target:bool = False, save_images:bool = True, geometric_ensemble:bool = False):\n",
        "    '''\n",
        "    Builds the dummy flags which replicates the behaviour of Terminal CLI execution (same as ArgParse)\n",
        "    args:\n",
        "      ckpt_path: Saved Model CheckPoint: Find all the checkpoints for pre trained models at https://console.cloud.google.com/storage/browser/gresearch/maxim/ckpt/\n",
        "      task: Task for which the model waas trained. Each task uses different Data and Checkpoints. Find the details of tasks and respective checkpoints details at: https://github.com/google-research/maxim#results-and-pre-trained-models\n",
        "      input_dir: Input Directory. We do not need it here as we are directly passing one image at a time\n",
        "      output_dir: Also not needed in out code\n",
        "      has_target: Used to calculate PSNR and SSIM calculation. Not needed in our case\n",
        "      save_images: Used in CLI command where images were saved in loop. Not needed in our case\n",
        "      geometric_ensemble: Was used in training part and as it is just an Inference part, it is not needed\n",
        "\n",
        "    '''\n",
        "    self.ckpt_path = ckpt_path\n",
        "    self.task = task\n",
        "    self.input_dir = input_dir\n",
        "    self.output_dir = output_dir\n",
        "    self.has_target = has_target\n",
        "    self.save_images = save_images\n",
        "    self.geometric_ensemble = geometric_ensemble\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JkgUJDR0daUP"
      },
      "source": [
        "# Refactored code from authors (`run_eval.py`)\n",
        "\n",
        "**NOTE**: This is not my code. I just changed the structure, redirected dependencies within modules, removed redundant imports and code and bla bla bla...."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9oUVYnXQK_WV"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\n",
        "def recover_tree(keys, values):\n",
        "  \"\"\"Recovers a tree as a nested dict from flat names and values.\n",
        "\n",
        "  This function is useful to analyze checkpoints that are saved by our programs\n",
        "  without need to access the exact source code of the experiment. In particular,\n",
        "  it can be used to extract an reuse various subtrees of the scheckpoint, e.g.\n",
        "  subtree of parameters.\n",
        "  Args:\n",
        "    keys: a list of keys, where '/' is used as separator between nodes.\n",
        "    values: a list of leaf values.\n",
        "  Returns:\n",
        "    A nested tree-like dict.\n",
        "  \"\"\"\n",
        "  tree = {}\n",
        "  sub_trees = collections.defaultdict(list)\n",
        "  for k, v in zip(keys, values):\n",
        "    if '/' not in k:\n",
        "      tree[k] = v\n",
        "    else:\n",
        "      k_left, k_right = k.split('/', 1)\n",
        "      sub_trees[k_left].append((k_right, v))\n",
        "  for k, kv_pairs in sub_trees.items():\n",
        "    k_subtree, v_subtree = zip(*kv_pairs)\n",
        "    tree[k] = recover_tree(k_subtree, v_subtree)\n",
        "  return tree\n",
        "\n",
        "\n",
        "def mod_padding_symmetric(image, factor=64):\n",
        "  \"\"\"Padding the image to be divided by factor.\"\"\"\n",
        "  height, width = image.shape[0], image.shape[1]\n",
        "  height_pad, width_pad = ((height + factor) // factor) * factor, (\n",
        "      (width + factor) // factor) * factor\n",
        "  padh = height_pad - height if height % factor != 0 else 0\n",
        "  padw = width_pad - width if width % factor != 0 else 0\n",
        "  image = jnp.pad(\n",
        "      image, [(padh // 2, padh // 2), (padw // 2, padw // 2), (0, 0)],\n",
        "      mode='reflect')\n",
        "  return image\n",
        "\n",
        "\n",
        "def get_params(ckpt_path):\n",
        "  \"\"\"Get params checkpoint.\"\"\"\n",
        "\n",
        "  with tf.io.gfile.GFile(ckpt_path, 'rb') as f:\n",
        "    data = f.read()\n",
        "  values = np.load(io.BytesIO(data))\n",
        "  params = recover_tree(*zip(*values.items()))\n",
        "  params = params['opt']['target']\n",
        "\n",
        "  return params\n",
        "\n",
        "\n",
        "def calculate_psnr(img1, img2, crop_border, test_y_channel=False):\n",
        "  \"\"\"Calculate PSNR (Peak Signal-to-Noise Ratio).\n",
        "\n",
        "  Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio\n",
        "  Args:\n",
        "    img1 (ndarray): Images with range [0, 255].\n",
        "    img2 (ndarray): Images with range [0, 255].\n",
        "    crop_border (int): Cropped pixels in each edge of an image. These\n",
        "        pixels are not involved in the PSNR calculation.\n",
        "    test_y_channel (bool): Test on Y channel of YCbCr. Default: False.\n",
        "  Returns:\n",
        "    float: psnr result.\n",
        "  \"\"\"\n",
        "  assert img1.shape == img2.shape, (\n",
        "      f'Image shapes are differnet: {img1.shape}, {img2.shape}.')\n",
        "  img1 = img1.astype(np.float64)\n",
        "  img2 = img2.astype(np.float64)\n",
        "\n",
        "  if crop_border != 0:\n",
        "    img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]\n",
        "    img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]\n",
        "\n",
        "  if test_y_channel:\n",
        "    img1 = to_y_channel(img1)\n",
        "    img2 = to_y_channel(img2)\n",
        "\n",
        "  mse = np.mean((img1 - img2)**2)\n",
        "  if mse == 0:\n",
        "    return float('inf')\n",
        "  return 20. * np.log10(255. / np.sqrt(mse))\n",
        "\n",
        "\n",
        "def _convert_input_type_range(img):\n",
        "  \"\"\"Convert the type and range of the input image.\n",
        "\n",
        "  It converts the input image to np.float32 type and range of [0, 1].\n",
        "  It is mainly used for pre-processing the input image in colorspace\n",
        "  convertion functions such as rgb2ycbcr and ycbcr2rgb.\n",
        "  Args:\n",
        "    img (ndarray): The input image. It accepts:\n",
        "        1. np.uint8 type with range [0, 255];\n",
        "        2. np.float32 type with range [0, 1].\n",
        "  Returns:\n",
        "      (ndarray): The converted image with type of np.float32 and range of\n",
        "          [0, 1].\n",
        "  \"\"\"\n",
        "  img_type = img.dtype\n",
        "  img = img.astype(np.float32)\n",
        "  if img_type == np.float32:\n",
        "    pass\n",
        "  elif img_type == np.uint8:\n",
        "    img /= 255.\n",
        "  else:\n",
        "    raise TypeError('The img type should be np.float32 or np.uint8, '\n",
        "                    f'but got {img_type}')\n",
        "  return img\n",
        "\n",
        "\n",
        "def _convert_output_type_range(img, dst_type):\n",
        "  \"\"\"Convert the type and range of the image according to dst_type.\n",
        "\n",
        "  It converts the image to desired type and range. If `dst_type` is np.uint8,\n",
        "  images will be converted to np.uint8 type with range [0, 255]. If\n",
        "  `dst_type` is np.float32, it converts the image to np.float32 type with\n",
        "  range [0, 1].\n",
        "  It is mainly used for post-processing images in colorspace convertion\n",
        "  functions such as rgb2ycbcr and ycbcr2rgb.\n",
        "  Args:\n",
        "    img (ndarray): The image to be converted with np.float32 type and\n",
        "        range [0, 255].\n",
        "    dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it\n",
        "        converts the image to np.uint8 type with range [0, 255]. If\n",
        "        dst_type is np.float32, it converts the image to np.float32 type\n",
        "        with range [0, 1].\n",
        "  Returns:\n",
        "    (ndarray): The converted image with desired type and range.\n",
        "  \"\"\"\n",
        "  if dst_type not in (np.uint8, np.float32):\n",
        "    raise TypeError('The dst_type should be np.float32 or np.uint8, '\n",
        "                    f'but got {dst_type}')\n",
        "  if dst_type == np.uint8:\n",
        "    img = img.round()\n",
        "  else:\n",
        "    img /= 255.\n",
        "\n",
        "  return img.astype(dst_type)\n",
        "\n",
        "\n",
        "def rgb2ycbcr(img, y_only=False):\n",
        "  \"\"\"Convert a RGB image to YCbCr image.\n",
        "\n",
        "  This function produces the same results as Matlab's `rgb2ycbcr` function.\n",
        "  It implements the ITU-R BT.601 conversion for standard-definition\n",
        "  television. See more details in\n",
        "  https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.\n",
        "  It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.\n",
        "  In OpenCV, it implements a JPEG conversion. See more details in\n",
        "  https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.\n",
        "\n",
        "  Args:\n",
        "    img (ndarray): The input image. It accepts:\n",
        "        1. np.uint8 type with range [0, 255];\n",
        "        2. np.float32 type with range [0, 1].\n",
        "    y_only (bool): Whether to only return Y channel. Default: False.\n",
        "  Returns:\n",
        "    ndarray: The converted YCbCr image. The output image has the same type\n",
        "        and range as input image.\n",
        "  \"\"\"\n",
        "  img_type = img.dtype\n",
        "  img = _convert_input_type_range(img)\n",
        "  if y_only:\n",
        "    out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0\n",
        "  else:\n",
        "    out_img = np.matmul(img,\n",
        "                        [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],\n",
        "                         [24.966, 112.0, -18.214]]) + [16, 128, 128]\n",
        "  out_img = _convert_output_type_range(out_img, img_type)\n",
        "  return out_img\n",
        "\n",
        "\n",
        "def to_y_channel(img):\n",
        "  \"\"\"Change to Y channel of YCbCr.\n",
        "\n",
        "  Args:\n",
        "    img (ndarray): Images with range [0, 255].\n",
        "  Returns:\n",
        "    (ndarray): Images with range [0, 255] (float type) without round.\n",
        "  \"\"\"\n",
        "  img = img.astype(np.float32) / 255.\n",
        "  if img.ndim == 3 and img.shape[2] == 3:\n",
        "    img = rgb2ycbcr(img, y_only=True)\n",
        "    img = img[..., None]\n",
        "  return img * 255.\n",
        "\n",
        "\n",
        "def augment_image(image, times=8):\n",
        "  \"\"\"Geometric augmentation.\"\"\"\n",
        "  if times == 4:  # only rotate image\n",
        "    images = []\n",
        "    for k in range(0, 4):\n",
        "      images.append(np.rot90(image, k=k))\n",
        "    images = np.stack(images, axis=0)\n",
        "  elif times == 8:  # roate and flip image\n",
        "    images = []\n",
        "    for k in range(0, 4):\n",
        "      images.append(np.rot90(image, k=k))\n",
        "    image = np.fliplr(image)\n",
        "    for k in range(0, 4):\n",
        "      images.append(np.rot90(image, k=k))\n",
        "    images = np.stack(images, axis=0)\n",
        "  else:\n",
        "    raise Exception(f'Error times: {times}')\n",
        "  return images\n",
        "\n",
        "\n",
        "def deaugment_image(images, times=8):\n",
        "  \"\"\"Reverse the geometric augmentation.\"\"\"\n",
        "\n",
        "  if times == 4:  # only rotate image\n",
        "    image = []\n",
        "    for k in range(0, 4):\n",
        "      image.append(np.rot90(images[k], k=4-k))\n",
        "    image = np.stack(image, axis=0)\n",
        "    image = np.mean(image, axis=0)\n",
        "  elif times == 8:  # roate and flip image\n",
        "    image = []\n",
        "    for k in range(0, 4):\n",
        "      image.append(np.rot90(images[k], k=4-k))\n",
        "    for k in range(0, 4):\n",
        "      image.append(np.fliplr(np.rot90(images[4+k], k=4-k)))\n",
        "    image = np.mean(image, axis=0)\n",
        "  else:\n",
        "    raise Exception(f'Error times: {times}')\n",
        "  return image\n",
        "\n",
        "\n",
        "def is_image_file(filename):\n",
        "  \"\"\"Check if it is an valid image file by extension.\"\"\"\n",
        "  return any(\n",
        "      filename.endswith(extension)\n",
        "      for extension in ['jpeg', 'JPEG', 'jpg', 'png', 'JPG', 'PNG', 'gif'])\n",
        "\n",
        "\n",
        "def save_img(img, pth):\n",
        "  \"\"\"Save an image to disk.\n",
        "\n",
        "  Args:\n",
        "    img: jnp.ndarry, [height, width, channels], img will be clipped to [0, 1]\n",
        "      before saved to pth.\n",
        "    pth: string, path to save the image to.\n",
        "  \"\"\"\n",
        "  Image.fromarray(np.array(\n",
        "      (np.clip(img, 0., 1.) * 255.).astype(jnp.uint8))).save(pth, 'PNG')\n",
        "\n",
        "\n",
        "def make_shape_even(image):\n",
        "  \"\"\"Pad the image to have even shapes.\"\"\"\n",
        "  height, width = image.shape[0], image.shape[1]\n",
        "  padh = 1 if height % 2 != 0 else 0\n",
        "  padw = 1 if width % 2 != 0 else 0\n",
        "  image = jnp.pad(image, [(0, padh), (0, padw), (0, 0)], mode='reflect')\n",
        "  return image\n",
        "\n",
        "\n",
        "# Refactored code --------------------------------------------------------------------------------------------------------------------\n",
        "\n",
        "def build_model(task = \"Enhancement\"):\n",
        "  model_mod = importlib.import_module(f'maxim.models.{_MODEL_FILENAME}')\n",
        "  model_configs = ml_collections.ConfigDict(_MODEL_CONFIGS)\n",
        "\n",
        "  model_configs.variant = _MODEL_VARIANT_DICT[task]\n",
        "\n",
        "  model = model_mod.Model(**model_configs)\n",
        "  return model\n",
        "\n",
        "\n",
        "def pre_process(input_file):\n",
        "  '''\n",
        "  Pre-process the image before sending to the model\n",
        "  '''\n",
        "  input_img = np.asarray(Image.open(input_file).convert('RGB'),np.float32) / 255.\n",
        "  # Padding images to have even shapes\n",
        "  height, width = input_img.shape[0], input_img.shape[1]\n",
        "  input_img = make_shape_even(input_img)\n",
        "  height_even, width_even = input_img.shape[0], input_img.shape[1]\n",
        "\n",
        "  # padding images to be multiplies of 64\n",
        "  input_img = mod_padding_symmetric(input_img, factor=64)\n",
        "  input_img = np.expand_dims(input_img, axis=0)\n",
        "\n",
        "  return input_img, height, width, height_even, width_even\n",
        "\n",
        "\n",
        "def predict(input_img):\n",
        "  # handle multi-stage outputs, obtain the last scale output of last stage\n",
        "  return model.apply({'params': flax.core.freeze(params)}, input_img)\n",
        "\n",
        "\n",
        "def post_process(preds, height, width, height_even, width_even):\n",
        "  '''\n",
        "  Post process the image coming out from prediction\n",
        "  '''\n",
        "  if isinstance(preds, list):\n",
        "    preds = preds[-1]\n",
        "    if isinstance(preds, list):\n",
        "      preds = preds[-1]\n",
        "\n",
        "  # De-ensemble by averaging inferenced results.\n",
        "  preds = np.array(preds[0], np.float32)\n",
        "\n",
        "  # unpad images to get the original resolution\n",
        "  new_height, new_width = preds.shape[0], preds.shape[1]\n",
        "  h_start = new_height // 2 - height_even // 2\n",
        "  h_end = h_start + height\n",
        "  w_start = new_width // 2 - width_even // 2\n",
        "  w_end = w_start + width\n",
        "  preds = preds[h_start:h_end, w_start:w_end, :]\n",
        "  return np.array((np.clip(preds, 0., 1.) * 255.).astype(jnp.uint8))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JOfe8u7_Wxks"
      },
      "source": [
        "# Default Configs and Model Building\n",
        "**Steps**:\n",
        "1. Get the name of `task` and the respective `ckpt` (pre-trained saved model for that task) [Follow this link for task name and model](https://github.com/google-research/maxim#results-and-pre-trained-models)\n",
        "2. Pass in the proper `task` and `ckpt_path` to the `DummyFlags`\n",
        "3. Build Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fcp68HNFf2Fy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "09761106-e8cd-4880-bea6-e6e0840effff"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1-BRKozXh81PtwoMZ9QN3kCAieLzozHIq\n",
            "To: /content/maxim/adobe.npz\n",
            "100%|██████████| 172M/172M [00:01<00:00, 166MB/s]\n"
          ]
        }
      ],
      "source": [
        "weight_drive_path = 'https://drive.google.com/uc?id=1-BRKozXh81PtwoMZ9QN3kCAieLzozHIq' # Path of the weights file which in the Google Drive\n",
        "MODEL_PATH = './adobe.npz' # name of the model to be saved as\n",
        "\n",
        "gdown.download(weight_drive_path, MODEL_PATH, quiet=False) # Download Model weights to your current instance\n",
        "\n",
        "\n",
        "FLAGS = DummyFlags(ckpt_path = MODEL_PATH, task = \"Enhancement\") # Path to your checkpoint and task name\n",
        "\n",
        "params = get_params(FLAGS.ckpt_path) # Parse the config\n",
        "\n",
        "model = build_model() # Build Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c_zxNUE4TugU"
      },
      "source": [
        "# Inference\n",
        "For Inference, you just need to pasd the *Image Path* to the the `predict` function. Result will be a `Numpy` array. You can easily save that by converting to `PIL` image.\n",
        "\n",
        "\n",
        "**Note**: You might get `OOM` or Out of memory issue which is not a big deal if you image size is too big. In that case, you just need to use the `resize` function\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jC35dD6ViBwZ"
      },
      "outputs": [],
      "source": [
        "# image_path = \"path/to/my/image.extension\" # your image path\n",
        "# enhanced_image_array = predict(image_path) # Get predictions\n",
        "\n",
        "# enhanced_pil_image = Image.fromarray(enhanced_image_array) # get PIL image from array\n",
        "# enhanced_pil_image.save(\"path/to/output/directory/image.extension\") # Save the image\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7x5dCbuOvIDz"
      },
      "source": [
        "# Test Images from Drive and Save\n",
        "\n",
        "**Note**: For huge number of images (say 50 or more), copy all the images from Google Drive to the current machine's drive else it will make the process so slow. And also for saving the enhanced image to drive, Get predictions for all the images at once, Save them here first and them copy a zip file to the drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W3YwEUHmtQSl"
      },
      "outputs": [],
      "source": [
        "# images = [\"../gdrive/My Drive/maxim/input/\"+i for i in os.listdir(\"../gdrive/My Drive/maxim/input/\") if i.endswith(('jpeg', 'png', 'jpg',\"PNG\",\"JPEG\",\"JPG\"))]\n",
        "\n",
        "# # _ = [resize(path, 1920, save_image=True) for path in images] # Resize Images to 1920 as the max dimension's size else it'll blow the GPU / CPU memory\n",
        "\n",
        "\n",
        "# for path in images:\n",
        "#   im = Image.fromarray(predict(path))\n",
        "#   im.save(\"../gdrive/My Drive/maxim/output/\"+path.split('/')[-1])\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9Wv9yIIamJL"
      },
      "source": [
        "# Visualization\n",
        "\n",
        "The below code demonstrates how to predict from Image URL. You can directly use `predict(image_path)`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q1T544sZan4d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 491
        },
        "outputId": "17f4b919-c811-4659-a178-95d02d8660dc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[1;30;43mThis cell output is too large and can only be displayed while logged in.\u001b[0m\n"
          ]
        }
      ],
      "source": [
        "import requests\n",
        "from io import BytesIO\n",
        "\n",
        "url = \"https://phototraces.b-cdn.net/wp-content/uploads/2021/02/id_Free_RAW_Photos_for_Editing_09_Uneditedd.jpg\"\n",
        "# url = \"https://phototraces.b-cdn.net/wp-content/uploads/2021/03/Free_RAW_Photos_for_Editing_13_Unedited.jpg\"\n",
        "\n",
        "image_bytes = BytesIO(requests.get(url).content)\n",
        "\n",
        "input_img, height, width, height_even, width_even = pre_process(image_bytes)\n",
        "preds = predict(input_img)\n",
        "result = post_process(preds, height, width, height_even, width_even)\n",
        "\n",
        "f, ax = plt.subplots(1,2, figsize = (35,20))\n",
        "\n",
        "ax[0].imshow(np.array(Image.open(image_bytes))) # Original image\n",
        "ax[1].imshow(result) # retouched image\n",
        "\n",
        "ax[0].set_title(\"Original Image\")\n",
        "ax[1].set_title(\"Enhanced Image\")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf_predict = tf.function(\n",
        "    jax2tf.convert(predict, enable_xla=False),\n",
        "    input_signature=[\n",
        "        tf.TensorSpec(shape=[1, 704, 1024, 3], dtype=tf.float32, name='input_image')\n",
        "    ],\n",
        "    autograph=False)"
      ],
      "metadata": {
        "id": "nZGyNC8pttv7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "converter = tf.lite.TFLiteConverter.from_concrete_functions(\n",
        "    [tf_predict.get_concrete_function()], tf_predict)\n",
        "\n",
        "converter.target_spec.supported_ops = [\n",
        "    tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.\n",
        "    tf.lite.OpsSet.SELECT_TF_OPS  # enable TensorFlow ops.\n",
        "]\n",
        "tflite_float_model = converter.convert()\n",
        "\n",
        "with open('./float_model.tflite', \"wb\") as f: f.write(tflite_float_model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "igxeGKfIyqLS",
        "outputId": "2c7d6ae3-8c84-421f-dc83-f6e0c69cc6c5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n",
            "WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_quantized_model = converter.convert()\n",
        "\n",
        "with open('./quantized.tflite', 'wb') as f: f.write(tflite_quantized_model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pTbjUpn-yxl1",
        "outputId": "fcb29b97-825c-4d6b-b921-0452bcf003a0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Load quantized TFLite model\n",
        "tflite_interpreter_quant = tf.lite.Interpreter(model_path='./maxim/quantized.tflite')\n",
        "\n",
        "# Learn about its input and output details\n",
        "input_details = tflite_interpreter_quant.get_input_details()\n",
        "output_details = tflite_interpreter_quant.get_output_details()\n",
        "\n",
        "# Resize input and output tensors to handle batch of desired size\n",
        "# tflite_interpreter_quant.resize_tensor_input(input_details[0]['index'], (1, 704, 1024, 3))\n",
        "# tflite_interpreter_quant.resize_tensor_input(output_details[0]['index'], (1, 176, 256, 3))\n",
        "tflite_interpreter_quant.allocate_tensors()\n",
        "\n",
        "input_details = tflite_interpreter_quant.get_input_details()\n",
        "output_details = tflite_interpreter_quant.get_output_details()\n",
        "\n",
        "\n",
        "# # Run inference\n",
        "val_image_batch = tf.random.normal(shape = (1, 704, 1024, 3), dtype = tf.float32)\n",
        "tflite_interpreter_quant.set_tensor(input_details[0]['index'], val_image_batch)\n",
        "\n",
        "tflite_interpreter_quant.invoke()\n",
        "\n",
        "tflite_q_model_predictions = tflite_interpreter_quant.get_tensor(output_details[0]['index'])\n",
        "print(\"\\nPrediction results shape:\", tflite_q_model_predictions.shape)"
      ],
      "metadata": {
        "id": "XY0M5SFg4Zw8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "RgsmrC0U7v5h"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}