205-vision-background-removal

205-vision-background-removal.ipynb-standalone.ipynb

{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Background Removal From Images With U$^2$-Net and OpenVINO\n\nThis notebook demostrates background removal in images with U$^2$-Net and [OpenVINO](https://github.com/openvinotoolkit/openvino)\n\nFor more information about U$^2$-Net, including source code and test data, see their Github page at https://github.com/xuebinqin/U-2-Net and their research paper: [U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection](https://arxiv.org/pdf/2005.09007.pdf)"}, {"id": "5e106969", "cell_type": "markdown", "source": "## Preparation\n\nInstall the requirements and download the files that are necessary for running this notebook.\n**NOTE:** installation may take a while.", "metadata": {}}, {"id": "efd68617", "cell_type": "code", "metadata": {}, "execution_count": null, "source": "# Install or upgrade required Python packages. Install specific ", "outputs": []}, {"id": "16c792d3", "cell_type": "code", "metadata": {}, "execution_count": null, "source": "# Download image and model files\nimport os\nimport pip\nimport urllib.parse\nimport urllib.request\nfrom pathlib import Path\n\nurls = ['https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/205-vision-background-removal-with-output.md', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/205-vision-background-removal-with-output.html', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/205-vision-background-removal-with-output_files/205-vision-background-removal-with-output_24_0.png', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/205-vision-background-removal-with-output_files/205-vision-background-removal-with-output_22_0.png', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/model/LICENSE', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/model/__init__.py', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/model/u2net.py', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/model/__pycache__/__init__.cpython-37.pyc', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/model/__pycache__/u2net.cpython-37.pyc', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/media/coco_cross.png', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/media/paris.jpg', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/media/coco_hollywood.jpg', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/media/coco_hollywood-wall.jpg', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/media/wall.jpg', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/saved_models/u2net_lite/u2net_lite.xml', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/saved_models/u2net_lite/u2net_lite.mapping', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/saved_models/u2net_lite/u2net_lite.onnx', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/saved_models/u2net_lite/u2net_lite.pth', 'https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/notebooks/205-vision-background-removal/saved_models/u2net_lite/u2net_lite.bin']\n\nfor url in urls:\n    save_path = Path(url).relative_to(fr\"https:/raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/205-vision-background-removal\")\n    os.makedirs(save_path.parent, exist_ok=True)\n    safe_url = urllib.parse.quote(url, safe=\":/\")\n\n    urllib.request.urlretrieve(safe_url, save_path.as_posix())", "outputs": []}, {"cell_type": "markdown", "metadata": {"tags": ["hide"]}, "source": "The PyTorch U$^2$-Net model is converted to ONNX and loaded with OpenVINO. The model source is 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 IR format, check out the [PyTorch/ONNX](../102-pytorch-onnx-to-openvino) notebook."}, {"cell_type": "markdown", "metadata": {"tags": ["hide"]}, "source": "## Prepare"}, {"cell_type": "markdown", "metadata": {"id": "QB4Yo-rGGLmV", "tags": ["hide"]}, "source": "### Import the PyTorch Library and U2NET"}, {"cell_type": "code", "execution_count": null, "metadata": {"id": "2ynWRum4iiTz"}, "outputs": [], "source": "import os\nimport sys\nimport time\nfrom collections import namedtuple\nfrom pathlib import Path\n\nimport cv2\nimport matplotlib.pyplot as plt\nimport mo_onnx\nimport numpy as np\nimport torch\nfrom IPython.display import HTML, FileLink, display\nfrom model.u2net import U2NET, U2NETP\nfrom openvino.inference_engine import IECore"}, {"cell_type": "markdown", "metadata": {"tags": ["hide"]}, "source": "### Settings\n\nThis notebook supports 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, "metadata": {}, "outputs": [], "source": "IMAGE_DIR = \"media\"\nmodel_config = namedtuple(\"ModelConfig\", [\"name\", \"url\", \"model\", \"model_args\"])\n\nu2net_lite = model_config(\n    \"u2net_lite\",\n    \"https://drive.google.com/uc?id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy\",\n    U2NETP,\n    ()\n)\nu2net = model_config(\n    \"u2net\",\n    \"https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ\",\n    U2NET,\n    (3, 1)\n)\nu2net_human_seg = model_config(\n    \"u2net_human_seg\",\n    \"https://drive.google.com/uc?id=1-Yg0cxgrNhHP-016FPdp902BR-kSsA4P\",\n    U2NET,\n    (3, 1)\n)\n\n# Set u2net_model to one of the three configurations listed above\nu2net_model = u2net_lite"}, {"cell_type": "code", "execution_count": null, "metadata": {"tags": ["hide_output", "hide"]}, "outputs": [], "source": "# The filenames of the downloaded and converted models\nMODEL_DIR = \"saved_models\"\nmodel_path = (\n    Path(MODEL_DIR)\n    / u2net_model.name\n    / Path(u2net_model.name).with_suffix(\".pth\")\n)\nonnx_path = model_path.with_suffix(\".onnx\")\nir_path = model_path.with_suffix(\".xml\")"}, {"cell_type": "markdown", "metadata": {"id": "u5xKw0hR0jq6", "tags": ["hide"]}, "source": "### Load the U2NET Model\n\nThe U$^2$-Net human segmentation model weights are stored on Google Drive. They will be downloaded if they have not been downloaded yet. The next cell loads the U2NET model and the pretrained weights."}, {"cell_type": "code", "execution_count": null, "metadata": {"tags": ["hide"]}, "outputs": [], "source": "if not model_path.exists():\n    import gdown\n\n    os.makedirs(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(u2net_model.url, output=model_file)\n        print(f\"Model weights have been downloaded to {model_path}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {"tags": ["hide"]}, "outputs": [], "source": "# Load the model\nnet = u2net_model.model(*u2net_model.model_args)\nnet.eval()\n\n# Load the weights\nprint(f\"Loading model weights from: '{model_path}'\")\nnet.load_state_dict(torch.load(model_path, map_location=\"cpu\"))\n\n# Save the model if it doesn't exist yet\nif not model_path.exists():\n    print(\"\\nSaving the model\")\n    torch.save(net.state_dict(), str(model_path))\n    print(f\"Model saved at {model_path}\")"}, {"cell_type": "markdown", "metadata": {"id": "Rhc_7EObUypw", "tags": ["hide"]}, "source": "## Convert PyTorch U$^2$-Net model to ONNX and IR\n\n### Convert PyTorch model to ONNX\n\nThe output for this cell will show some warnings. These are most likely harmless. Conversion succeeded if the last line of the output says `ONNX model exported to [filename].onnx.` "}, {"cell_type": "code", "execution_count": null, "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(net, dummy_input, onnx_path, opset_version=11)\n    print(f\"ONNX model exported to {onnx_path}.\")\nelse:\n    print(f\"ONNX model {onnx_path} already exists.\")"}, {"cell_type": "markdown", "metadata": {"id": "6JSoEIk60uxV", "tags": ["hide"]}, "source": "### Convert ONNX model to OpenVINO IR Format\n\nCall the OpenVINO Model Optimizer tool to convert the ONNX model to OpenVINO IR format, with FP16 precision. The models are saved to the current directory. We 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\nSee the [Model Optimizer Developer Guide](https://docs.openvinotoolkit.org/latest/openvino_docs_MO_DG_Deep_Learning_Model_Optimizer_DevGuide.html) for more information about Model Optimizer."}, {"cell_type": "markdown", "metadata": {"tags": ["hide"]}, "source": "Call the OpenVINO Model Optimizer tool to convert the ONNX model to OpenVINO IR, with FP16 precision. Executing this command may take a while. There may be some errors or warnings in the output. Model Optimization was succesful if the last lines of the output include `[ SUCCESS ] Generated IR version 10 model.\n`"}, {"cell_type": "code", "execution_count": null, "metadata": {"tags": ["hide"]}, "outputs": [], "source": "# Get the path to the Model Optimizer script\nmo_path = str(Path(mo_onnx.__file__))\n\n# Construct the command for Model Optimizer\nmo_command = f\"\"\"\"{sys.executable}\"\n                 \"{mo_path}\"\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                 --data_type FP16\n                 --output_dir \"{model_path.parent}\"\n                 \"\"\"\nmo_command = \" \".join(mo_command.split())\nprint(\"Model Optimizer command to convert the ONNX model to OpenVINO:\")\nprint(mo_command)"}, {"cell_type": "code", "execution_count": null, "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_command\nelse:\n    print(f\"IR model {ir_path} already exists.\")"}, {"cell_type": "markdown", "metadata": {"id": "JyD5EKka34Wd", "tags": ["hide"]}, "source": "## Load and Pre-process Input Image\n\nThe U2NET IR model expects images in RGB format. OpenCV reads images in BGR. We convert the images to RGB, resize them to 512 by 512 and transpose the dimensions to format that is expected by the IR model"}, {"cell_type": "code", "execution_count": null, "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\"\nimage = cv2.cvtColor(\n    cv2.imread(str(IMAGE_PATH)),\n    cv2.COLOR_BGR2RGB,\n)\n\nresized_image = cv2.resize(image, (512, 512))\n# Convert the image shape to shape and data type expected by network\n# for OpenVINO IR model: (1, 3, 512, 512)\ninput_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)"}, {"cell_type": "markdown", "metadata": {"id": "FnEiEbNq4Csh", "tags": ["hide"]}, "source": "## Do inference on IR model\n\nLoad the IR model to Inference Engine and do inference"}, {"cell_type": "code", "execution_count": null, "metadata": {"id": "otfT6EDk03KV", "tags": ["hide"]}, "outputs": [], "source": "# Load network to Inference Engine\nie = IECore()\nnet_ir = ie.read_network(model=ir_path)\nexec_net_ir = ie.load_network(network=net_ir, device_name=\"CPU\")\n# Get names of input and output layers\ninput_layer_ir = next(iter(exec_net_ir.input_info))\noutput_layer_ir = next(iter(exec_net_ir.outputs))\n\n# Run the Inference on the Input image...\nstart_time = time.perf_counter()\nres_ir = exec_net_ir.infer(inputs={input_layer_ir: input_image})\nres_ir = res_ir[output_layer_ir]\nend_time = time.perf_counter()\nprint(\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", "metadata": {"tags": ["hide"]}, "source": "## Visualize results\n\nShow the original image, the segmentation result, and the original image with the background removed. "}, {"cell_type": "code", "execution_count": null, "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# Network result has shape (1,1,512,512), np.squeeze converts this to (512, 512)\nresized_result = np.rint(\n    cv2.resize(np.squeeze(res_ir), (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)\nbg_removed_result = image.copy()\nbg_removed_result[resized_result == 0] = 255\n\nfig, ax = plt.subplots(1, 3, figsize=(20, 7))\nax[0].imshow(image)\nax[1].imshow(resized_result, cmap=\"gray\")\nax[2].imshow(bg_removed_result)\nfor a in ax:\n    a.axis(\"off\")"}, {"cell_type": "markdown", "metadata": {"tags": ["hide"]}, "source": "### Add a background image\n\nIn 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 image `background_image`\n- Resize this image to the same size as the original image\n- In the `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 the `bg_removed_result` from the previous step - the part of the original image that only contains foreground pixels - to the `background_image`."}, {"cell_type": "code", "execution_count": null, "metadata": {"tags": ["hide"]}, "outputs": [], "source": "BACKGROUND_FILE = \"media/wall.jpg\"\n\nbackground_image = cv2.cvtColor(cv2.imread(BACKGROUND_FILE), cv2.COLOR_BGR2RGB)\nbackground_image = cv2.resize(\n    background_image, (image.shape[1], image.shape[0])\n)\n\n# Set all the foreground pixels from the result to 0\n# in the background image and add the background-removed image\nbackground_image[resized_result == 1] = 0\nnew_image = background_image + bg_removed_result\n\n# Save the generated image\nnew_image_path = (\n    IMAGE_PATH.parent / f\"{IMAGE_PATH.stem}-{Path(BACKGROUND_FILE).stem}.jpg\"\n)\ncv2.imwrite(str(new_image_path), cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR))\n\n# Display the original image and the image with the new background side by side\nfig, ax = plt.subplots(1, 2, figsize=(18, 7))\nax[0].imshow(image)\nax[1].imshow(new_image)\nfor a in ax:\n    a.axis(\"off\")\nplt.show()\n\n# Create a link to download the image\nimage_link = FileLink(new_image_path)\nimage_link.html_link_str = \"<a href='%s' download>%s</a>\"\ndisplay(\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", "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.openvinotoolkit.org/latest/openvino_docs_IE_DG_ONNX_Support.html\n* Model Optimizer Documentation: https://docs.openvinotoolkit.org/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": {"colab": {"collapsed_sections": [], "name": "OpenVINO 2021.3 PIP installer - PyTorch Image Segmentation.ipynb", "provenance": [], "toc_visible": true}, "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.6.8"}}, "nbformat": 4, "nbformat_minor": 4}

apt.txt

libpython3.7-dev