notebooks/205-vision-background-removal/205-vision-background-removal.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93aa099e",
   "metadata": {},
   "source": [
    "# Image Background Removal with U^2-Net and OpenVINO™\n",
    "\n",
    "This notebook demonstrates background removal in images using U$^2$-Net and OpenVINO.\n",
    "\n",
    "For more information about U$^2$-Net, including source code and test data, see the [Github page](https://github.com/xuebinqin/U-2-Net) and the research paper: [U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection](https://arxiv.org/pdf/2005.09007.pdf).\n",
    "\n",
    "The PyTorch U$^2$-Net model is converted to ONNX and loaded with OpenVINO. The model source is [here](https://github.com/xuebinqin/U-2-Net). For a more detailed overview of loading PyTorch models in OpenVINO, including how to load an ONNX model in OpenVINO directly, without converting to OpenVINO IR format, see the [PyTorch/ONNX](../102-pytorch-onnx-to-openvino/102-pytorch-onnx-to-openvino.ipynb) notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e57b48f0",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Prepare\n",
    "\n",
    "### Import the PyTorch Library and U$^2$-Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8124fd11",
   "metadata": {
    "id": "2ynWRum4iiTz"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from collections import namedtuple\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from IPython.display import HTML, FileLink, display\n",
    "from model.u2net import U2NET, U2NETP\n",
    "from openvino.runtime import Core"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bae3e45",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Settings\n",
    "\n",
    "This tutorial supports using the original U$^2$-Net salient object detection model, as well as the smaller U2NETP version. Two sets of weights are supported for the original model: salient object detection and human segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e9c87b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_DIR = \"../data/image/\"\n",
    "model_config = namedtuple(\"ModelConfig\", [\"name\", \"url\", \"model\", \"model_args\"])\n",
    "\n",
    "u2net_lite = model_config(\n",
    "    name=\"u2net_lite\",\n",
    "    url=\"https://drive.google.com/uc?id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy\",\n",
    "    model=U2NETP,\n",
    "    model_args=(),\n",
    ")\n",
    "u2net = model_config(\n",
    "    name=\"u2net\",\n",
    "    url=\"https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ\",\n",
    "    model=U2NET,\n",
    "    model_args=(3, 1),\n",
    ")\n",
    "u2net_human_seg = model_config(\n",
    "    name=\"u2net_human_seg\",\n",
    "    url=\"https://drive.google.com/uc?id=1-Yg0cxgrNhHP-016FPdp902BR-kSsA4P\",\n",
    "    model=U2NET,\n",
    "    model_args=(3, 1),\n",
    ")\n",
    "\n",
    "# Set u2net_model to one of the three configurations listed above.\n",
    "u2net_model = u2net_lite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7799f33c",
   "metadata": {
    "tags": [
     "hide_output",
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "# The filenames of the downloaded and converted models.\n",
    "MODEL_DIR = \"model\"\n",
    "model_path = Path(MODEL_DIR) / u2net_model.name / Path(u2net_model.name).with_suffix(\".pth\")\n",
    "onnx_path = model_path.with_suffix(\".onnx\")\n",
    "ir_path = model_path.with_suffix(\".xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8522caea",
   "metadata": {
    "id": "u5xKw0hR0jq6",
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Load the U$^2$-Net Model\n",
    "\n",
    "The U$^2$-Net human segmentation model weights are stored on Google Drive. They will be downloaded if they are not present yet. The next cell loads the model and the pre-trained weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eae9467",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if not model_path.exists():\n",
    "    import gdown\n",
    "\n",
    "    os.makedirs(name=model_path.parent, exist_ok=True)\n",
    "    print(\"Start downloading model weights file... \")\n",
    "    with open(model_path, \"wb\") as model_file:\n",
    "        gdown.download(url=u2net_model.url, output=model_file)\n",
    "        print(f\"Model weights have been downloaded to {model_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0fb802",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "# Load the model.\n",
    "net = u2net_model.model(*u2net_model.model_args)\n",
    "net.eval()\n",
    "\n",
    "# Load the weights.\n",
    "print(f\"Loading model weights from: '{model_path}'\")\n",
    "net.load_state_dict(state_dict=torch.load(model_path, map_location=\"cpu\"))\n",
    "\n",
    "# Save the model if it does not exist yet.\n",
    "if not model_path.exists():\n",
    "    print(\"\\nSaving the model\")\n",
    "    torch.save(obj=net.state_dict(), f=str(model_path))\n",
    "    print(f\"Model saved at {model_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d55f18",
   "metadata": {
    "id": "Rhc_7EObUypw",
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Convert PyTorch U$^2$-Net model to ONNX and OpenVINO IR\n",
    "\n",
    "### Convert PyTorch model to ONNX\n",
    "\n",
    "The output for this cell will show some warnings. These are most likely harmless. When the conversion succeeds, the last line of the output will read `ONNX model exported to [filename].onnx.` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ef50a4c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ipQWpbgQUxoo",
    "outputId": "bbc1734a-c2a2-4261-ed45-264b9e3edd00",
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if not onnx_path.exists():\n",
    "    dummy_input = torch.randn(1, 3, 512, 512)\n",
    "    torch.onnx.export(model=net, args=dummy_input, f=onnx_path, opset_version=11)\n",
    "    print(f\"ONNX model exported to {onnx_path}.\")\n",
    "else:\n",
    "    print(f\"ONNX model {onnx_path} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "960b0c2f",
   "metadata": {
    "id": "6JSoEIk60uxV",
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Convert ONNX model to OpenVINO IR Format\n",
    "\n",
    "Use Model Optimizer to convert the ONNX model to OpenVINO IR format, with `FP16` precision. The models are saved to the current directory. Then, add the mean values to the model and scale the output with the standard deviation with `--scale_values`. With these options, it is not necessary to normalize input data before propagating it through the network. The mean and standard deviation values can be found in the [dataloader](https://github.com/xuebinqin/U-2-Net/blob/master/data_loader.py) file in the [U^2-Net repository](https://github.com/xuebinqin/U-2-Net/) and multiplied by 255 to support images with pixel values from 0-255.\n",
    "\n",
    "For more information, refer to the [Model Optimizer Developer Guide](https://docs.openvino.ai/latest/openvino_docs_MO_DG_Deep_Learning_Model_Optimizer_DevGuide.html).\n",
    "\n",
    "Executing the following command may take a while. There may be some errors or warnings in the output. When the model optimization is successful, the last lines of the output will include `[ SUCCESS ] Generated IR version 11 model.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "340fb5b5",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "# Construct the command for Model Optimizer.\n",
    "# Set log_level to CRITICAL to suppress warnings that can be ignored for this demo.\n",
    "mo_command = f\"\"\"mo\n",
    "                 --input_model \"{onnx_path}\"\n",
    "                 --input_shape \"[1,3, 512, 512]\"\n",
    "                 --mean_values=\"[123.675, 116.28 , 103.53]\"\n",
    "                 --scale_values=\"[58.395, 57.12 , 57.375]\"\n",
    "                 --compress_to_fp16\n",
    "                 --output_dir \"{model_path.parent}\"\n",
    "                 --log_level \"CRITICAL\"\n",
    "                 \"\"\"\n",
    "mo_command = \" \".join(mo_command.split())\n",
    "print(\"Model Optimizer command to convert the ONNX model to OpenVINO:\")\n",
    "print(mo_command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b28b55",
   "metadata": {
    "id": "6YUwrq7QWSzw",
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if not ir_path.exists():\n",
    "    print(\"Exporting ONNX model to IR... This may take a few minutes.\")\n",
    "    mo_result = %sx $mo_command\n",
    "    print(\"\\n\".join(mo_result))\n",
    "else:\n",
    "    print(f\"IR model {ir_path} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43ed7603",
   "metadata": {
    "id": "JyD5EKka34Wd",
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Load and Pre-Process Input Image\n",
    "\n",
    "While OpenCV reads images in `BGR` format, the OpenVINO IR model expects images in `RGB`. Therefore, convert the images to `RGB`, resize them to `512 x 512` and transpose the dimensions to the format that is expected by the OpenVINO IR model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a86e7fd",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DGFW5VXL3x9G",
    "outputId": "300eacff-c6de-4eb5-e99a-8def5260da1a",
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "IMAGE_PATH = Path(IMAGE_DIR) / \"coco_hollywood.jpg\"\n",
    "image = cv2.cvtColor(\n",
    "    src=cv2.imread(filename=str(IMAGE_PATH)),\n",
    "    code=cv2.COLOR_BGR2RGB,\n",
    ")\n",
    "\n",
    "resized_image = cv2.resize(src=image, dsize=(512, 512))\n",
    "# Convert the image shape to a shape and a data type expected by the network\n",
    "# for OpenVINO IR model: (1, 3, 512, 512).\n",
    "input_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae10f73a",
   "metadata": {
    "id": "FnEiEbNq4Csh",
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Do Inference on OpenVINO IR Model\n",
    "\n",
    "Load the OpenVINO IR model to OpenVINO Runtime and do inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6cff378",
   "metadata": {
    "id": "otfT6EDk03KV",
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "# Load the network to OpenVINO Runtime.\n",
    "ie = Core()\n",
    "model_ir = ie.read_model(model=ir_path)\n",
    "compiled_model_ir = ie.compile_model(model=model_ir, device_name=\"CPU\")\n",
    "# Get the names of input and output layers.\n",
    "input_layer_ir = compiled_model_ir.input(0)\n",
    "output_layer_ir = compiled_model_ir.output(0)\n",
    "\n",
    "# Do inference on the input image.\n",
    "start_time = time.perf_counter()\n",
    "result = compiled_model_ir([input_image])[output_layer_ir]\n",
    "end_time = time.perf_counter()\n",
    "print(\n",
    "    f\"Inference finished. Inference time: {end_time-start_time:.3f} seconds, \"\n",
    "    f\"FPS: {1/(end_time-start_time):.2f}.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4b7f850",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Visualize Results\n",
    "\n",
    "Show the original image, the segmentation result, and the original image with the background removed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f22737",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "# Resize the network result to the image shape and round the values\n",
    "# to 0 (background) and 1 (foreground).\n",
    "# The network result has (1,1,512,512) shape. The `np.squeeze` function converts this to (512, 512).\n",
    "resized_result = np.rint(\n",
    "    cv2.resize(src=np.squeeze(result), dsize=(image.shape[1], image.shape[0]))\n",
    ").astype(np.uint8)\n",
    "\n",
    "# Create a copy of the image and set all background values to 255 (white).\n",
    "bg_removed_result = image.copy()\n",
    "bg_removed_result[resized_result == 0] = 255\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(20, 7))\n",
    "ax[0].imshow(image)\n",
    "ax[1].imshow(resized_result, cmap=\"gray\")\n",
    "ax[2].imshow(bg_removed_result)\n",
    "for a in ax:\n",
    "    a.axis(\"off\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c29ee84",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Add a Background Image\n",
    "\n",
    "In the segmentation result, all foreground pixels have a value of 1, all background pixels a value of 0. Replace the background image as follows:\n",
    "\n",
    "- Load a new `background_image`.\n",
    "- Resize this image to the same size as the original image.\n",
    "- In `background_image`, set all the pixels where the resized segmentation result has a value of 1 - the foreground pixels in the original image - to 0.\n",
    "- Add `bg_removed_result` from the previous step - the part of the original image that only contains foreground pixels - to `background_image`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5628a0c",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "BACKGROUND_FILE = f\"{IMAGE_DIR}/wall.jpg\"\n",
    "OUTPUT_DIR = \"output\"\n",
    "\n",
    "os.makedirs(name=OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "background_image = cv2.cvtColor(src=cv2.imread(filename=BACKGROUND_FILE), code=cv2.COLOR_BGR2RGB)\n",
    "background_image = cv2.resize(src=background_image, dsize=(image.shape[1], image.shape[0]))\n",
    "\n",
    "# Set all the foreground pixels from the result to 0\n",
    "# in the background image and add the image with the background removed.\n",
    "background_image[resized_result == 1] = 0\n",
    "new_image = background_image + bg_removed_result\n",
    "\n",
    "# Save the generated image.\n",
    "new_image_path = Path(f\"{OUTPUT_DIR}/{IMAGE_PATH.stem}-{Path(BACKGROUND_FILE).stem}.jpg\")\n",
    "cv2.imwrite(filename=str(new_image_path), img=cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR))\n",
    "\n",
    "# Display the original image and the image with the new background side by side\n",
    "fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(18, 7))\n",
    "ax[0].imshow(image)\n",
    "ax[1].imshow(new_image)\n",
    "for a in ax:\n",
    "    a.axis(\"off\")\n",
    "plt.show()\n",
    "\n",
    "# Create a link to download the image.\n",
    "image_link = FileLink(new_image_path)\n",
    "image_link.html_link_str = \"<a href='%s' download>%s</a>\"\n",
    "display(\n",
    "    HTML(\n",
    "        f\"The generated image <code>{new_image_path.name}</code> is saved in \"\n",
    "        f\"the directory <code>{new_image_path.parent}</code>. You can also \"\n",
    "        \"download the image by clicking on this link: \"\n",
    "        f\"{image_link._repr_html_()}\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b11acf85",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## References\n",
    "\n",
    "* [PIP install openvino-dev](https://github.com/openvinotoolkit/openvino/blob/releases/2021/3/docs/install_guides/pypi-openvino-dev.md)\n",
    "* [OpenVINO ONNX support](https://docs.openvino.ai/2021.4/openvino_docs_IE_DG_ONNX_Support.html)\n",
    "* [Model Optimizer Documentation](https://docs.openvino.ai/latest/openvino_docs_MO_DG_prepare_model_convert_model_Converting_Model_General.html)\n",
    "* [U^2-Net](https://github.com/xuebinqin/U-2-Net)\n",
    "* U^2-Net research paper: [U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection](https://arxiv.org/pdf/2005.09007.pdf)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "ae617ccb002f72b3ab6d0069d721eac67ac2a969e83c083c4321cfcab0437cd1"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}