notebooks/110-ct-segmentation-quantize/110-ct-segmentation-quantize-nncf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantize a Segmentation Model and Show Live Inference\n",
    "\n",
    "## Kidney Segmentation with PyTorch Lightning and OpenVINO™ - Part 3\n",
    "\n",
    "This tutorial is part of a series on how to train, optimize, quantize and show live inference on a medical segmentation model. The goal is to accelerate inference on a kidney segmentation model. The [UNet](https://arxiv.org/abs/1505.04597) model is trained from scratch; the data is from [Kits19](https://github.com/neheller/kits19).\n",
    "\n",
    "This third tutorial in the series shows how to:\n",
    "\n",
    "- Convert an ONNX model to OpenVINO IR with [Model Optimizer](https://docs.openvino.ai/latest/openvino_docs_MO_DG_Deep_Learning_Model_Optimizer_DevGuide.html),\n",
    "- Quantize a PyTorch model with NNCF\n",
    "- Evaluate the F1 score metric of the original model and the quantized model\n",
    "- Benchmark performance of the original model and the quantized model\n",
    "- Show live inference with OpenVINO's async API\n",
    "\n",
    "All notebooks in this series:\n",
    "\n",
    "- [Data Preparation for 2D Segmentation of 3D Medical Data](data-preparation-ct-scan.ipynb)\n",
    "- [Train a 2D-UNet Medical Imaging Model with PyTorch Lightning](pytorch-monai-training.ipynb)\n",
    "- Convert and Quantize a Segmentation Model and Show Live Inference (this notebook)\n",
    "- [Live Inference and Benchmark CT-scan data](../210-ct-scan-live-inference/210-ct-scan-live-inference.ipynb) \n",
    "\n",
    "## Instructions\n",
    "\n",
    "This notebook needs a trained UNet model. We provide a pretrained model trained for 20 epochs with the full [Kits-19](https://github.com/neheller/kits19) frames dataset, which has an F1 score on the validation set of 0.9. The training code is available in [this notebook](pytorch-monai-training.ipynb). \n",
    "\n",
    "NNCF for PyTorch models requires a C++ compiler. On Windows, please install [Microsoft Visual Studio 2019](https://docs.microsoft.com/en-us/visualstudio/releases/2019/release-notes). During installation, choose Desktop development with C++ in the Workloads tab. On macOS, \n",
    "run `xcode-select –install` from a Terminal. On Linux, please install gcc.\n",
    "\n",
    "Running this notebook with the full dataset will take a long time. For demonstration purposes, this tutorial will download one converted CT scan and use that scan for quantization and inference. For production use, please use a representative dataset for quantizing the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# On Windows, try to find the directory that contains x64 cl.exe and add it to the PATH to enable PyTorch\n",
    "# to find the required C++ tools. This code assumes that Visual Studio is installed in the default\n",
    "# directory. If you have a different C++ compiler, please add the correct path to os.environ[\"PATH\"]\n",
    "# directly. Note that the C++ Redistributable is not enough to run this notebook.\n",
    "\n",
    "# Adding the path to os.environ[\"LIB\"] is not always required - it depends on the system's configuration\n",
    "\n",
    "import sys\n",
    "\n",
    "if sys.platform == \"win32\":\n",
    "    import distutils.command.build_ext\n",
    "    import os\n",
    "    from pathlib import Path\n",
    "\n",
    "    if sys.getwindowsversion().build >= 20000:  # Windows 11\n",
    "        search_path = \"**/Hostx64/x64/cl.exe\"\n",
    "    else:\n",
    "        search_path = \"**/Hostx86/x64/cl.exe\"\n",
    "\n",
    "    VS_INSTALL_DIR_2019 = r\"C:/Program Files (x86)/Microsoft Visual Studio\"\n",
    "    VS_INSTALL_DIR_2022 = r\"C:/Program Files/Microsoft Visual Studio\"\n",
    "\n",
    "    cl_paths_2019 = sorted(list(Path(VS_INSTALL_DIR_2019).glob(search_path)))\n",
    "    cl_paths_2022 = sorted(list(Path(VS_INSTALL_DIR_2022).glob(search_path)))\n",
    "    cl_paths = cl_paths_2019 + cl_paths_2022\n",
    "\n",
    "    if len(cl_paths) == 0:\n",
    "        raise ValueError(\n",
    "            \"Cannot find Visual Studio. This notebook requires an x64 C++ compiler. If you installed \"\n",
    "            \"a C++ compiler, please add the directory that contains cl.exe to `os.environ['PATH']`.\"\n",
    "        )\n",
    "    else:\n",
    "        # If multiple versions of MSVC are installed, get the most recent version\n",
    "        cl_path = cl_paths[-1]\n",
    "        vs_dir = str(cl_path.parent)\n",
    "        os.environ[\"PATH\"] += f\"{os.pathsep}{vs_dir}\"\n",
    "        # Code for finding the library dirs from\n",
    "        # https://stackoverflow.com/questions/47423246/get-pythons-lib-path\n",
    "        d = distutils.core.Distribution()\n",
    "        b = distutils.command.build_ext.build_ext(d)\n",
    "        b.finalize_options()\n",
    "        os.environ[\"LIB\"] = os.pathsep.join(b.library_dirs)\n",
    "        print(f\"Added {vs_dir} to PATH\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "import warnings\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import monai\n",
    "import numpy as np\n",
    "import torch\n",
    "from monai.transforms import LoadImage\n",
    "from nncf.common.utils.logger import set_log_level\n",
    "from openvino.runtime import Core\n",
    "from torch.jit import TracerWarning\n",
    "from torchmetrics import F1\n",
    "\n",
    "set_log_level(logging.ERROR)  # Disables all NNCF info and warning messages\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from models.custom_segmentation import SegmentationModel\n",
    "from notebook_utils import NotebookAlert, download_file, show_live_inference\n",
    "\n",
    "try:\n",
    "    import subprocess\n",
    "\n",
    "    from nncf import NNCFConfig\n",
    "    from nncf.torch import create_compressed_model, register_default_init_args\n",
    "except subprocess.CalledProcessError:\n",
    "    message = \"WARNING: Running this notebook requires an x64 C++ compiler.\"\n",
    "    NotebookAlert(message=message, alert_class=\"warning\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings\n",
    "\n",
    "By default, this notebook will download one CT scan from the KITS19 dataset, and use that for quantization. To use the full dataset, set `BASEDIR` to the path of the dataset, as prepared according to the [Data Preparation](data-preparation-ct-scan.ipynb) notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "BASEDIR = Path(\"kits19_frames_1\")\n",
    "# Uncomment the line below to use the full dataset, as prepared in the data preparation notebook\n",
    "# BASEDIR = Path(\"~/kits19/kits19_frames\").expanduser()\n",
    "MODEL_DIR = Path(\"model\")\n",
    "MODEL_DIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load PyTorch Model\n",
    "\n",
    "Download the pretrained model weights, load the PyTorch model and the state_dict that was saved after training. The model used in this notebook is a [BasicUnet](https://docs.monai.io/en/stable/networks.html#basicunet) model from [MONAI](https://monai.io). We provide a pretrained checkpoint. To see how this model yourself, check out the [training notebook](pytorch-monai-training.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "state_dict_url = \"https://github.com/helena-intel/openvino_notebooks/raw/110-nncf/notebooks/110-ct-segmentation-quantize/pretrained_model/unet_kits19_state_dict.pth\"\n",
    "state_dict_file = download_file(state_dict_url, directory=\"pretrained_model\")\n",
    "state_dict = torch.load(state_dict_file, map_location=torch.device(\"cpu\"))\n",
    "\n",
    "new_state_dict = {}\n",
    "for k, v in state_dict.items():\n",
    "    new_key = k.replace(\"_model.\", \"\")\n",
    "    new_state_dict[new_key] = v\n",
    "new_state_dict.pop(\"loss_function.pos_weight\")\n",
    "\n",
    "model = monai.networks.nets.BasicUNet(spatial_dims=2, in_channels=1, out_channels=1).eval()\n",
    "model.load_state_dict(new_state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We export the PyTorch model to ONNX and convert it to OpenVINO IR, for comparing the performance of the FP32 and INT8 model later in this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dummy_input = torch.randn(1, 1, 512, 512)\n",
    "fp32_onnx_path = MODEL_DIR / \"unet_kits19_fp32.onnx\"\n",
    "torch.onnx.export(model, dummy_input, fp32_onnx_path)\n",
    "!mo --input_model \"$fp32_onnx_path\" --output_dir $MODEL_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download CT-scan Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The CT scan case number. For example: 2 for data from the case_00002 directory\n",
    "# Currently only 117 is supported\n",
    "CASE = 117\n",
    "if not (BASEDIR / f\"case_{CASE:05d}\").exists():\n",
    "    BASEDIR.mkdir(exist_ok=True)\n",
    "    filename = download_file(\n",
    "        f\"https://storage.openvinotoolkit.org/data/test_data/openvino_notebooks/kits19/case_{CASE:05d}.zip\"\n",
    "    )\n",
    "    with zipfile.ZipFile(filename, \"r\") as zip_ref:\n",
    "        zip_ref.extractall(path=BASEDIR)\n",
    "    os.remove(filename)  # remove zipfile\n",
    "    print(f\"Downloaded and extracted data for case_{CASE:05d}\")\n",
    "else:\n",
    "    print(f\"Data for case_{CASE:05d} exists\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "\n",
    "The KitsDataset class in the next cell expects images and masks in the *basedir* directory, in a folder per patient. It is a simplified version of the DataSet class in the [training notebook](pytorch-monai-training.ipynb).\n",
    "\n",
    "Images are loaded with MONAI's [`LoadImage`](https://docs.monai.io/en/stable/transforms.html#loadimage), to align with the image loading method in the training notebook. This method rotates and flips the images. We define a `rotate_and_flip` method to display the images in the expected orientation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "return len(self.dataset)": "return 30"
    }
   },
   "outputs": [],
   "source": [
    "def rotate_and_flip(image):\n",
    "    \"\"\"Rotate `image` by 90 degrees and flip horizontally\"\"\"\n",
    "    return cv2.flip(cv2.rotate(image, rotateCode=cv2.ROTATE_90_CLOCKWISE), flipCode=1)\n",
    "\n",
    "\n",
    "class KitsDataset:\n",
    "    def __init__(self, basedir: str):\n",
    "        \"\"\"\n",
    "        Dataset class for prepared Kits19 data, for binary segmentation (background/kidney)\n",
    "        Source data should exist in basedir, in subdirectories case_00000 until case_00210,\n",
    "        with each subdirectory containing directories imaging_frames, with jpg images, and\n",
    "        segmentation_frames with segmentation masks as png files.\n",
    "        See https://github.com/openvinotoolkit/openvino_notebooks/blob/main/notebooks/110-ct-segmentation-quantize/data-preparation-ct-scan.ipynb\n",
    "\n",
    "        :param basedir: Directory that contains the prepared CT scans\n",
    "        \"\"\"\n",
    "        masks = sorted(BASEDIR.glob(\"case_*/segmentation_frames/*png\"))\n",
    "\n",
    "        self.basedir = basedir\n",
    "        self.dataset = masks\n",
    "        print(\n",
    "            f\"Created dataset with {len(self.dataset)} items. \"\n",
    "            f\"Base directory for data: {basedir}\"\n",
    "        )\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Get an item from the dataset at the specified index.\n",
    "\n",
    "        :return: (image, segmentation_mask)\n",
    "        \"\"\"\n",
    "        mask_path = self.dataset[index]\n",
    "        image_path = str(mask_path.with_suffix(\".jpg\")).replace(\n",
    "            \"segmentation_frames\", \"imaging_frames\"\n",
    "        )\n",
    "\n",
    "        # Load images with MONAI's LoadImage to match data loading in training notebook\n",
    "        mask = LoadImage(image_only=True, dtype=np.uint8)(str(mask_path)).numpy()\n",
    "        img = LoadImage(image_only=True, dtype=np.float32)(str(image_path)).numpy()\n",
    "\n",
    "        if img.shape[:2] != (512, 512):\n",
    "            img = cv2.resize(img.astype(np.uint8), (512, 512)).astype(np.float32)\n",
    "            mask = cv2.resize(mask, (512, 512))\n",
    "\n",
    "        input_image = np.expand_dims(img, axis=0)\n",
    "        return input_image, mask\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test that the data loader returns the expected output, we show an image and a mask. The image and mask are shown as returned by the dataloader, after resizing and preprocessing. Since this dataset contains a lot of slices without kidneys, we select a slice that contains at least 5000 kidney pixels to verify that the annotations look correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = KitsDataset(BASEDIR)\n",
    "# Find a slice that contains kidney annotations\n",
    "# item[0] is the annotation: (id, annotation_data)\n",
    "image_data, mask = next(item for item in dataset if np.count_nonzero(item[1]) > 5000)\n",
    "# Remove extra image dimension and rotate and flip the image for visualization\n",
    "image = rotate_and_flip(image_data.squeeze())\n",
    "\n",
    "# The data loader returns annotations as (index, mask) and mask in shape (H,W)\n",
    "mask = rotate_and_flip(mask)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n",
    "ax[0].imshow(image, cmap=\"gray\")\n",
    "ax[1].imshow(mask, cmap=\"gray\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Metric\n",
    "Define a metric to determine the performance of the model. \n",
    "\n",
    "For this demo, we use the [F1 score](https://en.wikipedia.org/wiki/F-score), or Dice coefficient, from the [TorchMetrics](https://torchmetrics.readthedocs.io/en/stable/) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_f1(model: torch.nn.Module, dataset: KitsDataset):\n",
    "    \"\"\"\n",
    "    Compute binary F1 score of `model` on `dataset`\n",
    "    F1 score metric is provided by the torchmetrics library\n",
    "    `model` is expected to be a binary segmentation model, images in the\n",
    "    dataset are expected in (N,C,H,W) format where N==C==1\n",
    "    \"\"\"\n",
    "    metric = F1(ignore_index=0)\n",
    "    with torch.no_grad():\n",
    "        for image, target in dataset:\n",
    "            input_image = torch.as_tensor(image).unsqueeze(0)\n",
    "            output = model(input_image)\n",
    "            label = torch.as_tensor(target.squeeze()).long()\n",
    "            prediction = torch.sigmoid(output.squeeze()).round().long()\n",
    "            metric.update(label.flatten(), prediction.flatten())\n",
    "    return metric.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Quantization\n",
    "\n",
    "Before quantizing the model, we compute the F1 score on the FP32 model for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fp32_f1 = compute_f1(model, dataset)\n",
    "print(f\"FP32 F1: {fp32_f1:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NNCF configuration can be defined in a json file or a dictionary. See the [NNCF quantization documentation](https://github.com/openvinotoolkit/nncf/blob/develop/docs/compression_algorithms/Quantization.md) for more information on the possible values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# NNCF uses the model loaded at the beginning of this notebook. If after quantizing the model, you\n",
    "# want to quantize with a different config, reload the model by uncommenting the next two lines\n",
    "#\n",
    "# model = monai.networks.nets.BasicUNet(spatial_dims=2, in_channels=1, out_channels=1).eval()\n",
    "# model.load_state_dict(new_state_dict)\n",
    "\n",
    "nncf_config_dict = {\n",
    "    \"input_info\": {\"sample_size\": [1, 1, 512, 512]},\n",
    "    \"target_device\": \"CPU\",\n",
    "    \"compression\": {\n",
    "        \"algorithm\": \"quantization\",\n",
    "        # performance preset uses symmetric weights and activations\n",
    "        \"preset\": \"performance\",\n",
    "        # Do not quantize LeakyReLU activations to allow the INT8 model to run on Intel GPU\n",
    "        \"ignored_scopes\": [\"{re}.*LeakyReLU*\"],\n",
    "    },\n",
    "}\n",
    "\n",
    "nncf_config = NNCFConfig.from_dict(nncf_config_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-24T14:05:30.690144Z",
     "iopub.status.busy": "2022-05-24T14:05:30.690006Z",
     "iopub.status.idle": "2022-05-24T14:05:30.692969Z",
     "shell.execute_reply": "2022-05-24T14:05:30.692533Z",
     "shell.execute_reply.started": "2022-05-24T14:05:30.690133Z"
    },
    "tags": []
   },
   "source": [
    "Create a quantized model from a pre-trained FP32 model and configuration object created with the dictionary defined in the previous cell and a DataLoader. See the [NNCF documentation](https://github.com/openvinotoolkit/nncf/blob/develop/docs/Usage.md) for more information. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_loader = torch.utils.data.DataLoader(dataset, batch_size=4)\n",
    "nncf_config = register_default_init_args(nncf_config, data_loader)\n",
    "compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `compressed_model` that was created in the previous cell is a PyTorch `nn.Module` that is wrapped by NNCF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compressed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we demonstrate post-training quantization with NNCF. To use the quantized model for faster inference with OpenVINO, we export the model to ONNX, and convert the ONNX model to OpenVINO's IR format.\n",
    "\n",
    "NNCF also supports quantization-aware training, and other algorithms than quantization. See the [NNCF documentation](https://github.com/openvinotoolkit/nncf/) in the NNCF repository for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "int8_onnx_path = MODEL_DIR / \"unet_kits19_int8.onnx\"\n",
    "warnings.filterwarnings(\"ignore\", category=TracerWarning)  # Ignore export warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "compression_ctrl.export_model(str(int8_onnx_path))\n",
    "print(f\"INT8 ONNX model exported to {int8_onnx_path}.\")\n",
    "\n",
    "!mo --input_model \"$int8_onnx_path\" --output_dir \"$MODEL_DIR\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Compare FP32 and INT8 Model\n",
    "\n",
    "## Compare File Size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fp32_ir_path = Path(fp32_onnx_path).with_suffix(\".xml\")\n",
    "int8_ir_path = Path(int8_onnx_path).with_suffix(\".xml\")\n",
    "\n",
    "original_model_size = fp32_ir_path.with_suffix(\".bin\").stat().st_size / 1024\n",
    "quantized_model_size = int8_ir_path.with_suffix(\".bin\").stat().st_size / 1024\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",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compare Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "int8_f1 = compute_f1(compressed_model, dataset)\n",
    "\n",
    "print(f\"FP32 F1: {fp32_f1:.3f}\")\n",
    "print(f\"INT8 F1: {int8_f1:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compare Performance of the Original and Quantized Models\n",
    "\n",
    "To measure the inference performance of the FP32 and INT8 models, we use [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html), OpenVINO's inference performance measurement tool. Benchmark tool is a command line application, part of OpenVINO development tools, that can be run in the notebook with `! benchmark_app` or `%sx benchmark_app`.\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 -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 all command line options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# ! benchmark_app --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"CPU\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "-t 15": "-t 3"
    }
   },
   "outputs": [],
   "source": [
    "# Benchmark FP32 model\n",
    "! benchmark_app -m $fp32_ir_path -d $device -t 15 -api sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "-t 15": "-t 3"
    }
   },
   "outputs": [],
   "source": [
    "# Benchmark INT8 model\n",
    "! benchmark_app -m $int8_ir_path -d $device -t 15 -api sync"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Visually Compare Inference Results\n",
    "\n",
    "Visualize the results of the model on four slices of the validation set. Compare the results of the FP32 IR model with the results of the quantized INT8 model and the reference segmentation annotation.\n",
    "\n",
    "Medical imaging datasets tend to be very imbalanced: most of the slices in a CT scan do not contain kidney data. The segmentation model should be good at finding kidneys where they exist (in medical terms: have good sensitivity) but also not find spurious kidneys that do not exist (have good specificity). In the next cell, we show four slices: two slices that have no kidney data, and two slices that contain kidney data. For this example, a slice has kidney data if at least 50 pixels in the slices are annotated as kidney.\n",
    "\n",
    "Run this cell again to show results on a different subset. The random seed is displayed to allow reproducing specific runs of this cell.\n",
    "\n",
    "> Note: the images are shown after optional augmenting and resizing. In the Kits19 dataset all but one of the cases has input shape `(512, 512)`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The sigmoid function is used to transform the result of the network\n",
    "# to binary segmentation masks\n",
    "def sigmoid(x):\n",
    "    return np.exp(-np.logaddexp(0, -x))\n",
    "\n",
    "\n",
    "num_images = 4\n",
    "colormap = \"gray\"\n",
    "\n",
    "# Load FP32 and INT8 models\n",
    "core = Core()\n",
    "fp_model = core.read_model(fp32_ir_path)\n",
    "int8_model = core.read_model(int8_ir_path)\n",
    "compiled_model_fp = core.compile_model(fp_model, device_name=\"CPU\")\n",
    "compiled_model_int8 = core.compile_model(int8_model, device_name=\"CPU\")\n",
    "output_layer_fp = compiled_model_fp.output(0)\n",
    "output_layer_int8 = compiled_model_int8.output(0)\n",
    "\n",
    "# Create subset of dataset\n",
    "background_slices = (item for item in dataset if np.count_nonzero(item[1]) == 0)\n",
    "kidney_slices = (item for item in dataset if np.count_nonzero(item[1]) > 50)\n",
    "data_subset = random.sample(list(background_slices), 2) + random.sample(list(kidney_slices), 2)\n",
    "\n",
    "# Set seed to current time. To reproduce specific results, copy the printed seed\n",
    "# and manually set `seed` to that value.\n",
    "seed = int(time.time())\n",
    "random.seed(seed)\n",
    "print(f\"Visualizing results with seed {seed}\")\n",
    "\n",
    "fig, ax = plt.subplots(nrows=num_images, ncols=4, figsize=(24, num_images * 4))\n",
    "for i, (image, mask) in enumerate(data_subset):\n",
    "    display_image = rotate_and_flip(image.squeeze())\n",
    "    target_mask = rotate_and_flip(mask).astype(np.uint8)\n",
    "    # Add batch dimension to image and do inference on FP and INT8 models\n",
    "    input_image = np.expand_dims(image, 0)\n",
    "    res_fp = compiled_model_fp([input_image])\n",
    "    res_int8 = compiled_model_int8([input_image])\n",
    "\n",
    "    # Process inference outputs and convert to binary segementation masks\n",
    "    result_mask_fp = sigmoid(res_fp[output_layer_fp]).squeeze().round().astype(np.uint8)\n",
    "    result_mask_int8 = sigmoid(res_int8[output_layer_int8]).squeeze().round().astype(np.uint8)\n",
    "    result_mask_fp = rotate_and_flip(result_mask_fp)\n",
    "    result_mask_int8 = rotate_and_flip(result_mask_int8)\n",
    "\n",
    "    # Display images, annotations, FP32 result and INT8 result\n",
    "    ax[i, 0].imshow(display_image, cmap=colormap)\n",
    "    ax[i, 1].imshow(target_mask, cmap=colormap)\n",
    "    ax[i, 2].imshow(result_mask_fp, cmap=colormap)\n",
    "    ax[i, 3].imshow(result_mask_int8, cmap=colormap)\n",
    "    ax[i, 2].set_title(\"Prediction on FP32 model\")\n",
    "    ax[i, 3].set_title(\"Prediction on INT8 model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Show Live Inference\n",
    "\n",
    "To show live inference on the model in the notebook, we use the asynchronous processing feature of OpenVINO.\n",
    "\n",
    "We use the `show_live_inference` function from [Notebook Utils](../utils/notebook_utils.ipynb) to show live inference. This function uses [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/)'s AsyncPipeline and Model API to perform asynchronous inference. After inference on the specified CT scan has completed, the total time and throughput (fps), including preprocessing and displaying, will be printed.\n",
    "\n",
    "> NOTE: you may experience flickering on Firefox. Please consider using Chrome or Edge to run this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Model and List of Image Files\n",
    "\n",
    "We load the segmentation model to OpenVINO Runtime with `SegmentationModel`, based on the [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/) Model API. This model implementation includes pre and post processing for the model. For `SegmentationModel` this includes the code to create an overlay of the segmentation mask on the original image/frame. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "122wcKhzXn3z",
    "tags": []
   },
   "outputs": [],
   "source": [
    "CASE = 117\n",
    "\n",
    "ie = Core()\n",
    "segmentation_model = SegmentationModel(\n",
    "    ie=ie, model_path=Path(int8_ir_path), sigmoid=True, rotate_and_flip=True\n",
    ")\n",
    "case_path = BASEDIR / f\"case_{CASE:05d}\"\n",
    "image_paths = sorted(case_path.glob(\"imaging_frames/*jpg\"))\n",
    "print(f\"{case_path.name}, {len(image_paths)} images\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Show Inference\n",
    "\n",
    "In the next cell, we run the `show_live_inference` function, which loads the `segmentation_model` to the specified `device` (using caching for faster model loading on GPU devices), loads the images, performs inference, and displays the results on the frames loaded in `images` in real-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "=image_paths": "=image_paths[:5]"
    }
   },
   "outputs": [],
   "source": [
    "# Possible options for device include \"CPU\", \"GPU\", \"AUTO\", \"MULTI:CPU,GPU\"\n",
    "device = \"CPU\"\n",
    "reader = LoadImage(image_only=True, dtype=np.uint8)\n",
    "show_live_inference(\n",
    "    ie=ie, image_paths=image_paths, model=segmentation_model, device=device, reader=reader\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "**OpenVINO**\n",
    "- [NNCF Repository](https://github.com/openvinotoolkit/nncf/)\n",
    "- [Neural Network Compression Framework for fast model inference](https://arxiv.org/abs/2002.08679)\n",
    "- [OpenVINO API Tutorial](../002-openvino-api/002-openvino-api.ipynb)\n",
    "- [OpenVINO PyPI (pip install openvino-dev)](https://pypi.org/project/openvino-dev/)\n",
    "\n",
    "**Kits19 Data**\n",
    "  - [Kits19 Challenge Homepage](https://kits19.grand-challenge.org/)\n",
    "  - [Kits19 Github Repository](https://github.com/neheller/kits19)\n",
    "  - [The KiTS19 Challenge Data: 300 Kidney Tumor Cases with Clinical Context, CT Semantic Segmentations, and Surgical Outcomes](https://arxiv.org/abs/1904.00445)\n",
    "  - [The state of the art in kidney and kidney tumor segmentation in contrast-enhanced CT imaging: Results of the KiTS19 challenge](https://www.sciencedirect.com/science/article/pii/S1361841520301857)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b3d35d6a072dc3a394ff616c9fe1a71d7b246bb5400d8f70c87361b733f2cd1b"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  },
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}