notebooks/202-vision-superresolution/202-vision-superresolution-image.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ab3e1c1",
   "metadata": {},
   "source": [
    "# Super Resolution with OpenVINO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a261517",
   "metadata": {},
   "source": [
    "Super Resolution is the process of enhancing the quality of an image by increasing the pixel count using deep learning. This notebook shows the Single Image Super Resolution (SISR) which takes just one low resolution image. We use a model called [single-image-super-resolution-1032](https://github.com/openvinotoolkit/open_model_zoo/tree/develop/models/intel/single-image-super-resolution-1032) which is available from the Open Model Zoo. It is based on the research paper cited below. \n",
    "\n",
    "Y. Liu et al., [\"An Attention-Based Approach for Single Image Super Resolution,\"](https://arxiv.org/abs/1807.06779) 2018 24th International Conference on Pattern Recognition (ICPR), 2018, pp. 2777-2784, doi: 10.1109/ICPR.2018.8545760."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c00428fe",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac2e7027",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d4b04ac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import urllib\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython.display import HTML, FileLink\n",
    "from IPython.display import Image as DisplayImage\n",
    "from IPython.display import Pretty, ProgressBar, clear_output, display\n",
    "from openvino.inference_engine import IECore\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb55a28c",
   "metadata": {},
   "source": [
    "### Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "322df115",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Device to use for inference. For example, \"CPU\", or \"GPU\"\n",
    "DEVICE = \"CPU\"\n",
    "# 1032: 4x superresolution, 1033: 3x superresolution\n",
    "MODEL_FILE = \"model/single-image-super-resolution-1032.xml\"\n",
    "model_name = os.path.basename(MODEL_FILE)\n",
    "model_xml_path = Path(MODEL_FILE).with_suffix(\".xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac90276e",
   "metadata": {},
   "source": [
    "### Functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1426c7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def write_text_on_image(image: np.ndarray, text: str) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Write the specified text in the top left corner of the image\n",
    "    as white text with a black border.\n",
    "\n",
    "    :param image: image as numpy arry with HWC shape, RGB or BGR\n",
    "    :param text: text to write\n",
    "    :return: image with written text, as numpy array\n",
    "    \"\"\"\n",
    "    font = cv2.FONT_HERSHEY_PLAIN\n",
    "    org = (20, 20)\n",
    "    font_scale = 4\n",
    "    font_color = (255, 255, 255)\n",
    "    line_type = 1\n",
    "    font_thickness = 2\n",
    "    text_color_bg = (0, 0, 0)\n",
    "    x, y = org\n",
    "\n",
    "    image = cv2.UMat(image)\n",
    "    (text_w, text_h), _ = cv2.getTextSize(\n",
    "        text, font, font_scale, font_thickness\n",
    "    )\n",
    "    result_im = cv2.rectangle(\n",
    "        image, org, (x + text_w, y + text_h), text_color_bg, -1\n",
    "    )\n",
    "\n",
    "    textim = cv2.putText(\n",
    "        result_im,\n",
    "        text,\n",
    "        (x, y + text_h + font_scale - 1),\n",
    "        font,\n",
    "        font_scale,\n",
    "        font_color,\n",
    "        font_thickness,\n",
    "        line_type,\n",
    "    )\n",
    "    return textim.get()\n",
    "\n",
    "\n",
    "def load_image(path: str) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Loads an image from `path` and returns it as BGR numpy array.\n",
    "\n",
    "    :param path: path to an image filename or url\n",
    "    :return: image as numpy array, with BGR channel order\n",
    "    \"\"\"\n",
    "    if path.startswith(\"http\"):\n",
    "        # Set User-Agent to Mozilla because some websites block requests\n",
    "        # with User-Agent Python\n",
    "        request = urllib.request.Request(\n",
    "            path, headers={\"User-Agent\": \"Mozilla/5.0\"}\n",
    "        )\n",
    "        response = urllib.request.urlopen(request)\n",
    "        array = np.asarray(bytearray(response.read()), dtype=\"uint8\")\n",
    "        image = cv2.imdecode(array, -1)  # Loads the image as BGR\n",
    "    else:\n",
    "        image = cv2.imread(path)\n",
    "    return image\n",
    "\n",
    "\n",
    "def convert_result_to_image(result) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Convert network result of floating point numbers to image with integer\n",
    "    values from 0-255. Values outside this range are clipped to 0 and 255.\n",
    "\n",
    "    :param result: a single superresolution network result in N,C,H,W shape\n",
    "    \"\"\"\n",
    "    result = result.squeeze(0).transpose(1, 2, 0)\n",
    "    result *= 255\n",
    "    result[result < 0] = 0\n",
    "    result[result > 255] = 255\n",
    "    result = result.astype(np.uint8)\n",
    "    return result\n",
    "\n",
    "\n",
    "def to_rgb(image_data) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Convert image_data from BGR to RGB\n",
    "    \"\"\"\n",
    "    return cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "944a341f",
   "metadata": {},
   "source": [
    "## Load the superresolution model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53821181",
   "metadata": {},
   "source": [
    "Load the model in Inference Engine with `ie.read_network` and load it to the specified device with `ie.load_network`\n",
    "\n",
    "The Super Resolution model expects two inputs: 1) the input image, 2) a bicubic interpolation of the input image to a size of 1920x1080. It returns the super resolution version of the image in 1920x1800 (for the default superresolution model (1032)). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f76cfb",
   "metadata": {},
   "source": [
    "Load the model in Inference Engine with `ie.read_network` and load it to the specified device with `ie.load_network`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96866b70",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ie = IECore()\n",
    "net = ie.read_network(str(model_xml_path))\n",
    "exec_net = ie.load_network(network=net, device_name=DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c836fa2",
   "metadata": {},
   "source": [
    "Get information about network inputs and outputs. The Super Resolution model expects two inputs: 1) the input image, 2) a bicubic interpolation of the input image to the target size 1920x1080. It returns the super resolution version of the image in 1920x1800. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f84b9e15",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Network inputs and outputs are dictionaries. Get the keys for the\n",
    "# dictionaries.\n",
    "original_image_key = list(exec_net.input_info)[0]\n",
    "bicubic_image_key = list(exec_net.input_info)[1]\n",
    "output_key = list(exec_net.outputs.keys())[0]\n",
    "\n",
    "# Get the expected input and target shape. `.dims[2:]` returns the height\n",
    "# and width. OpenCV's resize function expects the shape as (width, height),\n",
    "# so we reverse the shape with `[::-1]` and convert it to a tuple\n",
    "input_height, input_width = tuple(\n",
    "    exec_net.input_info[original_image_key].tensor_desc.dims[2:]\n",
    ")\n",
    "target_height, target_width = tuple(\n",
    "    exec_net.input_info[bicubic_image_key].tensor_desc.dims[2:]\n",
    ")\n",
    "\n",
    "upsample_factor = int(target_height / input_height)\n",
    "\n",
    "print(\n",
    "    f\"The network expects inputs with a width of {input_width}, \"\n",
    "    f\"height of {input_height}\"\n",
    ")\n",
    "print(\n",
    "    f\"The network returns images with a width of {target_width}, \"\n",
    "    f\"height of {target_height}\"\n",
    ")\n",
    "\n",
    "print(\n",
    "    f\"The image sides are upsampled by a factor {upsample_factor}. \"\n",
    "    f\"The new image is {upsample_factor**2} times as large as the \"\n",
    "    \"original image\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa08460",
   "metadata": {},
   "source": [
    "## Load and show the input image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ae46441",
   "metadata": {},
   "source": [
    "**NOTE:** For best results, use raw images (like TIFF, BMP or PNG). Compressed images (like JPEG) may appear distorted after processing with the super resolution model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cbbbd13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "IMAGE_PATH = Path(\"data/tower.jpg\")\n",
    "OUTPUT_PATH = Path(\"output/\")\n",
    "\n",
    "os.makedirs(str(OUTPUT_PATH), exist_ok=True)\n",
    "full_image = load_image(str(IMAGE_PATH))\n",
    "\n",
    "# Uncomment these lines to load a raw image as BGR\n",
    "# import rawpy\n",
    "# with rawpy.imread(IMAGE_PATH) as raw:\n",
    "#     full_image = raw.postprocess()[:,:,(2,1,0)]\n",
    "\n",
    "plt.imshow(to_rgb(full_image))\n",
    "print(\n",
    "    f\"Showing full image with width {full_image.shape[1]} \"\n",
    "    f\"and height {full_image.shape[0]}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eecb128",
   "metadata": {},
   "source": [
    "## Superresolution on one crop of the image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc5d16dd",
   "metadata": {},
   "source": [
    "### Take one crop of the input image\n",
    "\n",
    "Take a crop of network input size. Give the X (width) and Y (height) coordinates for the top left corner of the crop. Set the CROP_FACTOR variable to 2 to make a crop that is larger than the network input size (this only works with the 1032 superresolution model). The crop will be downsampled before propagating to the network. This is useful for very high resolution images, where a crop of network input size is too small to show enough information. It can also improve the result. Note though that with a CROP_FACTOR or 2 the net upsampling factor will be halved. If the superresolution network increases the side lengths of the image by a factor four, it upsamples a 480x270 crop to 1920x1080. With a CROP_FACTOR of two, a 960x540 crop is upsampled to the same 1920x1080: the side lengths are twice as large as the crop size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f86ec4d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set CROP_FACTOR to 2 to take crops with twice the input width and height\n",
    "# This only works with the 1032 (4x) superresolution model!\n",
    "# Set it to 1 to take crops of exactly input size\n",
    "CROP_FACTOR = 2\n",
    "adjusted_upsample_factor = upsample_factor // CROP_FACTOR\n",
    "\n",
    "image_id = \"flag\"  # A tag to recognize the saved images\n",
    "starty = 3200\n",
    "startx = 0\n",
    "\n",
    "# Perform the crop\n",
    "image_crop = full_image[\n",
    "    starty : starty + input_height * CROP_FACTOR,\n",
    "    startx : startx + input_width * CROP_FACTOR,\n",
    "]\n",
    "\n",
    "# Show the cropped image\n",
    "print(\n",
    "    f\"Showing image crop with width {image_crop.shape[1]} and \"\n",
    "    f\"height {image_crop.shape[0]}.\"\n",
    ")\n",
    "plt.imshow(to_rgb(image_crop));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5d8917b",
   "metadata": {},
   "source": [
    "### Reshape/resize crop for network input\n",
    "\n",
    "The input image is resized to network input size, and reshaped to (N,C,H,W) (H=height, W=width, C=number of channels, N=number of images). The image is also resized to network output size, with bicubic interpolation. This bicubic image is the second input to the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd2968e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Resize the image to the target shape with bicubic interpolation\n",
    "bicubic_image = cv2.resize(\n",
    "    image_crop, (target_width, target_height), interpolation=cv2.INTER_CUBIC\n",
    ")\n",
    "\n",
    "# If required, resize image to input image shape\n",
    "if CROP_FACTOR > 1:\n",
    "    image_crop = cv2.resize(image_crop, (input_width, input_height))\n",
    "\n",
    "# Reshape the images from (H,W,C) to (N,C,H,W)\n",
    "input_image_original = np.expand_dims(image_crop.transpose(2, 0, 1), axis=0)\n",
    "input_image_bicubic = np.expand_dims(bicubic_image.transpose(2, 0, 1), axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2e217a6",
   "metadata": {},
   "source": [
    "### Do inference\n",
    "\n",
    "Do inference and convert the inference result to an RGB image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b649dd4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "network_result = exec_net.infer(\n",
    "    inputs={\n",
    "        original_image_key: input_image_original,\n",
    "        bicubic_image_key: input_image_bicubic,\n",
    "    }\n",
    ")\n",
    "# Get inference result as numpy array and reshape to image shape and data type\n",
    "result_image = convert_result_to_image(network_result[output_key])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d16b4928",
   "metadata": {},
   "source": [
    "### Show and save results\n",
    "\n",
    "Show the bicubic image and the enhanced superresolution image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3349cbc2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(30, 15))\n",
    "ax[0].imshow(to_rgb(bicubic_image))\n",
    "ax[1].imshow(to_rgb(result_image))\n",
    "ax[0].set_title(\"Bicubic\")\n",
    "ax[1].set_title(\"Superresolution\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a643428",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Save superresolution and bicubic image crop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7efcad7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add text with \"SUPER\" or \"BICUBIC\" to the superresolution or bicubic image\n",
    "image_super = write_text_on_image(result_image, \"SUPER\")\n",
    "image_bicubic = write_text_on_image(bicubic_image, \"BICUBIC\")\n",
    "\n",
    "# Store the image and the results\n",
    "crop_image_path = Path(f\"{OUTPUT_PATH.stem}/{image_id}_{adjusted_upsample_factor}x_crop.png\")\n",
    "superres_image_path = Path(f\"{OUTPUT_PATH.stem}/{image_id}_{adjusted_upsample_factor}x_crop_superres.png\")\n",
    "bicubic_image_path = Path(f\"{OUTPUT_PATH.stem}/{image_id}_{adjusted_upsample_factor}x_crop_bicubic.png\")\n",
    "cv2.imwrite(str(crop_image_path), image_crop, [cv2.IMWRITE_PNG_COMPRESSION, 0])\n",
    "cv2.imwrite(\n",
    "    str(superres_image_path), image_super, [cv2.IMWRITE_PNG_COMPRESSION, 0]\n",
    ")\n",
    "cv2.imwrite(\n",
    "    str(bicubic_image_path), image_bicubic, [cv2.IMWRITE_PNG_COMPRESSION, 0]\n",
    ")\n",
    "print(f\"Images written to directory: {OUTPUT_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d217c577",
   "metadata": {},
   "source": [
    "#### Write animated gif with bicubic/superresolution comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2f19391",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_pil = Image.fromarray(to_rgb(image_super))\n",
    "bicubic_pil = Image.fromarray(to_rgb(image_bicubic))\n",
    "gif_image_path = Path(f\"{OUTPUT_PATH.stem}/{image_id}_comparison_{adjusted_upsample_factor}x.gif\")\n",
    "\n",
    "result_pil.save(\n",
    "    fp=str(gif_image_path),\n",
    "    format=\"GIF\",\n",
    "    append_images=[bicubic_pil],\n",
    "    save_all=True,\n",
    "    duration=1000,\n",
    "    loop=0,\n",
    ")\n",
    "\n",
    "# DisplayImage(str(gif_image_path)) doesn't work in Colab\n",
    "DisplayImage(open(gif_image_path, \"rb\").read(), width=1920 // 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e41243",
   "metadata": {},
   "source": [
    "#### Create video with sliding bicubic/superresolution comparison\n",
    "\n",
    "This may take a while. For the video, the superresolution and bicubic image are resized by a factor two to improve processing speed. This gives an indication of the superresolution effect. The video is saved as an .avi video. You can click on the link to download the video, or open it directly from the images directory, and play it locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b3b1db5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "FOURCC = cv2.VideoWriter_fourcc(*\"MJPG\")\n",
    "\n",
    "result_video_path = Path(f\"{OUTPUT_PATH.stem}/{image_id}_crop_comparison_{adjusted_upsample_factor}x.avi\")\n",
    "video_target_height, video_target_width = (\n",
    "    result_image.shape[0] // 2,\n",
    "    result_image.shape[1] // 2,\n",
    ")\n",
    "\n",
    "out_video = cv2.VideoWriter(\n",
    "    str(result_video_path),\n",
    "    FOURCC,\n",
    "    90,\n",
    "    (video_target_width, video_target_height),\n",
    ")\n",
    "\n",
    "resized_result_image = cv2.resize(\n",
    "    result_image, (video_target_width, video_target_height)\n",
    ")\n",
    "resized_bicubic_image = cv2.resize(\n",
    "    bicubic_image, (video_target_width, video_target_height)\n",
    ")\n",
    "\n",
    "progress_bar = ProgressBar(total=video_target_width)\n",
    "progress_bar.display()\n",
    "\n",
    "for i in range(video_target_width):\n",
    "    # Create a frame where the left part (until i pixels width) contains the\n",
    "    # superresolution image, and the right part (from i pixels width) contains\n",
    "    # the bicubic image\n",
    "    comparison_frame = np.hstack(\n",
    "        (\n",
    "            resized_result_image[:, :i, :],\n",
    "            resized_bicubic_image[:, i:, :],\n",
    "        )\n",
    "    )\n",
    "    # create a small black border line between the superresolution\n",
    "    # and bicubic part of the image\n",
    "    comparison_frame[:, i - 1 : i + 1, :] = 0\n",
    "    out_video.write(comparison_frame)\n",
    "    progress_bar.progress = i\n",
    "    progress_bar.update()\n",
    "out_video.release()\n",
    "clear_output()\n",
    "\n",
    "video_link = FileLink(result_video_path)\n",
    "video_link.html_link_str = \"<a href='%s' download>%s</a>\"\n",
    "display(HTML(f\"The video has been saved to {video_link._repr_html_()}\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12d53209",
   "metadata": {},
   "source": [
    "## Superresolution on full input image\n",
    "\n",
    "Superresolution on the full image is done by dividing the image into patches of equal size, doing superresolution on each path, and then stitching the resulting patches together again. For this demo, patches near the border of the image are ignored.\n",
    "\n",
    "Adjust the `CROPLINES` setting in the next cell if you see boundary effects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad91374e",
   "metadata": {},
   "source": [
    "### Compute patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc35dc09",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set the number of lines to crop from the network result to prevent\n",
    "# boundary effects. Should be an integer >= 1.\n",
    "CROPLINES = 10\n",
    "# See Superresolution on one crop of the image for description of CROP_FACTOR\n",
    "CROP_FACTOR = 2\n",
    "\n",
    "full_image_height, full_image_width = full_image.shape[:2]\n",
    "\n",
    "# Compute x and y coordinates of left top of image tiles\n",
    "x_coords = list(\n",
    "    range(0, full_image_width, input_width * CROP_FACTOR - CROPLINES * 2)\n",
    ")\n",
    "while full_image_width - x_coords[-1] < input_width * CROP_FACTOR:\n",
    "    x_coords.pop(-1)\n",
    "y_coords = list(\n",
    "    range(0, full_image_height, input_height * CROP_FACTOR - CROPLINES * 2)\n",
    ")\n",
    "while full_image_height - y_coords[-1] < input_height * CROP_FACTOR:\n",
    "    y_coords.pop(-1)\n",
    "\n",
    "# Compute the width and height to crop the full image. The full image is\n",
    "# cropped at the border to tiles of input size\n",
    "crop_width = x_coords[-1] + input_width * CROP_FACTOR\n",
    "crop_height = y_coords[-1] + input_height * CROP_FACTOR\n",
    "\n",
    "# Compute the width and height of the target superresolution image\n",
    "new_width = (\n",
    "    x_coords[-1] * (upsample_factor // CROP_FACTOR)\n",
    "    + target_width\n",
    "    - CROPLINES * 2 * (upsample_factor // CROP_FACTOR)\n",
    ")\n",
    "new_height = (\n",
    "    y_coords[-1] * (upsample_factor // CROP_FACTOR)\n",
    "    + target_height\n",
    "    - CROPLINES * 2 * (upsample_factor // CROP_FACTOR)\n",
    ")\n",
    "print(\n",
    "    f\"The output image will have a width of {new_width} \"\n",
    "    f\"and a height of {new_height}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cabe8da",
   "metadata": {},
   "source": [
    "### Do inference\n",
    "\n",
    "The code in this cell reads one patch of the image at a time. Each patch is reshaped to network input shape and upsampled with bicubic interpolation to target shape. Both the original and the bicubic image are propagated through the network. The network result is a numpy array with floating point values, with a shape of (1,3,1920,1080). This array is converted to an 8-bit image with shape (1080,1920,3) and written to `full_superresolution_image`. The bicubic image is written to `full_bicubic_image` for comparison. A progress bar shows the progress of the process. Inference time is measured, as well as total time to process each patch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84306b5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "start_time = time.perf_counter()\n",
    "patch_nr = 0\n",
    "num_patches = len(x_coords) * len(y_coords)\n",
    "progress_bar = ProgressBar(total=num_patches)\n",
    "progress_bar.display()\n",
    "\n",
    "# Crop image to fit tiles of input size\n",
    "full_image_crop = full_image.copy()[:crop_height, :crop_width, :]\n",
    "\n",
    "# Create empty array of target size.\n",
    "full_superresolution_image = np.empty(\n",
    "    (new_height, new_width, 3), dtype=np.uint8\n",
    ")\n",
    "\n",
    "# Create bicubic upsampled image of target size for comparison\n",
    "full_bicubic_image = cv2.resize(\n",
    "    full_image_crop[CROPLINES:-CROPLINES, CROPLINES:-CROPLINES, :],\n",
    "    (new_width, new_height),\n",
    "    interpolation=cv2.INTER_CUBIC,\n",
    ")\n",
    "\n",
    "total_inference_duration = 0\n",
    "for y in y_coords:\n",
    "    for x in x_coords:\n",
    "        patch_nr += 1\n",
    "\n",
    "        # Take a crop of the input image\n",
    "        image_crop = full_image_crop[\n",
    "            y : y + input_height * CROP_FACTOR,\n",
    "            x : x + input_width * CROP_FACTOR,\n",
    "        ]\n",
    "\n",
    "        # Resize the images to the target shape with bicubic interpolation\n",
    "        bicubic_image = cv2.resize(\n",
    "            image_crop,\n",
    "            (target_width, target_height),\n",
    "            interpolation=cv2.INTER_CUBIC,\n",
    "        )\n",
    "\n",
    "        if CROP_FACTOR > 1:\n",
    "            image_crop = cv2.resize(image_crop, (input_width, input_height))\n",
    "\n",
    "        input_image_original = np.expand_dims(\n",
    "            image_crop.transpose(2, 0, 1), axis=0\n",
    "        )\n",
    "\n",
    "        input_image_bicubic = np.expand_dims(\n",
    "            bicubic_image.transpose(2, 0, 1), axis=0\n",
    "        )\n",
    "\n",
    "        # Do inference\n",
    "        inference_start_time = time.perf_counter()\n",
    "\n",
    "        network_result = exec_net.infer(\n",
    "            inputs={\n",
    "                original_image_key: input_image_original,\n",
    "                bicubic_image_key: input_image_bicubic,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        inference_stop_time = time.perf_counter()\n",
    "        inference_duration = inference_stop_time - inference_start_time\n",
    "        total_inference_duration += inference_duration\n",
    "\n",
    "        # Reshape inference result to image shape and data type\n",
    "        result = network_result[output_key]\n",
    "        result_image = convert_result_to_image(result)\n",
    "\n",
    "        # Add the inference result of this patch to the full superresolution\n",
    "        # image\n",
    "        adjusted_upsample_factor = upsample_factor // CROP_FACTOR\n",
    "        new_y = y * adjusted_upsample_factor\n",
    "        new_x = x * adjusted_upsample_factor\n",
    "        full_superresolution_image[\n",
    "            new_y : new_y\n",
    "            + target_height\n",
    "            - CROPLINES * adjusted_upsample_factor * 2,\n",
    "            new_x : new_x\n",
    "            + target_width\n",
    "            - CROPLINES * adjusted_upsample_factor * 2,\n",
    "        ] = result_image[\n",
    "            CROPLINES\n",
    "            * adjusted_upsample_factor : -CROPLINES\n",
    "            * adjusted_upsample_factor,\n",
    "            CROPLINES\n",
    "            * adjusted_upsample_factor : -CROPLINES\n",
    "            * adjusted_upsample_factor,\n",
    "            :,\n",
    "        ]\n",
    "\n",
    "        progress_bar.progress = patch_nr\n",
    "        progress_bar.update()\n",
    "\n",
    "        if patch_nr % 10 == 0:\n",
    "            clear_output(wait=True)\n",
    "            progress_bar.display()\n",
    "            display(\n",
    "                Pretty(\n",
    "                    f\"Processed patch {patch_nr}/{num_patches}. \"\n",
    "                    f\"Inference time: {inference_duration:.2f} seconds \"\n",
    "                    f\"({1/inference_duration:.2f} FPS)\"\n",
    "                )\n",
    "            )\n",
    "\n",
    "end_time = time.perf_counter()\n",
    "duration = end_time - start_time\n",
    "clear_output(wait=True)\n",
    "print(\n",
    "    f\"Processed {num_patches} patches in {duration:.2f} seconds. \"\n",
    "    f\"Total patches per second (including processing): \"\n",
    "    f\"{num_patches/duration:.2f}.\\nInference patches per second: \"\n",
    "    f\"{num_patches/total_inference_duration:.2f} \"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61760bac",
   "metadata": {},
   "source": [
    "### Save superresolution image and the bicubic image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c98569a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "full_superresolution_image_path = Path(f\"{OUTPUT_PATH.stem}/full_superres_{adjusted_upsample_factor}x.jpg\")\n",
    "full_bicubic_image_path = Path(f\"{OUTPUT_PATH.stem}/full_bicubic_{adjusted_upsample_factor}x.jpg\")\n",
    "\n",
    "cv2.imwrite(str(full_superresolution_image_path), full_superresolution_image)\n",
    "cv2.imwrite(str(full_bicubic_image_path), full_bicubic_image);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a262e821",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bicubic_link = FileLink(full_bicubic_image_path)\n",
    "image_link = FileLink(full_superresolution_image_path)\n",
    "bicubic_link.html_link_str = \"<a href='%s' download>%s</a>\"\n",
    "image_link.html_link_str = \"<a href='%s' download>%s</a>\"\n",
    "display(\n",
    "    HTML(\n",
    "        \"The images are saved in the images directory. You can also download \"\n",
    "        \"them by clicking on these links:\"\n",
    "        f\"<ul><li>{image_link._repr_html_()}<li>{bicubic_link._repr_html_()}\"\n",
    "    )\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openvino_env",
   "language": "python",
   "name": "openvino_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}