notebooks/110-ct-segmentation-quantize/pytorch-monai-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a Kidney Segmentation Model with MONAI and PyTorch Lightning\n",
    "\n",
    "## Kidney Segmentation with PyTorch Lightning and OpenVINO™ - Part 2\n",
    "\n",
    "This tutorial is a part of the 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 second tutorial in the series shows how to:\n",
    "\n",
    "- train a 2D segmentation model with MONAI and PyTorch Lightning,\n",
    "- define a Metric that can be used for both training and quantization,\n",
    "- visualize training results,\n",
    "- convert the trained model to OpenVINO Intermediate Representation (IR).\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 (this notebook)\n",
    "- [Convert and Quantize a Segmentation Model and Show Live Inference](110-ct-segmentation-quantize-nncf.ipynb)\n",
    "- [Live Inference and Benchmark CT-scan data](110-ct-scan-live-inference.ipynb) \n",
    "\n",
    "## Instructions\n",
    "\n",
    "This notebook needs the KiTS19 dataset, prepared according to the instructions in the [Data Preparation for 2D Segmentation of 3D Medical Data](data-preparation-ct-scan.ipynb) tutorial. Set `BASEDIR` to the directory that contains the kits19_frames directory in the cell below.\n",
    "\n",
    "To install the requirements for running this notebook, follow the instructions in the [`README.md` file](README.md).\n",
    "\n",
    "> **TIP:** Training the model can take a long time. If you want to run the code with a script instead of in the notebook, you can export the notebook to a Python script, using the following command: `jupyter nbconvert --TagRemovePreprocessor.remove_cell_tags=hide --to script pytorch-monai-training.ipynb`. This will export the code without the visualization and TensorBoard cells.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Table of contents:\n",
    "- [Settings](#Settings)\n",
    "- [Imports](#Imports)\n",
    "- [Define Dataset and DataModule](#Define-Dataset-and-DataModule)\n",
    "- [PyTorch Lightning Model](#PyTorch-Lightning-Model)\n",
    "- [TensorBoard](#TensorBoard)\n",
    "- [Train the Model](#Train-the-Model)\n",
    "- [Show Inference on the Trained Model](#Show-Inference-on-the-Trained-Model)\n",
    "- [Convert PyTorch model to OpenVINO IR](#Convert-PyTorch-model-to-OpenVINO-IR)\n",
    "- [Next Steps](#Next-Steps)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"monai>=0.9.1,<1.0.0\" nibabel pytorch_lightning \"openvino>=2023.1.0\" --extra-index-url https://download.pytorch.org/whl/cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "USE_CUDA = torch.cuda.is_available()": "USE_CUDA = False",
     "~/kits19": "kits19"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "\n",
    "# The directory with the dataset, as prepared in the data preparation notebook.\n",
    "BASEDIR = Path(\"~/kits19/kits19_frames/\").expanduser()\n",
    "# Set to True to use CUDA for training - this requires an NVIDIA GPU, installed CUDA drivers\n",
    "# and a PyTorch version with CUDA enabled.\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "# A directory where the saved model will be stored.\n",
    "MODEL_DIR = \"model\"\n",
    "\n",
    "# Check if BASEDIR contains imaging and segmentation frames.\n",
    "assert len(list(BASEDIR.glob(\"**/segmentation_frames*/*png\"))) > 0\n",
    "assert len(list(BASEDIR.glob(\"**/imaging_frames*/*jpg\"))) > 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import random\n",
    "import time\n",
    "\n",
    "import dateutil.relativedelta\n",
    "import matplotlib.pyplot as plt\n",
    "import monai\n",
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "from monai.data import ImageDataset\n",
    "from monai.transforms import (\n",
    "    AddChannel,\n",
    "    Compose,\n",
    "    EnsureType,\n",
    "    LabelToMask,\n",
    "    RandGaussianNoise,\n",
    "    RandHistogramShift,\n",
    "    RandRotate,\n",
    "    Resize,\n",
    ")\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "from skimage import measure\n",
    "from torch.utils.data import DataLoader\n",
    "from torchmetrics.classification import BinaryF1Score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Dataset and DataModule\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "In this step, a MONAI [`ImageDataset`](https://docs.monai.io/en/latest/data.html#imagedataset) is created to load and transform the data, and a [PyTorch Lightning DataModule](https://pytorch-lightning.readthedocs.io/en/stable/extensions/datamodules.html) for accessing this data during the training. The dataset returns data as a tuple consisting of `(image, mask, image_metadata, mask_metadata)`.\n",
    "\n",
    "Use [MONAI Transforms](https://docs.monai.io/en/latest/transforms.html) to transform and augment the data during the training: randomly rotate the data, add noise, and shift pixel values. MONAI's `ImageDataset` ensures that the random seed for the image and segmentation mask transform are the same. Therefore, for the random rotation transforms, an image and a mask will be rotated in the same way. During validation, make sure that the dimensions and the data type are correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class DataModule(pl.LightningDataModule):\n",
    "    def __init__(self, data_dir: str, batch_size):\n",
    "        super().__init__()\n",
    "        self.batch_size = batch_size\n",
    "        self.data_dir = data_dir\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        random.seed(1.414213)\n",
    "        data_path = Path(self.data_dir)\n",
    "        segs = sorted(data_path.glob(\"case_*/segmentation_frames/*png\"))\n",
    "        images = sorted(data_path.glob(\"case_*/imaging_frames/*jpg\"))\n",
    "        val_indices = [7, 10, 14, 15, 30, 60, 71, 74, 75, 81, 92, 93,\n",
    "                       106, 115, 117, 119, 134, 161, 192, 196, 203]  # fmt: skip\n",
    "        test_indices = [2, 5, 37, 103, 174]\n",
    "        train_indices = set(list(range(210))) - set(val_indices) - set(test_indices)\n",
    "\n",
    "        val_cases = [f\"case_{i:05d}\" for i in val_indices]\n",
    "        train_cases = [f\"case_{i:05d}\" for i in train_indices]\n",
    "\n",
    "        train_segs = [seg for seg in segs if Path(seg).parents[1].name in train_cases]\n",
    "        val_segs = [seg for seg in segs if Path(seg).parents[1].name in val_cases]\n",
    "\n",
    "        train_images = [im for im in images if Path(im).parents[1].name in train_cases]\n",
    "        val_images = [im for im in images if Path(im).parents[1].name in val_cases]\n",
    "\n",
    "        # Define transforms for an image and a segmentation mask.\n",
    "        train_imtrans = Compose(\n",
    "            [\n",
    "                AddChannel(),\n",
    "                Resize((512, 512)),\n",
    "                RandRotate(0.78, prob=0.5),\n",
    "                RandGaussianNoise(prob=0.5),\n",
    "                RandHistogramShift(),\n",
    "                EnsureType(),\n",
    "            ]\n",
    "        )\n",
    "        train_segtrans = Compose(\n",
    "            [\n",
    "                AddChannel(),\n",
    "                Resize((512, 512)),\n",
    "                RandRotate(0.78, prob=0.5),\n",
    "                LabelToMask(select_labels=[1]),\n",
    "                EnsureType(),\n",
    "            ]\n",
    "        )\n",
    "        val_imtrans = Compose([AddChannel(), EnsureType()])\n",
    "        val_segtrans = Compose(\n",
    "            [\n",
    "                AddChannel(),\n",
    "                LabelToMask(select_labels=[1]),\n",
    "                EnsureType(),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.dataset_train = ImageDataset(\n",
    "            image_files=train_images,\n",
    "            seg_files=train_segs,\n",
    "            transform=train_imtrans,\n",
    "            seg_transform=train_segtrans,\n",
    "            image_only=False,\n",
    "            transform_with_metadata=False,\n",
    "        )\n",
    "        self.dataset_val = ImageDataset(\n",
    "            image_files=val_images,\n",
    "            seg_files=val_segs,\n",
    "            transform=val_imtrans,\n",
    "            seg_transform=val_segtrans,\n",
    "            image_only=False,\n",
    "            transform_with_metadata=False,\n",
    "        )\n",
    "\n",
    "        print(f\"Setup train dataset: {len(self.dataset_train)} items\")\n",
    "        print(f\"Setup val dataset: {len(self.dataset_val)} items\")\n",
    "\n",
    "        assert len(self.dataset_train) > 0, \"Train dataset is empty.\"\n",
    "        assert len(self.dataset_val) > 0, \"Val dataset is empty\"\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.dataset_train,\n",
    "            batch_size=self.batch_size,\n",
    "            shuffle=True,\n",
    "            # Set num_workers to 0 to prevent issues in Jupyter. Increase this in the production code.\n",
    "            num_workers=0,\n",
    "            pin_memory=torch.cuda.is_available(),\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.dataset_val,\n",
    "            batch_size=self.batch_size,\n",
    "            num_workers=0,\n",
    "            shuffle=False,\n",
    "            pin_memory=torch.cuda.is_available(),\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return self.val_dataloader()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the code below to check if the dataset looks correct by creating an instance of `DataModule` and plotting a few sample images of the train dataset, after augmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plotdataset(dataset, num=4):\n",
    "    assert num < 16\n",
    "    fig, ax = plt.subplots(3, num, figsize=(3 * num, 9))\n",
    "    dataset_items = random.choices(dataset, k=num)\n",
    "\n",
    "    for i, (image, mask, im_meta, mask_meta) in enumerate(dataset_items):\n",
    "        image = image.long().squeeze(0).cpu()\n",
    "        image_name = Path(im_meta[\"filename_or_obj\"]).stem\n",
    "        mask = mask.long().squeeze(0).cpu()\n",
    "        contours = measure.find_contours(mask.cpu().numpy(), 0.4)\n",
    "\n",
    "        ax[0, i].imshow(image, cmap=\"gray\")\n",
    "        ax[0, i].set_title(image_name)\n",
    "        ax[1, i].imshow(mask, cmap=\"gray\")\n",
    "        for n, contour in enumerate(contours):\n",
    "            ax[2, i].imshow(image)\n",
    "            ax[2, i].plot(contour[:, 1], contour[:, 0], linewidth=2, color=\"red\")\n",
    "\n",
    "    for axi in ax.ravel():\n",
    "        axi.axis(\"off\")\n",
    "    fig.suptitle(f\"Dataset directory: {Path(im_meta['filename_or_obj']).parents[2]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "dm = DataModule(data_dir=BASEDIR, batch_size=8)\n",
    "dm.setup()\n",
    "ds = dm.dataset_train\n",
    "ds.transform.transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "# Run this cell again to see different randomly selected images.\n",
    "plotdataset(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Lightning Model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Create a PyTorch Lightning [`LightningModule`](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html) to train a MONAI [BasicUNet](https://docs.monai.io/en/latest/networks.html#basicunet) model.\n",
    "\n",
    "Use Binary Cross Entropy Loss for the loss function and the Adam Optimizer with the default learning rate of 0.001. The evaluation metric is the Binary F1/Dice score `BinaryF1Score` from `torchmetrics`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class MonaiModel(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self._model = monai.networks.nets.BasicUNet(spatial_dims=2, in_channels=1, out_channels=1)\n",
    "        # https://docs.monai.io/en/latest/highlights.html?deterministic-training-for-reproducibility\n",
    "        monai.utils.set_determinism(seed=2.71828, additional_settings=None)\n",
    "\n",
    "        # https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html\n",
    "        # Set pos_weight to 0.5 to favor precision over recall\n",
    "        self.loss_function = torch.nn.BCEWithLogitsLoss(pos_weight=torch.as_tensor([0.5]))\n",
    "        self.metric = BinaryF1Score()\n",
    "\n",
    "        self.best_val_dice = 0\n",
    "        self.best_val_epoch = 0\n",
    "\n",
    "        self.validation_step_outputs = []\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self._model(x)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self._model.parameters())\n",
    "        return optimizer\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        images, labels, _, _ = batch\n",
    "        labels = labels.float()\n",
    "        output = self.forward(images)\n",
    "        loss = self.loss_function(output, labels)\n",
    "        self.log(\"train_loss\", loss.item())\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        images, labels, _, _ = batch\n",
    "        labels = labels.float()\n",
    "        output = self.forward(images)\n",
    "        loss = self.loss_function(output, labels)\n",
    "\n",
    "        # Compute statistics for metric computation\n",
    "        y_true = labels.long()\n",
    "        y_pred = torch.sigmoid(output).round().long()\n",
    "\n",
    "        self.metric.update(y_true, y_pred)\n",
    "        self.log(\"val_loss\", loss)\n",
    "        self.validation_step_outputs.append({\"val_loss\": loss, \"val_number\": len(output)})\n",
    "        return {\"val_loss\": loss, \"val_number\": len(output)}\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        val_loss, num_items = 0, 0\n",
    "\n",
    "        for output in self.validation_step_outputs:\n",
    "            val_loss += output[\"val_loss\"].sum().item()\n",
    "            num_items += output[\"val_number\"]\n",
    "        mean_val_dice = self.metric.compute()\n",
    "        self.metric.reset()\n",
    "\n",
    "        mean_val_loss = torch.tensor(val_loss / num_items)\n",
    "\n",
    "        self.logger.experiment.add_scalar(\"Loss/Validation\", mean_val_loss, self.current_epoch)\n",
    "        self.logger.experiment.add_scalar(\"F1-score/Validation\", mean_val_dice, self.current_epoch)\n",
    "        self.log(\"F1\", mean_val_dice, prog_bar=True, logger=False)\n",
    "\n",
    "        if mean_val_dice > self.best_val_dice:\n",
    "            self.best_val_dice = mean_val_dice\n",
    "            self.best_val_epoch = self.current_epoch\n",
    "        self.validation_step_outputs.clear()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## TensorBoard\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "During the training loop, loss and metric information is logged to [TensorBoard](https://www.tensorflow.org/tensorboard/get_started). With the TensorBoard Jupyter extension, you can see these in the notebook. \n",
    "\n",
    "When running this cell for the first time, there will be a message that reads: *No dashboards are active for the current data set.* Once training has started, click the *reload* button to see the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ],
    "test_replace": {
     "%tensorboard": "# %tensorboard"
    }
   },
   "outputs": [],
   "source": [
    "# The --bind_all parameter enables you to access TensorBoard if you run this notebook on a remote computer.\n",
    "%tensorboard --logdir tb_logs --bind_all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Create instances of a PyTorch Lightning Model and DataModule, as well as a Logger and a `ModelCheckpoint`. Adjust the batch size for the DataModule if you have a GPU with enough memory. \n",
    "\n",
    "The `TensorBoardLogger` enables logging to TensorBoard; the `ModelCheckpoint` saves the top three models with the best F1 score.\n",
    "\n",
    "For more information on the PyTorch Lightning options, refer to the [PyTorch Lightning documentation](https://pytorch-lightning.readthedocs.io/en/latest/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "batch_size=8": "batch_size=3"
    }
   },
   "outputs": [],
   "source": [
    "model = MonaiModel()\n",
    "data = DataModule(data_dir=BASEDIR, batch_size=8)\n",
    "\n",
    "logger = TensorBoardLogger(\"tb_logs\", name=\"kits19_monai\")\n",
    "checkpoint_callback = ModelCheckpoint(monitor=\"F1\", mode=\"max\", save_top_k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running the next cell will start the training. By default, training runs for 15 epochs on GPU, and for 1 epoch on CPU. Running for 1 epoch will not result in a good model.This is only useful for testing purposes.\n",
    "\n",
    "During the training, a progress bar will be shown with the best F1 score so far. It is possible to stop the training before reaching `max_epochs` by clicking the *stop* button in the Jupyter toolbar at the top of the notebook. The training will stop gracefully and the model that was saved during the training will be used for showing inference and the ONNX export in the next cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "test_replace": {
     "limit_train_batches=0.5": "limit_train_batches=0.1",
     "limit_val_batches=0.5": "limit_val_batches=0.1",
     "max_epochs=15": "max_epochs=1"
    }
   },
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(\n",
    "    max_epochs=15 if USE_CUDA else 1,\n",
    "    accelerator=\"gpu\" if USE_CUDA else \"cpu\",\n",
    "    devices=\"auto\",\n",
    "    logger=logger,\n",
    "    precision=16 if USE_CUDA else 32,\n",
    "    limit_train_batches=0.5,\n",
    "    limit_val_batches=0.5,\n",
    "    callbacks=[checkpoint_callback],\n",
    "    fast_dev_run=False,  # Set to True to quickly test the PyTorch Lightning model.\n",
    ")\n",
    "\n",
    "start = datetime.datetime.now()\n",
    "print(start.strftime(\"%H:%M:%S\"))\n",
    "try:\n",
    "    trainer.fit(model, data)\n",
    "finally:\n",
    "    end = datetime.datetime.now()\n",
    "    print(end.strftime(\"%H:%M:%S\"))\n",
    "    delta = dateutil.relativedelta.relativedelta(end, start)\n",
    "    print(f\"Training duration: {delta.hours:02d}:{delta.minutes:02d}:{delta.seconds:02d}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Show Inference on the Trained Model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The F1 score gives an indication of the quality of the model. However, it is useful to show model outputs on a few random images as well. Therefore, load the model from the best checkpoint, and visualize model outputs on four randomly selected images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "checkpoint_path = checkpoint_callback.best_model_path\n",
    "assert checkpoint_path != \"\", \"No checkpoint saved. Please train the model for at least one epoch.\"\n",
    "print(f\"checkpoint_path: {checkpoint_path}\")\n",
    "\n",
    "best_model = model.load_from_checkpoint(checkpoint_path)\n",
    "valmodel = best_model._model\n",
    "valmodel.eval().cpu();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "dataset_items = random.choices(data.dataset_val, k=4)\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=4, ncols=3, figsize=(24, 16))\n",
    "for i, (image, mask, im_meta, mask_meta) in enumerate(dataset_items):\n",
    "    input_image = image.unsqueeze(0)\n",
    "    image_name = Path(im_meta[\"filename_or_obj\"]).stem\n",
    "    with torch.no_grad():\n",
    "        res = valmodel(input_image)\n",
    "    target_mask = mask.short()[0, ::]\n",
    "    result_mask = torch.sigmoid(res).round().short()[0, 0, ::]\n",
    "\n",
    "    ax[i, 0].imshow(image[0, ::], cmap=\"gray\")\n",
    "    ax[i, 1].imshow(target_mask, cmap=\"gray\")\n",
    "    ax[i, 2].imshow(result_mask, cmap=\"gray\")\n",
    "    ax[i, 0].set_title(image_name)\n",
    "    ax[i, 1].set_title(\"Annotation\")\n",
    "    ax[i, 2].set_title(\"Prediction\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert PyTorch model to OpenVINO IR\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import openvino as ov\n",
    "\n",
    "Path(MODEL_DIR).mkdir(exist_ok=True)\n",
    "ir_path = Path(MODEL_DIR) / \"unet_kits19.xml\"\n",
    "dummy_input = torch.randn(1, 1, 512, 512)\n",
    "ov_model = ov.convert_model(valmodel, example_input=dummy_input)\n",
    "ov.save_model(ov_model, ir_path)\n",
    "print(f\"OpenVINO model saved to {ir_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "Open the [110-ct-segmentation-quantize](../notebooks/110-ct-segmentation-quantize/110-ct-segmentation-quantize.ipynb) notebook to convert the ONNX model to OpenVINO IR, quantize the IR model the model with NNCF with the [Post-training Quantization with NNCF Tool](https://docs.openvino.ai/nightly/basic_quantization_flow.html) API in OpenVINO and show live inference in the notebook."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "cec18e25feb9469b5ff1085a8097bdcd86db6a4ac301d6aeff87d0f3e7ce4ca5"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}