notebooks/125-torchvision-zoo-to-openvino/125-lraspp-segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Semantic segmentation with LRASPP MobileNet v3 and OpenVINO\n",
    "The [`torchvision.models`](https://pytorch.org/vision/stable/models.html) subpackage contains definitions of models for addressing different tasks, including: image classification, pixelwise semantic segmentation, object detection, instance segmentation, person keypoint detection, video classification, and optical flow. Throughout this notebook we will show how to use one of them.\n",
    "The LRASPP model is based on the [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244) paper. According to the paper, Searching for MobileNetV3, LR-ASPP or Lite Reduced Atrous Spatial Pyramid Pooling has a lightweight and efficient segmentation decoder architecture. he model is pre-trained on the [MS COCO](https://cocodataset.org/#home) dataset. Instead of training on all 80 classes, the segmentation model has been trained on 20 classes from the [PASCAL VOC](http://host.robots.ox.ac.uk/pascal/VOC/) dataset:\n",
    "***background*, *aeroplane*, *bicycle*, *bird*, *boat*, *bottle*, *bus*, *car*, *cat*, *chair*, *cow*, *dining table*, *dog*, *horse*, *motorbike*, *person*, *potted plant*, *sheep*, *sofa*, *train*, *tv monitor***\n",
    "\n",
    "More information about the model is available in the [torchvision documentation](https://pytorch.org/vision/main/models/lraspp.html)\n",
    "\n",
    "### Table of content:\n",
    "- [Prerequisites](#Prerequisites)\n",
    "- [Get a test image](#Get-a-test-image)\n",
    "- [Download and prepare model](#Download-and-prepare-model)\n",
    "- [Define a preprocessing and prepare an input data](#Define-a-preprocessing-and-prepare-an-input-data)\n",
    "- [Run an inference on the PyTorch model](#Run-an-inference-on-the-PyTorch-model)\n",
    "- [Convert the original model to OpenVINO IR Format](#Convert-the-original-model-to-OpenVINO-IR-Format)\n",
    "- [Run an inference on the OpenVINO model](#Run-an-inference-on-the-OpenVINO-model)\n",
    "- [Show Results](#Show-results)\n",
    "- [Show results for the OpenVINO IR model](#Show-results-for-the-OpenVINO-IR-model)\n",
    "    \n",
    "    \n",
    "    "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "55a1a04d5eaf0bea"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Prerequisites"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3b593e3f8d32cfbb"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "%pip install -q \"torch==2.0.1\" \"torchvision==0.15.2\"\n",
    "%pip install -q matplotlib\n",
    "%pip install -q --pre \"openvino>2023.1.0\""
   ],
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "id": "2515d4f6502681dc"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff2235b9bd1ba39",
   "metadata": {
    "collapsed": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import openvino as ov\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Get a test image\n",
    "First of all lets get a test image from an open dataset. "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "972f2c2b61244b97"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "\n",
    "from torchvision.io import read_image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "img_path = 'cats_image.jpeg'\n",
    "urllib.request.urlretrieve(\n",
    "    url='https://huggingface.co/datasets/huggingface/cats-image/resolve/main/cats_image.jpeg',\n",
    "    filename=img_path\n",
    ")\n",
    "image = read_image(img_path)\n",
    "display(transforms.ToPILImage()(image))"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "cbea76f60b13916e"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Download and prepare a model\n",
    "Define width and height of the image that will be used by the network during inference. According to the input transforms function, the model is pre-trained on images with a height of 480 and width of 640."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5ca420fd1cb14d78"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d4bd4bda02d787d",
   "metadata": {
    "collapsed": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "IMAGE_WIDTH = 640\n",
    "IMAGE_HEIGHT = 480"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Torchvision provides a mechanism of [listing and retrieving available models](https://pytorch.org/vision/stable/models.html#listing-and-retrieving-available-models). "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4e0f326d23440076"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "\n",
    "# List available models\n",
    "all_models = models.list_models()\n",
    "# List of models by type\n",
    "segmentation_models = models.list_models(module=models.segmentation)\n",
    "\n",
    "print(segmentation_models)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "12f6376d4abb7e61"
  },
  {
   "cell_type": "markdown",
   "source": [
    "We will use `lraspp_mobilenet_v3_large`. You can get a model by name using `models.get_model(\"lraspp_mobilenet_v3_large\", weights='DEFAULT')` or call a [corresponding function](https://pytorch.org/vision/stable/models/lraspp.html) directly. We will use `torchvision.models.segmentation.lraspp_mobilenet_v3_large`. You can directly pass pre-trained model weights to the model initialization function using weights enum LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1. It is a default weights. To get all available weights for the model you can call `weights_enum = models.get_model_weights(\"lraspp_mobilenet_v3_large\")`, but there is only one for this model."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "79835c4d4992ccaf"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "342be705e3152644",
   "metadata": {
    "collapsed": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "weights = models.segmentation.LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1\n",
    "model = models.segmentation.lraspp_mobilenet_v3_large(weights=weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Define a preprocessing and prepare an input data\n",
    "You can use `torchvision.transforms` to make a preprocessing or use[preprocessing transforms from the model wight](https://pytorch.org/vision/stable/models.html#using-the-pre-trained-models)."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f4b8fd0b41ac9b2c"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "preprocess = models.segmentation.LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1.transforms()\n",
    "preprocess.resize_size = (IMAGE_HEIGHT, IMAGE_WIDTH)  # change to an image size\n",
    "\n",
    "input_data = preprocess(image)\n",
    "input_data = np.expand_dims(input_data, axis=0)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "d27f89fc36134e34"
  },
  {
   "cell_type": "markdown",
   "id": "d6298d1512d6b745",
   "metadata": {},
   "source": [
    "## Run an inference on the PyTorch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51f312be8df5b9c6",
   "metadata": {
    "collapsed": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    result_torch = model(torch.as_tensor(input_data).float())['out']"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Convert the original model to OpenVINO IR Format\n",
    "\n",
    "To convert the original model to OpenVINO IR with `FP16` precision, use model conversion API. The models are saved inside the current directory. For more information on how to convert models, see this [page](https://docs.openvino.ai/2023.0/openvino_docs_model_processing_introduction.html)."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "577e673162569b63"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "ov_model_xml_path = Path('models/ov_lraspp_model.xml')\n",
    "\n",
    "\n",
    "if not ov_model_xml_path.exists():\n",
    "    ov_model_xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH)\n",
    "    ov_model = ov.convert_model(model, example_input=dummy_input)\n",
    "    ov.save_model(ov_model, ov_model_xml_path)\n",
    "else:\n",
    "    print(f\"IR model {ov_model_xml_path} already exists.\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6617e45f1df952f1"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Run an inference on the OpenVINO model"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "46be2fdb568b766"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Select device from dropdown list for running inference using OpenVINO"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "d4a126aeffa159e7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "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"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "677f739a6cc6f60a"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Run an inference"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7e269f98a3f07826"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e593740244684cd",
   "metadata": {
    "collapsed": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "compiled_model = core.compile_model(ov_model_xml_path, device_name=device.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f422b3698230c78d",
   "metadata": {
    "collapsed": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "res_ir = compiled_model(input_data)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Show results\n",
    "Confirm that the segmentation results look as expected by comparing model predictions on the OpenVINO IR and PyTorch models."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f3ef83667c402f6c"
  },
  {
   "cell_type": "markdown",
   "source": [
    "You can use [pytorch tutorial](https://pytorch.org/vision/0.12/auto_examples/plot_visualization_utils.html#sphx-glr-auto-examples-plot-visualization-utils-py) to visualize segmentation masks. Below is a simple example how to visualize the image with a `cat` mask for the PyTorch model."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f848b2893dd90725"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f95b993e440baf5c",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torchvision.transforms.functional as F\n",
    "\n",
    "\n",
    "plt.rcParams[\"savefig.bbox\"] = 'tight'\n",
    "\n",
    "\n",
    "def show(imgs):\n",
    "    if not isinstance(imgs, list):\n",
    "        imgs = [imgs]\n",
    "    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)\n",
    "    for i, img in enumerate(imgs):\n",
    "        img = img.detach()\n",
    "        img = F.to_pil_image(img)\n",
    "        axs[0, i].imshow(np.asarray(img))\n",
    "        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Prepare and display a cat mask."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4765c948799bec0b"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5215fdcc9b759cb",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sem_classes = [\n",
    "    '__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',\n",
    "    'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',\n",
    "    'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'\n",
    "]\n",
    "sem_class_to_idx = {cls: idx for (idx, cls) in enumerate(sem_classes)}\n",
    "\n",
    "normalized_mask = torch.nn.functional.softmax(result_torch, dim=1)\n",
    "\n",
    "cat_mask = normalized_mask[0, sem_class_to_idx['cat']]\n",
    "\n",
    "show(cat_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "The [draw_segmentation_masks()](https://pytorch.org/vision/0.12/generated/torchvision.utils.draw_segmentation_masks.html#torchvision.utils.draw_segmentation_masks)function can be used to plots those masks on top of the original image. This function expects the masks to be boolean masks, but our masks above contain probabilities in [0, 1]. To get boolean masks, we can do the following:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8464c91d92ebf5df"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class_dim = 1\n",
    "boolean_cat_mask = (normalized_mask.argmax(class_dim) == sem_class_to_idx['cat'])"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a9aabbd316208e32"
  },
  {
   "cell_type": "markdown",
   "source": [
    "And now we can plot a boolean mask on top of the original image."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1bb3b7c9190ea902"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from torchvision.utils import draw_segmentation_masks\n",
    "\n",
    "show(draw_segmentation_masks(image, masks=boolean_cat_mask, alpha=0.7, colors='yellow'))"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e6311727d89b5f7e"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Show results for the OpenVINO IR model"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "93435536ca0f2980"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "normalized_mask = torch.nn.functional.softmax(torch.from_numpy(res_ir), dim=1)\n",
    "boolean_cat_mask = (normalized_mask.argmax(class_dim) == sem_class_to_idx['cat'])\n",
    "show(draw_segmentation_masks(image, masks=boolean_cat_mask, alpha=0.7, colors='yellow'))"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "56b9409b96eecea1"
  }
 ],
 "metadata": {
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}