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",
    "- export the trained model to the ONNX.\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.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](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",
    "\n",
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install nibabel"
   ]
  },
  {
   "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"
   ]
  },
  {
   "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 compression.api import Metric\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Dataset and DataModule\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 Lighnting 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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Metric\n",
    "\n",
    "For the metric computation, implement a Binary F1 or a Dice score, following the [Post-Training Optimization Tool](https://docs.openvino.ai/2021.4/pot_README.html) interface in OpenVINO. This enables reusing the same metric class when quantizing the model later in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BinaryF1(Metric):\n",
    "    \"\"\"\n",
    "    Metric to compute a F1/Dice score for binary segmentation. F1 is computed as\n",
    "    (2 * precision * recall) / (precision + recall) where precision is computed as\n",
    "    the ratio of pixels that were correctly predicted as true and all actual true pixels,\n",
    "    and recall as the ratio of pixels that were correctly predicted as true and all\n",
    "    predicted true pixels.\n",
    "\n",
    "    See https://en.wikipedia.org/wiki/F-score\n",
    "    \"\"\"\n",
    "\n",
    "    # Required methods\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self._name = \"F1\"\n",
    "        self.y_true = 0\n",
    "        self.y_pred = 0\n",
    "        self.correct_true = 0\n",
    "\n",
    "    @property\n",
    "    def value(self):\n",
    "        \"\"\"\n",
    "        Returns a metric value for the last model output.\n",
    "        Possible format: {metric_name: [metric_values_per_image]}\n",
    "        \"\"\"\n",
    "        return {self._name: [None]}\n",
    "\n",
    "    @property\n",
    "    def avg_value(self):\n",
    "        \"\"\"\n",
    "        Returns an average metric value for all model outputs.\n",
    "        Possible format: {metric_name: metric_value}\n",
    "        \"\"\"\n",
    "        if self.y_pred == 0:\n",
    "            recall = 0\n",
    "        else:\n",
    "            recall = self.correct_true / self.y_pred\n",
    "\n",
    "        if self.y_true == 0:\n",
    "            precision = 0\n",
    "        else:\n",
    "            precision = self.correct_true / self.y_true\n",
    "\n",
    "        if precision + recall == 0:\n",
    "            f1 = 0\n",
    "        else:\n",
    "            f1 = (2 * precision * recall) / (precision + recall)\n",
    "        return {self._name: f1}\n",
    "\n",
    "    def update(self, output, target):\n",
    "        \"\"\"\n",
    "        :param output: model output\n",
    "        :param target: annotations for model output\n",
    "        \"\"\"\n",
    "        self.y_true += torch.sum(output)\n",
    "        self.y_pred += torch.sum(target)\n",
    "\n",
    "        correct_true = torch.sum((output == target).short() * (target == 1).short()).float()\n",
    "\n",
    "        self.correct_true += correct_true\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"Resets metric\"\"\"\n",
    "        self.y_true = 0\n",
    "        self.y_pred = 0\n",
    "        self.correct_true = 0\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\": \"F1\"}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Lightning Model\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 defined above."
   ]
  },
  {
   "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 = BinaryF1()\n",
    "\n",
    "        self.best_val_dice = 0\n",
    "        self.best_val_epoch = 0\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",
    "        self.metric.update(output=y_pred, target=y_true)\n",
    "\n",
    "        self.log(\"val_loss\", loss)\n",
    "        return {\"val_loss\": loss, \"val_number\": len(output)}\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        val_loss, num_items = 0, 0\n",
    "\n",
    "        for output in outputs:\n",
    "            val_loss += output[\"val_loss\"].sum().item()\n",
    "            num_items += output[\"val_number\"]\n",
    "        mean_val_dice = self.metric.avg_value[\"F1\"]\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## TensorBoard\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",
    "\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",
    "    gpus=1 if USE_CUDA else 0,\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}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Show Inference on the Trained Model\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": [
    "## Export the Trained Model to ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "Path(MODEL_DIR).mkdir(exist_ok=True)\n",
    "onnx_path = Path(MODEL_DIR) / \"unet_kits19.onnx\"\n",
    "dummy_input = torch.randn(1, 1, 512, 512)\n",
    "torch.onnx.export(valmodel, dummy_input, onnx_path, opset_version=10)\n",
    "print(f\"Exported ONNX model to {onnx_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps"
   ]
  },
  {
   "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 with the [Post-Training Optimization Tool](https://docs.openvino.ai/2021.4/pot_README.html) in OpenVINO and show live inference in the notebook."
   ]
  }
 ],
 "metadata": {
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}