notebooks/111-detection-quantization/111-detection-quantization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2aa95237-5d60-437b-b080-b7118cd639da",
   "metadata": {},
   "source": [
    "# Object Detection Quantization\n",
    "\n",
    "This tutorial shows how to quantize an object detection model, using [Post-Training Optimization Tool API](https://docs.openvino.ai/2021.4/pot_compression_api_README.html) in OpenVINO™. \n",
    "\n",
    "For demonstration purposes, a very small dataset of 10 images presenting people at the airport is used. The images have been resized from the original resolution of 1920x1080 to 960x540. For any real use cases, a representative dataset of about 300 images is recommended. The tutorial uses the [person-detection-retail-0013](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/intel/person-detection-retail-0013) model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18a66ffc-6cf1-4483-9ab0-94ef8376acad",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa10aa83-c62c-41c0-a28f-338022364646",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import sys\n",
    "import time\n",
    "from pathlib import Path\n",
    "from typing import Sequence, Tuple\n",
    "\n",
    "import addict\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchmetrics\n",
    "from compression.api import DataLoader, Metric\n",
    "from compression.engines.ie_engine import IEEngine\n",
    "from compression.graph import load_model, save_model\n",
    "from compression.graph.model_utils import compress_model_weights\n",
    "from compression.pipeline.initializer import create_pipeline\n",
    "from openvino.runtime import Core\n",
    "from yaspin import yaspin\n",
    "\n",
    "sys.path.append(\"../utils\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "033bed78-0daa-4501-b41b-5912d81e6719",
   "metadata": {},
   "source": [
    "### Download Model\n",
    "\n",
    "Download the model from Open Model Zoo if it is not already in the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55acb9b9-51f5-49f3-b796-6e68b8ba56fb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ir_path = Path(\"intel/person-detection-retail-0013/FP32/person-detection-retail-0013.xml\")\n",
    "\n",
    "if not ir_path.exists():\n",
    "    ! omz_downloader --name \"person-detection-retail-0013\" --precisions FP32"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f676a2f9-6e92-4374-af2d-20dcc3a966c8",
   "metadata": {},
   "source": [
    "### Load Model\n",
    "\n",
    "Load the OpenVINO IR model, and get information about network inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e4108c1-6c3c-4410-8ca2-11d6a2779583",
   "metadata": {},
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "model = ie.read_model(model=ir_path)\n",
    "compiled_model = ie.compile_model(model=model, device_name=\"CPU\")\n",
    "input_layer = compiled_model.input(0)\n",
    "output_layer = compiled_model.output(0)\n",
    "input_size = input_layer.shape\n",
    "_, _, input_height, input_width = input_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b3101a2-84a6-46ab-8fe0-66bd0c4d94c7",
   "metadata": {},
   "source": [
    "## Post-Training Optimization Tool (POT) Quantization\n",
    "\n",
    "The Post-Training Optimization Tool (POT) `compression` API defines base classes for `Metric` and `DataLoader`. This notebook uses a custom `Metric` and `DataLoader` class that implement all the required methods.\n",
    "\n",
    "To implement `Metric` and `Dataloader`, you need to know the outputs of the model and the annotation format.\n",
    "\n",
    "The dataset in this example uses annotations in `JSON` format, with keys: `['categories', 'annotations', 'images']`. The `annotations` key is a list of dictionaries, with one item per annotation. Such item contains a `boxes` key, which holds the prediction boxes, in the `[xmin, xmax, ymin, ymax]` format. In this dataset, there is only one label: \"person\". \n",
    "\n",
    "The [model documentation](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/intel/person-detection-retail-0013) specifies that the model returns an array of shape `[1, 1, 200, 7]` where 200 is the number of detected boxes. Each detection has the format of `[image_id, label, conf, x_min, y_min, x_max, y_max]`. For this dataset, the label of `1` indicates a person.\n",
    "\n",
    "### Configuration\n",
    "\n",
    "#### Dataset\n",
    "\n",
    "The `DetectionDataLoader` class follows `compression.api.DataLoader` interface in POT, which should implement `__init__`, `__getitem__` and `__len__`, where `__getitem__` should return data as `(annotation, image)` or optionally `(annotation, image, metadata)`, with the annotation as `(index, label)`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54e4f0b3-5d83-433a-912b-65b6c84c7460",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DetectionDataLoader(DataLoader):\n",
    "    def __init__(self, basedir: str, target_size: Tuple[int, int]):\n",
    "        \"\"\"\n",
    "        :param basedir: Directory that contains images and annotation as \"annotation.json\"\n",
    "        :param target_size: Tuple of (width, height) to resize images to.\n",
    "        \"\"\"\n",
    "        self.images = sorted(Path(basedir).glob(\"*.jpg\"))\n",
    "        self.target_size = target_size\n",
    "        with open(f\"{basedir}/annotation_person_train.json\") as f:\n",
    "            self.annotations = json.load(f)\n",
    "        self.image_ids = {\n",
    "            Path(item[\"file_name\"]).name: item[\"id\"]\n",
    "            for item in self.annotations[\"images\"]\n",
    "        }\n",
    "\n",
    "        for image_filename in self.images:\n",
    "            annotations = [\n",
    "                item\n",
    "                for item in self.annotations[\"annotations\"]\n",
    "                if item[\"image_id\"] == self.image_ids[Path(image_filename).name]\n",
    "            ]\n",
    "            assert (\n",
    "                len(annotations) != 0\n",
    "            ), f\"No annotations found for image id {image_filename}\"\n",
    "\n",
    "        print(\n",
    "            f\"Created dataset with {len(self.images)} items. Data directory: {basedir}\"\n",
    "        )\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Get an item from the dataset at the specified index.\n",
    "        Detection boxes are converted from absolute coordinates to relative coordinates\n",
    "        between 0 and 1 by dividing xmin, xmax by image width and ymin, ymax by image height.\n",
    "\n",
    "        :return: (annotation, input_image, metadata) where annotation is (index, target_annotation)\n",
    "                 with target_annotation as a dictionary with keys category_id, image_width, image_height\n",
    "                 and bbox, containing the relative bounding box coordinates [xmin, ymin, xmax, ymax]\n",
    "                 (with values between 0 and 1) and metadata a dictionary: {\"filename\": path_to_image}\n",
    "        \"\"\"\n",
    "        image_path = self.images[index]\n",
    "        image = cv2.imread(str(image_path))\n",
    "        image = cv2.resize(image, self.target_size)\n",
    "        image_id = self.image_ids[Path(image_path).name]\n",
    "\n",
    "        # The `image_info` key contains height and width of the annotated image.\n",
    "        image_info = [\n",
    "            image for image in self.annotations[\"images\"] if image[\"id\"] == image_id\n",
    "        ][0]\n",
    "        # The `image_annotations` key contains the boxes and labels for the image.\n",
    "        image_annotations = [\n",
    "            item\n",
    "            for item in self.annotations[\"annotations\"]\n",
    "            if item[\"image_id\"] == image_id\n",
    "        ]\n",
    "\n",
    "        # The annotations are in xmin, ymin, width, height format. Convert to\n",
    "        # xmin, ymin, xmax, ymax and normalize to image width and height as\n",
    "        # stored in the annotation.\n",
    "        target_annotations = []\n",
    "        for annotation in image_annotations:\n",
    "            xmin, ymin, width, height = annotation[\"bbox\"]\n",
    "            xmax = xmin + width\n",
    "            ymax = ymin + height\n",
    "            xmin /= image_info[\"width\"]\n",
    "            ymin /= image_info[\"height\"]\n",
    "            xmax /= image_info[\"width\"]\n",
    "            ymax /= image_info[\"height\"]\n",
    "            target_annotation = {}\n",
    "            target_annotation[\"category_id\"] = annotation[\"category_id\"]\n",
    "            target_annotation[\"image_width\"] = image_info[\"width\"]\n",
    "            target_annotation[\"image_height\"] = image_info[\"height\"]\n",
    "            target_annotation[\"bbox\"] = [xmin, ymin, xmax, ymax]\n",
    "            target_annotations.append(target_annotation)\n",
    "\n",
    "        item_annotation = (index, target_annotations)\n",
    "        input_image = np.expand_dims(image.transpose(2, 0, 1), axis=0).astype(\n",
    "            np.float32\n",
    "        )\n",
    "        return (\n",
    "            item_annotation,\n",
    "            input_image,\n",
    "            {\"filename\": str(image_path), \"shape\": image.shape},\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cb85bf1-db51-4d5e-9238-4c24c3e13a8a",
   "metadata": {},
   "source": [
    "#### Metric\n",
    "\n",
    "Define a metric to determine the performance of a model. For the Default Quantization algorithm used in this notebook, defining a metric is optional, but it can be used to compare the quantized `INT8` model with the original FP OpenVINO IR model.\n",
    "\n",
    "This tutorial uses the Mean Average Precision (MAP) metric from [TorchMetrics](https://torchmetrics.readthedocs.io/en/latest/references/modules.html#detection-metrics).\n",
    "\n",
    "A metric for POT inherits from `compression.api.Metric` and should implement all the methods in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe75bfd1-9a87-449e-a918-5f839132f61a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MAPMetric(Metric):\n",
    "    def __init__(self, map_value=\"map\"):\n",
    "        \"\"\"\n",
    "        Mean Average Precision Metric. Wraps torchmetrics implementation, see\n",
    "        https://torchmetrics.readthedocs.io/en/latest/references/modules.html#map\n",
    "\n",
    "        :map_value: specific metric to return. Default: \"map\"\n",
    "                    Change `to one of the values in the list below to return a different value\n",
    "                    ['mar_1', 'mar_10', 'mar_100', 'mar_small', 'mar_medium', 'mar_large',\n",
    "                     'map', 'map_50', 'map_75', 'map_small', 'map_medium', 'map_large']\n",
    "                    See torchmetrics documentation for more details.\n",
    "        \"\"\"\n",
    "        assert (\n",
    "            map_value\n",
    "            in torchmetrics.detection.map.MARMetricResults.__slots__\n",
    "            + torchmetrics.detection.map.MAPMetricResults.__slots__\n",
    "        )\n",
    "\n",
    "        self._name = map_value\n",
    "        self.metric = torchmetrics.detection.map.MAP()\n",
    "        super().__init__()\n",
    "\n",
    "    @property\n",
    "    def value(self):\n",
    "        \"\"\"\n",
    "        Returns metric value for the last model output.\n",
    "        Possible format: {metric_name: [metric_values_per_image]}\n",
    "        \"\"\"\n",
    "        return {self._name: [0]}\n",
    "\n",
    "    @property\n",
    "    def avg_value(self):\n",
    "        \"\"\"\n",
    "        Returns average metric value for all model outputs.\n",
    "        Possible format: {metric_name: metric_value}\n",
    "        \"\"\"\n",
    "        return {self._name: self.metric.compute()[self._name].item()}\n",
    "\n",
    "    def update(self, output, target):\n",
    "        \"\"\"\n",
    "        Convert network output and labels to the format that torchmetrics' MAP\n",
    "        implementation expects, and call `metric.update()`.\n",
    "\n",
    "        :param output: model output\n",
    "        :param target: annotations for model output\n",
    "        \"\"\"\n",
    "        targetboxes = []\n",
    "        targetlabels = []\n",
    "        predboxes = []\n",
    "        predlabels = []\n",
    "        scores = []\n",
    "\n",
    "        image_width = target[0][0][\"image_width\"]\n",
    "        image_height = target[0][0][\"image_height\"]\n",
    "\n",
    "        for single_target in target[0]:\n",
    "            txmin, tymin, txmax, tymax = single_target[\"bbox\"]\n",
    "            category = single_target[\"category_id\"]\n",
    "            txmin *= image_width\n",
    "            txmax *= image_width\n",
    "            tymin *= image_height\n",
    "            tymax *= image_height\n",
    "\n",
    "            targetbox = [round(txmin), round(tymin), round(txmax), round(tymax)]\n",
    "            targetboxes.append(targetbox)\n",
    "            targetlabels.append(category)\n",
    "\n",
    "        for single_output in output:\n",
    "            for pred in single_output[0, 0, ::]:\n",
    "                image_id, label, conf, xmin, ymin, xmax, ymax = pred\n",
    "                xmin *= image_width\n",
    "                xmax *= image_width\n",
    "                ymin *= image_height\n",
    "                ymax *= image_height\n",
    "\n",
    "                predbox = [round(xmin), round(ymin), round(xmax), round(ymax)]\n",
    "                predboxes.append(predbox)\n",
    "                predlabels.append(label)\n",
    "                scores.append(conf)\n",
    "\n",
    "        preds = [\n",
    "            dict(\n",
    "                boxes=torch.Tensor(predboxes).float(),\n",
    "                labels=torch.Tensor(predlabels).short(),\n",
    "                scores=torch.Tensor(scores),\n",
    "            )\n",
    "        ]\n",
    "        targets = [\n",
    "            dict(\n",
    "                boxes=torch.Tensor(targetboxes).float(),\n",
    "                labels=torch.Tensor(targetlabels).short(),\n",
    "            )\n",
    "        ]\n",
    "        self.metric.update(preds, targets)\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Resets metric\n",
    "        \"\"\"\n",
    "        self.metric.reset()\n",
    "\n",
    "    def get_attributes(self):\n",
    "        \"\"\"\n",
    "        Returns a dictionary of metric attributes {metric_name: {attribute_name: value}}.\n",
    "        Required attributes: 'direction': 'higher-better' or 'higher-worse'\n",
    "                             'type': metric type\n",
    "        \"\"\"\n",
    "        return {self._name: {\"direction\": \"higher-better\", \"type\": \"mAP\"}}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "983102d5-fc76-46e2-983c-7df4e8b6765e",
   "metadata": {},
   "source": [
    "#### Quantization Config\n",
    "\n",
    "The POT methods expect configuration dictionaries as arguments, which are defined in the cell below. The `ir_path` variable points to the `.xml` file of the OpenVINO IR model. It is defined at the top of the notebook. This tutorial uses the `DefaultQuantization` algorithm.\n",
    "\n",
    "For more information about the settings and best practices, refer to the [Post-Training Optimization Best Practices](https://docs.openvino.ai/2021.4/pot_docs_BestPractices.html) and the main [POT documentation](https://docs.openvino.ai/2021.4/pot_README.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58cb1224-dfcb-4283-ae42-507fd4ed4157",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model config specifies the name of the model and paths to `.xml` and `.bin` files of the model.\n",
    "model_config = addict.Dict(\n",
    "    {\n",
    "        \"model_name\": ir_path.stem,\n",
    "        \"model\": ir_path,\n",
    "        \"weights\": ir_path.with_suffix(\".bin\"),\n",
    "    }\n",
    ")\n",
    "\n",
    "# Engine config\n",
    "engine_config = addict.Dict({\"device\": \"CPU\"})\n",
    "\n",
    "# Standard DefaultQuantization config. For this tutorial, stat_subset_size is ignored\n",
    "# because there are fewer than 300 images. For production use, 300 is recommended.\n",
    "default_algorithms = [\n",
    "    {\n",
    "        \"name\": \"DefaultQuantization\",\n",
    "        \"stat_subset_size\": 300,\n",
    "        \"params\": {\n",
    "            \"target_device\": \"ANY\",\n",
    "            \"preset\": \"mixed\",  # Choose between \"mixed\" and \"performance\".\n",
    "        },\n",
    "    }\n",
    "]\n",
    "\n",
    "print(f\"model_config: {model_config}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15df6f12-c8cc-4b94-bc04-322863d61e80",
   "metadata": {},
   "source": [
    "### Run Quantization Pipeline\n",
    "\n",
    "The POT pipeline uses the functions: `load_model()`, `IEEngine`, and `create_pipeline()`. The `load_model()` function loads an OpenVINO IR model specified in `model_config`. `IEEngine` is a POT implementation of OpenVINO Runtime that will be passed to the POT pipeline created by the `create_pipeline()` function. The POT classes and functions expect a config argument. These configs are created in the Config section in the cell above. `MAPMetric` metric and `DetectionDataLoader` have been defined earlier in this notebook.\n",
    "\n",
    "Creating and running the POT quantization pipeline takes just two lines of code. First, create the pipeline with the `create_pipeline` function, and then run that pipeline with `pipeline.run()`. To reuse the quantized model later, compress the model weights and save the compressed model to a disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9583d34e-5563-4b3c-b278-736ba721dc3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Create the data loader.\n",
    "data_loader = DetectionDataLoader(\n",
    "    basedir=\"data\", target_size=(input_width, input_height)\n",
    ")\n",
    "\n",
    "# Step 2: Load the model.\n",
    "ir_model = load_model(model_config=model_config)\n",
    "\n",
    "# Step 3: Initialize the metric.\n",
    "# For DefaultQuantization, specifying a metric is optional: metric can be set to None.\n",
    "metric = MAPMetric(map_value=\"map\")\n",
    "\n",
    "# Step 4: Initialize the engine for metric calculation and statistics collection.\n",
    "engine = IEEngine(config=engine_config, data_loader=data_loader, metric=metric)\n",
    "\n",
    "# Step 5: Create a pipeline of compression algorithms.\n",
    "# The `default_algorithms` parameter is defined in the Config cell above this cell.\n",
    "pipeline = create_pipeline(default_algorithms, engine)\n",
    "\n",
    "# Step 6: Execute the pipeline to quantize the model.\n",
    "algorithm_name = pipeline.algo_seq[0].name\n",
    "with yaspin(\n",
    "    text=f\"Executing POT pipeline on {model_config['model']} with {algorithm_name}\"\n",
    ") as sp:\n",
    "    start_time = time.perf_counter()\n",
    "    compressed_model = pipeline.run(ir_model)\n",
    "    end_time = time.perf_counter()\n",
    "    sp.ok(\"✔\")\n",
    "print(f\"Quantization finished in {end_time - start_time:.2f} seconds\")\n",
    "\n",
    "# Step 7 (Optional): Compress model weights to quantized precision\n",
    "#                    in order to reduce the size of the final `.bin` file.\n",
    "compress_model_weights(compressed_model)\n",
    "\n",
    "# Step 8: Save the compressed model to the desired path.\n",
    "# Set `save_path` to the directory where the compressed model should be stored.\n",
    "preset = pipeline._algo_seq[0].config[\"preset\"]\n",
    "algorithm_name = pipeline.algo_seq[0].name\n",
    "compressed_model_paths = save_model(\n",
    "    model=compressed_model,\n",
    "    save_path=\"optimized_model\",\n",
    "    model_name=f\"{ir_model.name}_{preset}_{algorithm_name}\",\n",
    ")\n",
    "\n",
    "compressed_model_path = compressed_model_paths[0][\"model\"]\n",
    "print(\"The quantized model is stored at\", compressed_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3babde6-5e33-4ec5-9c88-06e8629c26a1",
   "metadata": {},
   "source": [
    "## Compare Metric of Floating Point and Quantized Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1dface1-ccfc-4ff9-b0c1-66fbed93d006",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the mAP on the quantized model and compare with the mAP on the FP16 OpenVINO IR model.\n",
    "ir_model = load_model(model_config=model_config)\n",
    "evaluation_pipeline = create_pipeline(algo_config=dict(), engine=engine)\n",
    "\n",
    "with yaspin(text=\"Evaluating original IR model\") as sp:\n",
    "    original_metric = evaluation_pipeline.evaluate(ir_model)\n",
    "\n",
    "with yaspin(text=\"Evaluating quantized IR model\") as sp:\n",
    "    quantized_metric = pipeline.evaluate(compressed_model)\n",
    "\n",
    "if original_metric:\n",
    "    for key, value in original_metric.items():\n",
    "        print(f\"The {key} score of the original FP16 model is {value:.5f}\")\n",
    "\n",
    "if quantized_metric:\n",
    "    for key, value in quantized_metric.items():\n",
    "        print(f\"The {key} score of the quantized INT8 model is {value:.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a43956df-b28c-46bc-87fa-c000615e03b8",
   "metadata": {},
   "source": [
    "## Visualize Results\n",
    "\n",
    "Compare the annotated boxes (green) with the results of the floating point (red) and quantized (green) models. First, define a helper function to draw the boxes on an image, using the specified color. Then, do inference on five images and show the results. The figure shows three images for every input image: the left image shows the annotation and both `FP` and `INT8` predictions, the middle image shows the floating point model prediction separately, and the image to the right shows the quantized model prediction. The mAP score of the prediction is shown with each prediction. Predicted boxes with a confidence value of at least 0.5 will be shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39cea6ce-dd5c-4e2c-bc6a-511a7c91c670",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_boxes_on_image(\n",
    "    box: Sequence[float], image: np.ndarray, color: str, scale: bool = True\n",
    "):\n",
    "    \"\"\"\n",
    "    Draw a `box` on an `image` with a `color`, optionally scaling the box from normalized\n",
    "    coordinates (between 0 and 1) to image coordinates.\n",
    "    This is a utility function for binary detection where all boxes belong to one category.\n",
    "\n",
    "    :param box: Box coordinates as [xmin, ymin, xmax, ymax]\n",
    "    :param image: numpy array of RGB image\n",
    "    :param color: Box color, \"red\", \"green\" or \"blue\"\n",
    "    \"param scale: If True, scale normalized box coordinates to absolute coordinates based\n",
    "                  on the image size.\n",
    "    \"\"\"\n",
    "    colors = {\"red\": (255, 0, 64), \"green\": (0, 255, 0), \"yellow\": (255, 255, 128)}\n",
    "    assert color in colors, f\"{color} is not defined yet. Defined colors are: {colors}\"\n",
    "    image_height, image_width, _ = image.shape\n",
    "    x_min, y_min, x_max, y_max = box\n",
    "    if scale:\n",
    "        x_min *= image_width\n",
    "        x_max *= image_width\n",
    "        y_min *= image_height\n",
    "        y_max *= image_height\n",
    "\n",
    "    image = cv2.rectangle(\n",
    "        img=image,\n",
    "        pt1=(round(x_min), round(y_min)),\n",
    "        pt2=(round(x_max), round(y_max)),\n",
    "        color=colors[color],\n",
    "        thickness=2,\n",
    "    )\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb9c5da-c4e1-4d3c-bddf-50a07eee6c2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change `map_value` to one of the values in the list below to show a different metric.\n",
    "# See https://torchmetrics.readthedocs.io/en/latest/references/modules.html#map\n",
    "# ('map', 'map_50', 'map_75', 'map_small', 'map_medium', 'map_large'\n",
    "#  'mar_1', 'mar_10', 'mar_100', 'mar_small', 'mar_medium', 'mar_large')\n",
    "\n",
    "map_value = \"map\"\n",
    "confidence_threshold = 0.5\n",
    "num_images = 4\n",
    "\n",
    "# FP prediction\n",
    "fp_model = ie.read_model(model=ir_path)\n",
    "fp_compiled_model = ie.compile_model(model=fp_model, device_name=\"CPU\")\n",
    "input_layer_fp = fp_compiled_model.input(0)\n",
    "output_layer_fp = fp_compiled_model.output(0)\n",
    "\n",
    "# INT8 prediction\n",
    "int8_model = ie.read_model(model=compressed_model_path)\n",
    "int8_compiled_model = ie.compile_model(model=int8_model, device_name=\"CPU\")\n",
    "input_layer_int8 = int8_compiled_model.input(0)\n",
    "output_layer_int8 = int8_compiled_model.output(0)\n",
    "\n",
    "fig, axs = plt.subplots(nrows=num_images, ncols=3, figsize=(16, 14), squeeze=False)\n",
    "for i in range(num_images):\n",
    "    annotation, input_image, metadata = data_loader[i]\n",
    "    image = cv2.cvtColor(\n",
    "        src=cv2.imread(filename=metadata[\"filename\"]), code=cv2.COLOR_BGR2RGB\n",
    "    )\n",
    "    orig_image = image.copy()\n",
    "    resized_image = cv2.resize(image, (input_width, input_height))\n",
    "    target_annotation = annotation[1]\n",
    "\n",
    "    fp_res = fp_compiled_model([input_image])[output_layer_fp]\n",
    "    \n",
    "    fp_metric = MAPMetric(map_value=map_value)\n",
    "    fp_metric.update(output=[fp_res], target=[target_annotation])\n",
    "\n",
    "    for item in fp_res[0, 0, ::]:\n",
    "        _, _, conf, xmin, xmax, ymin, ymax = item\n",
    "        if conf > confidence_threshold:\n",
    "            total_image = draw_boxes_on_image([xmin, xmax, ymin, ymax], image, \"red\")\n",
    "\n",
    "    axs[i, 1].imshow(total_image)\n",
    "\n",
    "    int8_res = int8_compiled_model([input_image])[output_layer_int8]\n",
    "    int8_metric = MAPMetric(map_value=map_value)\n",
    "    int8_metric.update(output=[int8_res], target=[target_annotation])\n",
    "\n",
    "    for item in int8_res[0, 0, ::]:\n",
    "        _, _, conf, xmin, xmax, ymin, ymax = item\n",
    "        if conf > confidence_threshold:\n",
    "            total_image = draw_boxes_on_image(\n",
    "                [xmin, xmax, ymin, ymax], total_image, \"yellow\"\n",
    "            )\n",
    "            int8_image = draw_boxes_on_image(\n",
    "                [xmin, xmax, ymin, ymax], orig_image, \"yellow\"\n",
    "            )\n",
    "\n",
    "    axs[i, 2].imshow(int8_image)\n",
    "\n",
    "    # Annotation\n",
    "    for annotation in target_annotation:\n",
    "        total_image = draw_boxes_on_image(annotation[\"bbox\"], total_image, \"green\")\n",
    "\n",
    "    axs[i, 0].imshow(image)\n",
    "    axs[i, 0].set_title(Path(metadata[\"filename\"]).stem)\n",
    "    axs[i, 1].set_title(f\"FP32 mAP: {fp_metric.avg_value[map_value]:.3f}\")\n",
    "    axs[i, 2].set_title(f\"INT8 mAP: {int8_metric.avg_value[map_value]:.3f}\")\n",
    "    fig.suptitle(\n",
    "        \"Annotated (green) and detected boxes on FP (red) and INT8 (yellow) model\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c18abd6-393e-4376-a3fc-d389e5b73ff3",
   "metadata": {},
   "source": [
    "## Compare the Size of the Original and Quantized Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d599fa7-bf93-4476-a183-7cad01d2e4ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "original_model_size = Path(ir_path).with_suffix(\".bin\").stat().st_size / 1024\n",
    "quantized_model_size = (\n",
    "    Path(compressed_model_path).with_suffix(\".bin\").stat().st_size / 1024\n",
    ")\n",
    "\n",
    "print(f\"FP32 model size: {original_model_size:.2f} KB\")\n",
    "print(f\"INT8 model size: {quantized_model_size:.2f} KB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "322e899a-5bcf-4759-964e-9c9721d39928",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compare Performance of the Original and Quantized Models\n",
    "\n",
    "To measure the inference performance of OpenVINO `FP16` and `INT8` models, use [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) in OpenVINO.\n",
    "\n",
    "> **Note**: The `benchmark_app` tool is able to measure the performance of the OpenVINO Intermediate Representation (OpenVINO IR) models only. For more accurate performance, run `benchmark_app` in a terminal/command prompt after closing other applications. Run `benchmark_app -m model.xml -d CPU` to benchmark async inference on CPU for one minute. Change `CPU` to `GPU` to benchmark on GPU. Run `benchmark_app --help` to see an overview of all command-line options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2249c7c-9ba6-4f31-a30e-b09f2dd26a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! benchmark_app --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5280f83f-7642-4f1f-aa63-07e2fb3e786d",
   "metadata": {
    "tags": [],
    "test_replace": {
     "-t 15": "-t 3"
    }
   },
   "outputs": [],
   "source": [
    "# Benchmark FP16 model\n",
    "! benchmark_app -m $ir_path -d CPU -t 15 -api sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ab9bc3-372d-4ab3-8c4b-cde35a7f9215",
   "metadata": {
    "tags": [],
    "test_replace": {
     "-t 15": "-t 3"
    }
   },
   "outputs": [],
   "source": [
    "# Benchmark INT8 model\n",
    "! benchmark_app -m $compressed_model_path -d CPU -t 15 -api sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba7471f6-2aa1-494d-8b4b-e818cc23a008",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Benchmark INT8 model on MULTI:CPU,GPU device (requires an Intel integrated GPU)\n",
    "ie = Core()\n",
    "if \"GPU\" in ie.available_devices:\n",
    "    ! benchmark_app -m $compressed_model_path -d MULTI:CPU,GPU -t 15 -api sync -b 4"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "5ac51ce60ed7a11cc2408f3858df5a3450d75f1c9bd0c072efa487aad689faf9"
  },
  "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
}