{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9536ae60-8a7d-4b36-8caf-1e68136752cc",
   "metadata": {},
   "source": [
    "# Deblur Photos with DeblurGAN-v2 and OpenVINO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "049dd5be-3fa1-4539-b37b-93981c22c480",
   "metadata": {},
   "source": [
    "This tutorial demonstrates Single Image Motion Deblurring with DeblurGAN-v2 in OpenVINO, by first converting the [VITA-Group/DeblurGANv2](https://github.com/VITA-Group/DeblurGANv2) model to OpenVINO's Intermediate Representation (IR) format. Model information can be found [here](https://docs.openvino.ai/latest/omz_models_model_deblurgan_v2.html).\n",
    "\n",
    "### What is deblurring?\n",
    "\n",
    "Deblurring is the task of removing motion blurs that usually occur in photos shot with hand-held cameras when there are moving objects in the scene. Blurs not only reduce the human perception about the quality of the image, but also complicate computer vision analyses. \n",
    "\n",
    "Kupyn, O., Martyniuk, T., Wu, J., & Wang, Z. (2019). [Deblurgan-v2: Deblurring (orders-of-magnitude) faster and better.](https://openaccess.thecvf.com/content_ICCV_2019/html/Kupyn_DeblurGAN-v2_Deblurring_Orders-of-Magnitude_Faster_and_Better_ICCV_2019_paper.html) In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 8878-8887)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae78b058-b51f-4bac-97f6-8189d6cf99c8",
   "metadata": {},
   "source": [
    "## Preparations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a507de7-4d9f-4f48-90e2-43678b116362",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c506fef8-b0ea-4d69-84ed-4c908e11e536",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import shutil\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from IPython.display import Markdown, display\n",
    "from openvino.runtime import Core\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from notebook_utils import load_image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c5f62c-24f7-40f4-9f4b-8c1d21bb2e48",
   "metadata": {},
   "source": [
    "### Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "377103cf-2c93-4c73-9392-f3765ce9f5d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Device to use for inference. For example, \"CPU\", or \"GPU\"\n",
    "DEVICE = \"CPU\"\n",
    "\n",
    "# Directory to save the model in\n",
    "model_dir = Path(\"model\")\n",
    "model_dir.mkdir(exist_ok=True)\n",
    "\n",
    "# Output directory\n",
    "output_dir = Path(\"output\")\n",
    "output_dir.mkdir(exist_ok=True)\n",
    "\n",
    "# The name of the Open Model Zoo model\n",
    "model_name = \"deblurgan-v2\"\n",
    "model_xml_path = model_dir / Path(f\"{model_name}.xml\")\n",
    "\n",
    "precision = \"FP16\"\n",
    "base_model_dir = Path(\"~/open_model_zoo_models\").expanduser()\n",
    "omz_cache_dir = Path(\"~/open_model_zoo_cache\").expanduser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d49ef301-bedd-493e-9434-531f89c16b7c",
   "metadata": {},
   "source": [
    "### Download DeblurGAN-v2 Model from Open Model Zoo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef6cb03e-a4be-45e6-bac4-749f916eba31",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "download_command = (\n",
    "    f\"omz_downloader --name {model_name} --output_dir\"\n",
    "    f\" {base_model_dir} --cache_dir {omz_cache_dir}\"\n",
    ")\n",
    "display(Markdown(f\"Download command: `{download_command}`\"))\n",
    "display(Markdown(f\"Downloading {model_name}...\"))\n",
    "! $download_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "449e4325-8441-480d-b9fa-fb2e0562a3cd",
   "metadata": {},
   "source": [
    "### Convert DeblurGAN-v2 Model to OpenVINO IR format\n",
    "\n",
    "Model Conversion may take a while. Conversion succeeded if the last lines of the output include `[ SUCCESS ] Generated IR version 11 model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fc62ac5-f2ad-4357-911a-c9f59cbdae3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "convert_command = (\n",
    "    f\"omz_converter --name {model_name} --precisions {precision} \"\n",
    "    f\"--download_dir {base_model_dir} --output_dir {base_model_dir}\"\n",
    ")\n",
    "display(Markdown(f\"Convert command: `{convert_command}`\"))\n",
    "display(Markdown(f\"Converting {model_name}...\"))\n",
    "\n",
    "! $convert_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "531c9e67-93dc-4238-b8ae-1742de498b4a",
   "metadata": {},
   "source": [
    "### Copy Model\n",
    "\n",
    "Copy the DeblurGAN-v2 Model to the specified `model_dir` directory in the settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd5c51b-b63c-4ace-8e6a-9eccd5e02b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_info = %sx omz_info_dumper --name $model_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b979b368-dd15-4d64-8c48-e6d5378fb4a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_info_json = json.loads(model_info.get_nlstr())[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a15f4c33-d7fc-47a4-9e2b-9a1c94a51dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_downloaded_dir = (\n",
    "    base_model_dir / Path(model_info_json[\"subdirectory\"]) / Path(precision)\n",
    ")\n",
    "\n",
    "for model_file in model_downloaded_dir.iterdir():\n",
    "    try:\n",
    "        shutil.copyfile(model_file, model_dir / model_file.name)\n",
    "    except FileExistsError:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0a38ab9-fea0-434d-b855-113799b51011",
   "metadata": {},
   "source": [
    "## Load the Model\n",
    "\n",
    "Load and compile the DeblurGAN-v2 model in the Inference Engine with `ie.read_model` and compile it for the specified device with `ie.compile_model`. Get input and output keys and the expected input shape for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baea1ef6-08fa-48e6-bac9-56e57eede3c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "model = ie.read_model(model=model_xml_path)\n",
    "compiled_model = ie.compile_model(model=model, device_name=DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b260f8b-9034-4dcc-b894-472125bf542c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input_layer = compiled_model.input(0)\n",
    "model_output_layer = compiled_model.output(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86f845ef-1be3-4b22-9f75-23a5ec985017",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0bad1b4-c0ac-4ad3-aeda-2e2035a52c21",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_output_layer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00f1d1c5-64e8-432e-b90a-6dac37d643ee",
   "metadata": {},
   "source": [
    "## Deblur Image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0861896-7579-461c-80de-720833616b68",
   "metadata": {},
   "source": [
    "### Load, resize and reshape input image\n",
    "\n",
    "The input image is read using the default `load_image` function from notebooks.utils, resized to meet the network's expected input sizes, and reshaped to (N, C, H, W), where N=number of images in the batch, C=number of channels, H=height, and W=width."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1b4718a-8711-4bcd-a382-08e88c227d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image filename (local path or URL)\n",
    "filename = \"https://user-images.githubusercontent.com/41332813/166901955-5d813e4c-a895-4da2-a36c-a96b3f627ebb.png\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af40f278-06a7-4f70-9217-372097c250f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load input image\n",
    "image = load_image(filename)\n",
    "\n",
    "# In case immage is stored in RGBA format\n",
    "image = np.asarray(Image.fromarray(image).convert('RGB'))\n",
    "\n",
    "# N,C,H,W = batch size, number of channels, height, width\n",
    "N, C, H, W = model_input_layer.shape\n",
    "\n",
    "# Resize image to meet network expected input sizes\n",
    "resized_image = cv2.resize(image, (W, H))\n",
    "\n",
    "# Reshape to network input shape\n",
    "input_image = np.expand_dims(resized_image.transpose(2, 0, 1), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491a66bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "144d2a65-6c70-4c9b-ac51-6d0dff371ab0",
   "metadata": {},
   "source": [
    "### Do inference on input image\n",
    "\n",
    "Do the inference, convert the result to an image shape and resize it to the original image size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "932a5f8c-3159-413a-8ba7-ed63c7398116",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference\n",
    "result = compiled_model([input_image])[model_output_layer]\n",
    "\n",
    "# Convert result to an image shape\n",
    "result_image = result[0].transpose((1, 2, 0))\n",
    "\n",
    "# Resize to original image size\n",
    "resized_result_image = cv2.resize(result_image, image.shape[:2][::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e65705c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f583b254-841d-42eb-9545-17d7b2f23726",
   "metadata": {},
   "source": [
    "### Load groundtruth image\n",
    "\n",
    "Load the groundtruth image for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0599e4f-fc00-4d50-95c6-07bd37ba9312",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load grountruth image\n",
    "groundtruth_image = load_image(\"https://user-images.githubusercontent.com/41332813/166900611-83be8ae5-23ed-4426-b535-fad960795687.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d344c818-690c-43af-8b16-8a111b950b2e",
   "metadata": {},
   "source": [
    "### Display results\n",
    "\n",
    "Images are in BGR format, so they will be converted to the RGB format in order to be properly displayed by matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e80530b2-59c3-42de-81f8-5c5305709538",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create subplot(r,c) by providing the no. of rows (r),\n",
    "# number of columns (c) and figure size\n",
    "f, ax = plt.subplots(1, 3, figsize=(20, 15))\n",
    "\n",
    "# Use the created array and display the images horizontally.\n",
    "ax[0].set_title(\"Blurred\")\n",
    "ax[0].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "\n",
    "ax[1].set_title(\"DeblurGAN-v2\")\n",
    "ax[1].imshow(cv2.cvtColor(resized_result_image, cv2.COLOR_BGR2RGB))\n",
    "\n",
    "ax[2].set_title(\"Groundtruth\")\n",
    "ax[2].imshow(cv2.cvtColor(groundtruth_image, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a148c1-ca1b-4a22-a98f-953b0873ad53",
   "metadata": {},
   "source": [
    "### Save the deblurred image\n",
    "\n",
    "Save the output image of the DeblurGAN-v2 model in the `output_dir` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "085b6b04-f082-4ad0-9250-16bc54c3ed8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "savename = \"deblurred.png\"\n",
    "plt.imsave(\n",
    "    output_dir / Path(savename),\n",
    "    cv2.cvtColor(resized_result_image, cv2.COLOR_BGR2RGB),\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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}