notebooks/003-hello-segmentation/003-hello-segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "67fab52a",
   "metadata": {},
   "source": [
    "# Hello Image Segmentation\n",
    "\n",
    "A very basic introduction to using segmentation models with OpenVINO™.\n",
    "\n",
    "In this tutorial, a pre-trained [road-segmentation-adas-0001](https://docs.openvino.ai/latest/omz_models_model_road_segmentation_adas_0001.html) model from the [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/) is used. ADAS stands for Advanced Driver Assistance Services. The model recognizes four classes: background, road, curb and mark."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2f2f808",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "485ef549",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sys\n",
    "from openvino.runtime import Core\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from notebook_utils import segmentation_map_to_image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a340210",
   "metadata": {},
   "source": [
    "## Load the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26a398a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "\n",
    "model = ie.read_model(model=\"model/road-segmentation-adas-0001.xml\")\n",
    "compiled_model = ie.compile_model(model=model, device_name=\"CPU\")\n",
    "\n",
    "input_layer_ir = compiled_model.input(0)\n",
    "output_layer_ir = compiled_model.output(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c18397f8",
   "metadata": {},
   "source": [
    "## Load an Image\n",
    "A sample image from the [Mapillary Vistas](https://www.mapillary.com/dataset/vistas) dataset is provided. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23134719",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The segmentation network expects images in BGR format.\n",
    "image = cv2.imread(\"../data/image/empty_road_mapillary.jpg\")\n",
    "\n",
    "rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "image_h, image_w, _ = image.shape\n",
    "\n",
    "# N,C,H,W = batch size, number of channels, height, width.\n",
    "N, C, H, W = input_layer_ir.shape\n",
    "\n",
    "# OpenCV resize expects the destination size as (width, height).\n",
    "resized_image = cv2.resize(image, (W, H))\n",
    "\n",
    "# Reshape to the network input shape.\n",
    "input_image = np.expand_dims(\n",
    "    resized_image.transpose(2, 0, 1), 0\n",
    ")  \n",
    "plt.imshow(rgb_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df0e7703",
   "metadata": {},
   "source": [
    "## Do Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d7c5a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the inference.\n",
    "result = compiled_model([input_image])[output_layer_ir]\n",
    "\n",
    "# Prepare data for visualization.\n",
    "segmentation_mask = np.argmax(result, axis=1)\n",
    "plt.imshow(segmentation_mask.transpose(1, 2, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9503a83b",
   "metadata": {},
   "source": [
    "## Prepare Data for Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb32243d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define colormap, each color represents a class.\n",
    "colormap = np.array([[68, 1, 84], [48, 103, 141], [53, 183, 120], [199, 216, 52]])\n",
    "\n",
    "# Define the transparency of the segmentation mask on the photo.\n",
    "alpha = 0.3\n",
    "\n",
    "# Use function from notebook_utils.py to transform mask to an RGB image.\n",
    "mask = segmentation_map_to_image(segmentation_mask, colormap)\n",
    "resized_mask = cv2.resize(mask, (image_w, image_h))\n",
    "\n",
    "# Create an image with mask.\n",
    "image_with_mask = cv2.addWeighted(resized_mask, alpha, rgb_image, 1 - alpha, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49d16e71",
   "metadata": {},
   "source": [
    "## Visualize data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce182b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define titles with images.\n",
    "data = {\"Base Photo\": rgb_image, \"Segmentation\": mask, \"Masked Photo\": image_with_mask}\n",
    "\n",
    "# Create a subplot to visualize images.\n",
    "fig, axs = plt.subplots(1, len(data.items()), figsize=(15, 10))\n",
    "\n",
    "# Fill the subplot.\n",
    "for ax, (name, image) in zip(axs, data.items()):\n",
    "    ax.axis('off')\n",
    "    ax.set_title(name)\n",
    "    ax.imshow(image)\n",
    "\n",
    "# Display an image.\n",
    "plt.show(fig)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "ae617ccb002f72b3ab6d0069d721eac67ac2a969e83c083c4321cfcab0437cd1"
  },
  "kernelspec": {
   "display_name": "openvino_env",
   "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}