notebooks/004-hello-detection/004-hello-detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Hello World Text Detection\n",
    "\n",
    "A very basic introduction to OpenVINO that shows how to do text detection on a given IR model.\n",
    "\n",
    "We use the [horizontal-text-detection-0001](https://docs.openvinotoolkit.org/latest/omz_models_model_horizontal_text_detection_0001.html) model from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). It detects texts in images and returns blob of data in shape of [100, 5]. For each detection description has format [x_min, y_min, x_max, y_max, conf].\n"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Imports"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from openvino.inference_engine import IECore\n",
    "from os.path import isfile"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load the network"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "ie = IECore()\n",
    "model_folder = \"model\"\n",
    "model_name = \"horizontal-text-detection-0001\"\n",
    "model_extensions = (\"bin\", \"xml\")\n",
    "\n",
    "for extension in model_extensions:\n",
    "    if not isfile(f'{model_folder}/{model_name}.{extension}'):\n",
    "        raise FileNotFoundError(f\"Missing model file! Please download missing file: {model_name}.{extension}\")\n",
    "\n",
    "net = ie.read_network(\n",
    "    model=\"model/horizontal-text-detection-0001.xml\"\n",
    ")\n",
    "exec_net = ie.load_network(net, \"CPU\")\n",
    "\n",
    "output_layer_ir = next(iter(exec_net.outputs))\n",
    "input_layer_ir = next(iter(exec_net.input_info))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load an Image"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Text detection models expects image in BGR format\n",
    "image = cv2.imread(\"data/intel_rnb.jpg\")\n",
    "\n",
    "# N,C,H,W = batch size, number of channels, height, width\n",
    "N, C, H, W = net.input_info[input_layer_ir].tensor_desc.dims\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(\n",
    "    resized_image.transpose(2, 0, 1), 0\n",
    ")  \n",
    "\n",
    "plt.imshow(image)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Do Inference"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "result = exec_net.infer(inputs={input_layer_ir: input_image})\n",
    "\n",
    "# Extract list of boxes from results\n",
    "boxes = result['boxes']\n",
    "\n",
    "# Remove zero only boxes\n",
    "boxes = boxes[~np.all(boxes==0, axis=1)]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Visualize data"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# For each detection, the description has the format: [x_min, y_min, x_max, y_max, conf]\n",
    "# Image passed here is in BGR format with changed width and height. To display it in colors expected by matplotlib we use cvtColor funtion\n",
    "\n",
    "def convert_result_to_image(bgr_image, resized_image, boxes, threshold=0.3, conf_labels=True): \n",
    "    # Helper function to multiply shape by ratio\n",
    "    def multiply_by_ratio(ratio_x, ratio_y, box):\n",
    "        return [max(shape * ratio_y, 10) if idx % 2 else shape * ratio_x for idx, shape in enumerate(box[:-1])]\n",
    "\n",
    "    # Define colors for boxes and descriptions\n",
    "    colors = {'red': (255, 0, 0), 'green': (0, 255, 0)} \n",
    " \n",
    "    # Fetch image shapes to calculate ratio\n",
    "    (real_y, real_x), (resized_y, resized_x) = image.shape[:2], resized_image.shape[:2]\n",
    "    ratio_x, ratio_y = real_x/resized_x, real_y/resized_y\n",
    "\n",
    "    # Convert base image from bgr to rgb format\n",
    "    rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB) \n",
    "\n",
    "    # Iterate through non-zero boxes\n",
    "    for box in boxes: \n",
    "        # Pick confidence factor from last place in array\n",
    "        conf = box[-1]\n",
    "        if conf > threshold: \n",
    "            # Convert float to int and multiply position of each box by x and y ratio\n",
    "            (x_min, y_min, x_max, y_max) = map(int, multiply_by_ratio(ratio_x, ratio_y, box)) \n",
    "\n",
    "            # Draw box based on position, parameters in rectangle function are: image, start_point, end_point, color, thickness \n",
    "            rgb_image = cv2.rectangle( \n",
    "                rgb_image, \n",
    "                (x_min, y_min), \n",
    "                (x_max, y_max), \n",
    "                colors['green'], \n",
    "                3\n",
    "            ) \n",
    "\n",
    "            # Add text to image based on position and confidence, parameters in putText function are: image, text, bottomleft_corner_textfield, font, font_scale, color, thickness, line_type \n",
    "            if conf_labels:\n",
    "                rgb_image = cv2.putText( \n",
    "                    rgb_image, \n",
    "                    f\"{conf:.2f}\", \n",
    "                    (x_min, y_min - 10), \n",
    "                    cv2.FONT_HERSHEY_SIMPLEX, \n",
    "                    0.8, \n",
    "                    colors['red'], \n",
    "                    1, \n",
    "                    cv2.LINE_AA\n",
    "                ) \n",
    "            \n",
    "    return rgb_image"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "plt.figure(figsize=(10,6))\n",
    "plt.axis('off')\n",
    "plt.imshow(convert_result_to_image(image, resized_image, boxes, conf_labels=False))"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}