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 OpenVINO's [Post-Training Optimization Tool API](https://docs.openvino.ai/2021.4/pot_compression_api_README.html). \n",
    "\n",
    "For demonstration purposes, we use a very small dataset of 10 images presenting people at the airport. 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 model used is: [person-detection-retail-0013](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/intel/person-detection-retail-0013)"
   ]
  },
  {
   "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\")\n",
    "from notebook_utils import benchmark_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "033bed78-0daa-4501-b41b-5912d81e6719",
   "metadata": {},
   "source": [
    "### Download Model\n",
    "\n",
    "Download the model from Open Model Zoo, if it does not yet exist. "
   ]
  },
  {
   "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 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`. In this notebook, we use a custom Metric and DataLoader class that implement all the required methods.\n",
    "\n",
    "To implement the Metric and Dataloader, we 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']`. `annotations` 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 POT's `compression.api.DataLoader` interface, which should implement `__init__`, `__getitem__` and `__len__`, where `__getitem__` should return data as `(annotation, image)` or optionally `(annotation, image, metadata)`, with 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",
    "        # image_info 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",
    "        # image_annotations 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",
    "        # 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 model's performance. 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 IR model.\n",
    "\n",
    "In this tutorial we use 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",
    "POT methods expect configuration dictionaries as arguments, which are defined in the cell below. The variable `ir_path` points to the IR model's xml file. It is defined at the top of the notebook. In this tutorial, we use the DefaultQuantization algorithm.\n",
    "\n",
    "See [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) page for more information about the settings and best practices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58cb1224-dfcb-4283-ae42-507fd4ed4157",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model config specifies the model name and paths to model .xml and .bin file\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()`. `load_model()` loads an IR model specified in `model_config`. `IEEngine` is a POT implementation of Inference Engine that will be passed to the POT pipeline created by `create_pipeline()`. The POT classes and functions expect a config argument. These configs are created in the Config section in the cell above. The 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. We create the pipeline with the `create_pipeline` function, and then run that pipeline with `pipeline.run()`. To reuse the quantized model later, we compress the model weights and save the compressed model to disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9583d34e-5563-4b3c-b278-736ba721dc3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: create data loader\n",
    "data_loader = DetectionDataLoader(\n",
    "    basedir=\"data\", target_size=(input_width, input_height)\n",
    ")\n",
    "\n",
    "# Step 2: load 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",
    "# algorithms 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 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 `box` on `image` with `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 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 inference performance of the FP16 and INT8 models, we use OpenVINO's benchmarking solution, the [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html). It can be run in the notebook with: `! benchmark_app` or `%sx benchmark_app`.\n",
    "\n",
    "In this tutorial, we use a wrapper function from [Notebook Utils](https://github.com/openvinotoolkit/openvino_notebooks/blob/main/notebooks/utils/notebook_utils.ipynb). It prints the `benchmark_app` command with the chosen parameters.\n",
    "\n",
    "> NOTE: For the most accurate performance estimation, we recommended running `benchmark_app` in a terminal/command prompt after closing other applications. Run `benchmark_app --help` to see 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": "63452dbe-b1c4-4f2c-a5a7-77a021c7b64c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# benchmark_model??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5280f83f-7642-4f1f-aa63-07e2fb3e786d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benchmark FP16 model\n",
    "benchmark_model(model_path=ir_path, device=\"CPU\", seconds=15, api=\"async\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ab9bc3-372d-4ab3-8c4b-cde35a7f9215",
   "metadata": {
    "tags": [],
    "test_replace": {
     "seconds=15": "seconds=3"
    }
   },
   "outputs": [],
   "source": [
    "# Benchmark INT8 model\n",
    "benchmark_model(model_path=compressed_model_path, device=\"CPU\", seconds=15, api=\"async\")"
   ]
  },
  {
   "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_model(\n",
    "        model_path=compressed_model_path,\n",
    "        device=\"MULTI:CPU,GPU\",\n",
    "        seconds=15,\n",
    "        api=\"async\",\n",
    "        batch=4,\n",
    "    )"
   ]
  }
 ],
 "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
}