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

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74e3c567",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Video Super Resolution with OpenVINO™\n",
    "Super Resolution is the process of enhancing the quality of an image by increasing the pixel count using deep learning. This notebook applies Single Image Super Resolution (SISR) to frames in a 360p (480×360) video in 360p resolution. A model called [single-image-super-resolution-1032](https://docs.openvino.ai/2023.0/omz_models_model_single_image_super_resolution_1032.html), which is available in Open Model Zoo, is used in this tutorial. 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.\n",
    "\n",
    "> **NOTE**: The Single Image Super Resolution (SISR) model used in this demo is not optimized for a video. Results may vary depending on the video.\n",
    "\n",
    "#### Table of content:\n",
    "- [Preparation](#Preparation-Uparrow)\n",
    "    - [Install requirements](#Install-requirements-Uparrow)\n",
    "    - [Imports](#Imports-Uparrow)\n",
    "    - [Settings](#Settings-Uparrow)\n",
    "        - [Select inference device](#Select-inference-device-Uparrow)\n",
    "    - [Functions](#Functions-Uparrow)\n",
    "- [Load the Superresolution Model](#Load-the-Superresolution-Model-Uparrow)\n",
    "- [Superresolution on Video](#Superresolution-on-Video-Uparrow)\n",
    "    - [Settings](#Settings-Uparrow)\n",
    "    - [Download and Prepare Video](#Download-and-Prepare-Video-Uparrow)\n",
    "    - [Do Inference](#Do-Inference-Uparrow)\n",
    "    - [Show Side-by-Side Video of Bicubic and Superresolution Version](#Show-Side-by-Side-Video-of-Bicubic-and-Superresolution-Version-Uparrow)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4026645b",
   "metadata": {},
   "source": [
    "## Preparation [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "### Install requirements [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9f658ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"openvino>=2023.1.0\"\n",
    "%pip install -q opencv-python\n",
    "%pip install -q \"pytube>=12.1.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7e726e7",
   "metadata": {},
   "source": [
    "### Imports [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4ef55816",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "from IPython.display import (\n",
    "    HTML,\n",
    "    FileLink,\n",
    "    Pretty,\n",
    "    ProgressBar,\n",
    "    Video,\n",
    "    clear_output,\n",
    "    display,\n",
    ")\n",
    "import openvino as ov\n",
    "from pytube import YouTube"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c6824d93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a download file helper function\n",
    "def download_file(url: str, path: Path) -> None:\n",
    "    \"\"\"Download file.\"\"\"\n",
    "    import urllib.request\n",
    "    path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    urllib.request.urlretrieve(url, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4891215e",
   "metadata": {},
   "source": [
    "### Settings [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275aab16-ab78-4586-a494-64a22ee31a2b",
   "metadata": {},
   "source": [
    "#### Select inference device [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bce615e3-486f-41bb-b769-a50f97cc3f59",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f45054e9779e4520802ee8c810ef4895",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', index=2, options=('CPU', 'GPU', 'AUTO'), value='AUTO')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "core = ov.Core()\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='AUTO',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "582d9cde",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 1032: 4x superresolution, 1033: 3x superresolution\n",
    "model_name = 'single-image-super-resolution-1032'\n",
    "\n",
    "base_model_dir = Path('./model').expanduser()\n",
    "\n",
    "model_xml_name = f'{model_name}.xml'\n",
    "model_bin_name = f'{model_name}.bin'\n",
    "\n",
    "model_xml_path = base_model_dir / model_xml_name\n",
    "model_bin_path = base_model_dir / model_bin_name\n",
    "\n",
    "if not model_xml_path.exists():\n",
    "    base_url = f'https://storage.openvinotoolkit.org/repositories/open_model_zoo/2023.0/models_bin/1/{model_name}/FP16/'\n",
    "    model_xml_url = base_url + model_xml_name\n",
    "    model_bin_url = base_url + model_bin_name\n",
    "\n",
    "    download_file(model_xml_url, model_xml_path)\n",
    "    download_file(model_bin_url, model_bin_path)\n",
    "else:\n",
    "    print(f'{model_name} already downloaded to {base_model_dir}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdd4c452",
   "metadata": {},
   "source": [
    "### Functions [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4818ca56",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b8dc8ea",
   "metadata": {},
   "source": [
    "## Load the Superresolution Model [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Load the model in OpenVINO Runtime with `core.read_model` and compile it for the specified device with `core.compile_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c8fee7d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "core = ov.Core()\n",
    "model = core.read_model(model=model_xml_path)\n",
    "compiled_model = core.compile_model(model=model, device_name=device.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d897ad82",
   "metadata": {},
   "source": [
    "Get information about network inputs and outputs. The Super Resolution model expects two inputs: the input image and a bicubic interpolation of the input image to the target size of 1920x1080. It returns the super resolution version of the image in 1920x1080."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b23a5d80",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Network inputs and outputs are dictionaries. Get the keys for the\n",
    "# dictionaries.\n",
    "original_image_key, bicubic_image_key = compiled_model.inputs\n",
    "output_key = compiled_model.output(0)\n",
    "\n",
    "# Get the expected input and target shape. The `.dims[2:]` function returns the height\n",
    "# and width.The `resize` function of  OpenCV expects the shape as (width, height),\n",
    "# so reverse the shape with `[::-1]` and convert it to a tuple.\n",
    "input_height, input_width = list(original_image_key.shape)[2:]\n",
    "target_height, target_width = list(bicubic_image_key.shape)[2:]\n",
    "\n",
    "upsample_factor = int(target_height / input_height)\n",
    "\n",
    "print(f\"The network expects inputs with a width of {input_width}, \" f\"height of {input_height}\")\n",
    "print(f\"The network returns images with a width of {target_width}, \" f\"height of {target_height}\")\n",
    "\n",
    "print(\n",
    "    f\"The image sides are upsampled by a factor of {upsample_factor}. \"\n",
    "    f\"The new image is {upsample_factor**2} times as large as the \"\n",
    "    \"original image\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "564fc343",
   "metadata": {},
   "source": [
    "## Superresolution on Video [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Download a YouTube video with `PyTube` and enhance the video quality with superresolution. \n",
    "\n",
    "By default, only the first 100 frames of the video are processed. Change `NUM_FRAMES` in the cell below to modify this. \n",
    "\n",
    "> **NOTE**: The resulting video does not contain audio. The input video should be a landscape video and have an input resolution of 360p (640x360) for the 1032 model, or 480p (720x480) for the 1033 model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8356c7b",
   "metadata": {},
   "source": [
    "### Settings [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "af99b289",
   "metadata": {
    "tags": [],
    "test_replace": {
     "NUM_FRAMES = 100": "NUM_FRAMES = 3"
    }
   },
   "outputs": [],
   "source": [
    "OUTPUT_DIR = \"output\"\n",
    "\n",
    "Path(OUTPUT_DIR).mkdir(exist_ok=True)\n",
    "# Maximum number of frames to read from the input video. Set to 0 to read all frames.\n",
    "NUM_FRAMES = 100\n",
    "# The format for saving the result videos. The `vp09` codec is slow, but widely available.\n",
    "# If you have FFMPEG installed, you can change FOURCC to `*\"THEO\"` to improve video writing speed.\n",
    "FOURCC = cv2.VideoWriter_fourcc(*\"vp09\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a630d4",
   "metadata": {},
   "source": [
    "### Download and Prepare Video [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "76b55b9c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Use pytube to download a video. It downloads to the videos subdirectory.\n",
    "# You can also place a local video there and comment out the following lines\n",
    "VIDEO_URL = \"https://www.youtube.com/watch?v=V8yS3WIkOrA\"\n",
    "yt = YouTube(VIDEO_URL)\n",
    "# Use `yt.streams` to see all available streams. See the PyTube documentation\n",
    "# https://python-pytube.readthedocs.io/en/latest/api.html for advanced\n",
    "# filtering options\n",
    "stream = yt.streams.filter(resolution=\"360p\").first()\n",
    "filename = Path(stream.default_filename.encode(\"ascii\", \"ignore\").decode(\"ascii\")).stem\n",
    "stream.download(output_path=OUTPUT_DIR, filename=filename)\n",
    "print(f\"Video {filename} downloaded to {OUTPUT_DIR}\")\n",
    "\n",
    "# Create Path objects for the input video and the resulting videos.\n",
    "video_path = Path(stream.get_file_path(filename, OUTPUT_DIR))\n",
    "\n",
    "# Path names for the result videos.\n",
    "superres_video_path = Path(f\"{OUTPUT_DIR}/{video_path.stem}_superres.mp4\")\n",
    "bicubic_video_path = Path(f\"{OUTPUT_DIR}/{video_path.stem}_bicubic.mp4\")\n",
    "comparison_video_path = Path(f\"{OUTPUT_DIR}/{video_path.stem}_superres_comparison.mp4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "678b6cf5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Open the video and get the dimensions and the FPS.\n",
    "cap = cv2.VideoCapture(filename=str(video_path))\n",
    "ret, image = cap.read()\n",
    "if not ret:\n",
    "    raise ValueError(f\"The video at '{video_path}' cannot be read.\")\n",
    "fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)\n",
    "\n",
    "if NUM_FRAMES == 0:\n",
    "    total_frames = frame_count\n",
    "else:\n",
    "    total_frames = min(frame_count, NUM_FRAMES)\n",
    "\n",
    "original_frame_height, original_frame_width = image.shape[:2]\n",
    "\n",
    "cap.release()\n",
    "print(\n",
    "    f\"The input video has a frame width of {original_frame_width}, \"\n",
    "    f\"frame height of {original_frame_height} and runs at {fps:.2f} fps\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fd26468",
   "metadata": {},
   "source": [
    "Create a superresolution video, a bicubic video and a comparison video. The superresolution video contains the enhanced video, upsampled with superresolution, the bicubic video is the input video upsampled with bicubic interpolation, the comparison video sets the bicubic video and the superresolution side by side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a6951656-1230-480a-a185-df603efc93d7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "superres_video = cv2.VideoWriter(\n",
    "    filename=str(superres_video_path),\n",
    "    fourcc=FOURCC,\n",
    "    fps=fps,\n",
    "    frameSize=(target_width, target_height),\n",
    ")\n",
    "bicubic_video = cv2.VideoWriter(\n",
    "    filename=str(bicubic_video_path),\n",
    "    fourcc=FOURCC,\n",
    "    fps=fps,\n",
    "    frameSize=(target_width, target_height),\n",
    ")\n",
    "comparison_video = cv2.VideoWriter(\n",
    "    filename=str(comparison_video_path),\n",
    "    fourcc=FOURCC,\n",
    "    fps=fps,\n",
    "    frameSize=(target_width * 2, target_height),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27b5c979",
   "metadata": {},
   "source": [
    "### Do Inference [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Read video frames and enhance them with superresolution. Save the superresolution video, the bicubic video and the comparison video to a file.\n",
    "\n",
    "The code below reads the video frame by frame. Each frame is resized and reshaped to the network input shape and upsampled with bicubic interpolation to the target shape. Both the original and the bicubic images 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 the `(1080,1920,3)` shape and written to a `superres_video`. The bicubic image is written to a `bicubic_video` for comparison. Finally, the bicubic and result frames are combined side by side and written to a `comparison_video`. A progress bar shows the progress of the process. Both inference time and total time to process each frame are measured. That also includes inference time as well as the time it takes to process and write the video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "12ac8de5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "start_time = time.perf_counter()\n",
    "frame_nr = 0\n",
    "total_inference_duration = 0\n",
    "\n",
    "progress_bar = ProgressBar(total=total_frames)\n",
    "progress_bar.display()\n",
    "\n",
    "cap = cv2.VideoCapture(filename=str(video_path))\n",
    "try:\n",
    "    while cap.isOpened():\n",
    "        ret, image = cap.read()\n",
    "        if not ret:\n",
    "            cap.release()\n",
    "            break\n",
    "\n",
    "        if frame_nr >= total_frames:\n",
    "            break\n",
    "\n",
    "        # Resize the input image to the network shape and convert it from (H,W,C) to\n",
    "        # (N,C,H,W).\n",
    "        resized_image = cv2.resize(src=image, dsize=(input_width, input_height))\n",
    "        input_image_original = np.expand_dims(resized_image.transpose(2, 0, 1), axis=0)\n",
    "\n",
    "        # Resize and reshape the image to the target shape with bicubic\n",
    "        # interpolation.\n",
    "        bicubic_image = cv2.resize(\n",
    "            src=image, dsize=(target_width, target_height), interpolation=cv2.INTER_CUBIC\n",
    "        )\n",
    "        input_image_bicubic = np.expand_dims(bicubic_image.transpose(2, 0, 1), axis=0)\n",
    "\n",
    "        # Do inference.\n",
    "        inference_start_time = time.perf_counter()\n",
    "        result = compiled_model(\n",
    "            {\n",
    "                original_image_key.any_name: input_image_original,\n",
    "                bicubic_image_key.any_name: input_image_bicubic,\n",
    "            }\n",
    "        )[output_key]\n",
    "        inference_stop_time = time.perf_counter()\n",
    "        inference_duration = inference_stop_time - inference_start_time\n",
    "        total_inference_duration += inference_duration\n",
    "\n",
    "        # Transform the inference result into an image.\n",
    "        result_frame = convert_result_to_image(result=result)\n",
    "\n",
    "        # Write the result image and the bicubic image to a video file.\n",
    "        superres_video.write(image=result_frame)\n",
    "        bicubic_video.write(image=bicubic_image)\n",
    "\n",
    "        stacked_frame = np.hstack((bicubic_image, result_frame))\n",
    "        comparison_video.write(image=stacked_frame)\n",
    "\n",
    "        frame_nr = frame_nr + 1\n",
    "\n",
    "        # Update the progress bar and the status message.\n",
    "        progress_bar.progress = frame_nr\n",
    "        progress_bar.update()\n",
    "        if frame_nr % 10 == 0 or frame_nr == total_frames:\n",
    "            clear_output(wait=True)\n",
    "            progress_bar.display()\n",
    "            display(\n",
    "                Pretty(\n",
    "                    f\"Processed frame {frame_nr}. Inference time: \"\n",
    "                    f\"{inference_duration:.2f} seconds \"\n",
    "                    f\"({1/inference_duration:.2f} FPS)\"\n",
    "                )\n",
    "            )\n",
    "\n",
    "\n",
    "except KeyboardInterrupt:\n",
    "    print(\"Processing interrupted.\")\n",
    "finally:\n",
    "    superres_video.release()\n",
    "    bicubic_video.release()\n",
    "    comparison_video.release()\n",
    "    end_time = time.perf_counter()\n",
    "    duration = end_time - start_time\n",
    "    print(f\"Video's saved to {comparison_video_path.parent} directory.\")\n",
    "    print(\n",
    "        f\"Processed {frame_nr} frames in {duration:.2f} seconds. Total FPS \"\n",
    "        f\"(including video processing): {frame_nr/duration:.2f}. \"\n",
    "        f\"Inference FPS: {frame_nr/total_inference_duration:.2f}.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61022f66",
   "metadata": {},
   "source": [
    "### Show Side-by-Side Video of Bicubic and Superresolution Version [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b32ba25d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not comparison_video_path.exists():\n",
    "    raise ValueError(\"The comparison video does not exist.\")\n",
    "else:\n",
    "    video_link = FileLink(comparison_video_path)\n",
    "    video_link.html_link_str = \"<a href='%s' download>%s</a>\"\n",
    "    display(\n",
    "        HTML(\n",
    "            f\"Showing side by side comparison. If you cannot see the video in \"\n",
    "            \"your browser, please click on the following link to download \"\n",
    "            f\"the video<br>{video_link._repr_html_()}\"\n",
    "        )\n",
    "    )\n",
    "    display(Video(comparison_video_path, width=800, embed=True))"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "bb0b397daf458ed78ef5a7e21732498aa92824cb15d3098f5341da903a887e15"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}