notebooks/SpleenLiverSegmentation/SpleenSeg_Pretrained-4_27.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1452463e",
   "metadata": {},
   "source": [
    "# Spleen Segmentation With NVIDIA Pretrained Model\n",
    "\n",
    "## Overview\n",
    "NVIDIA offers pre-trained models that can be downloaded and applied to various tasks. Here we will do Spleen segmentation using a pretrained model with a Unet architecture found [here](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/monaitoolkit/models/monai_spleen_ct_segmentation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "170d7954",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "+ Learn how to interact with NVIDIA pre trained models\n",
    "+ Learn how to conduct medical image segmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7c2472",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "You only need access to a Sagemaker environment, and make sure you are runnning a GPU. If not, stop your instance and resize to a T4 GPU machine."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93aa3269",
   "metadata": {},
   "source": [
    "## Get Started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66cdaf2",
   "metadata": {},
   "source": [
    "### Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82db674f",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install 'monai[all]'\n",
    "! pip install matplotlib "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb1228b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "540e5d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MONAI version: 0.6.0+38.gf6ad4ba5\n",
    "# Numpy version: 1.21.1\n",
    "# Pytorch version: 1.9.0\n",
    "# Pytorch Ignite version: 0.4.5\n",
    "# Nibabel version: 3.2.1\n",
    "# scikit-image version: 0.18.2\n",
    "# Pillow version: 8.3.1\n",
    "# Tensorboard version: 2.5.0\n",
    "# gdown version: 3.13.0\n",
    "# TorchVision version: 0.10.0+cu111\n",
    "# tqdm version: 4.61.2\n",
    "# lmdb version: 1.2.1\n",
    "# psutil version: 5.8.0\n",
    "# pandas version: 1.3.0\n",
    "# einops version: 0.3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07510582",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import glob\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "#import plotly.graph_objects as go\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from monai.apps import download_and_extract\n",
    "from monai.networks.nets import UNet\n",
    "from monai.networks.layers import Norm\n",
    "from monai.losses import DiceFocalLoss\n",
    "from monai.metrics import DiceMetric\n",
    "from monai.inferers import sliding_window_inference\n",
    "from monai.data import (\n",
    "    LMDBDataset,\n",
    "    DataLoader,\n",
    "    decollate_batch,\n",
    "    ImageDataset,\n",
    "    Dataset\n",
    ")\n",
    "from monai.apps import load_from_mmar\n",
    "from monai.transforms import (\n",
    "    AsDiscrete,\n",
    "    EnsureChannelFirstd,\n",
    "    Compose,\n",
    "    LoadImaged,\n",
    "    ScaleIntensityRanged,\n",
    "    Spacingd,\n",
    "    Orientationd,\n",
    "    CropForegroundd,\n",
    "    RandCropByPosNegLabeld,\n",
    "    RandAffined,\n",
    "    RandRotated,\n",
    "    EnsureType,\n",
    "    EnsureTyped,\n",
    ")\n",
    "from monai.utils import first, set_determinism\n",
    "from monai.apps.mmars import RemoteMMARKeys\n",
    "from monai.config import print_config\n",
    "\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f523cbf",
   "metadata": {},
   "source": [
    "#### Running a pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0be7401d",
   "metadata": {},
   "outputs": [],
   "source": [
    "PRETRAINED = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9f3e5f3",
   "metadata": {},
   "source": [
    "#### Create the directory for storing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311c3282",
   "metadata": {},
   "outputs": [],
   "source": [
    "directory = \"monai_data/\"\n",
    "root_dir = tempfile.mkdtemp() if directory is None else directory\n",
    "print(root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38463a18",
   "metadata": {},
   "source": [
    "#### Download the public dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da7cfede",
   "metadata": {},
   "outputs": [],
   "source": [
    "resource = \"https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar\"\n",
    "md5 = \"410d4a301da4e5b2f6f86ec3ddba524e\"\n",
    "\n",
    "compressed_file = os.path.join(root_dir, \"Task09_Spleen.tar\")\n",
    "download_and_extract(resource, compressed_file, root_dir, md5)\n",
    "data_dir = os.path.join(root_dir, \"Task09_Spleen\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fae7c51b",
   "metadata": {},
   "source": [
    "#### Create Date Dictionaries and separate files from training and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2515b177",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images = sorted(\n",
    "    glob.glob(os.path.join(data_dir, \"imagesTr\", \"*.nii.gz\")))\n",
    "train_labels = sorted(\n",
    "    glob.glob(os.path.join(data_dir, \"labelsTr\", \"*.nii.gz\")))\n",
    "data_dicts = [\n",
    "    {\"image\": image_name, \"label\": label_name}\n",
    "    for image_name, label_name in zip(train_images, train_labels)\n",
    "]\n",
    "train_files, val_files = data_dicts[:-9], data_dicts[-9:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974fc5aa",
   "metadata": {},
   "source": [
    "#### Define your transformations for training and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2357d35d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transforms = Compose( #Transformations for training dataset\n",
    "    [\n",
    "        LoadImaged(keys=[\"image\", \"label\"]), #Load dictionary based images and labels\n",
    "        EnsureChannelFirstd(keys=[\"image\", \"label\"]), #Ensures the first channel of each image is the channel dimension\n",
    "        Spacingd(keys=[\"image\", \"label\"], pixdim=( #Change spacing of voxels to be same across images\n",
    "            1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"), #Correct the orientation of images (Right, Anterior, Superior)\n",
    "        ScaleIntensityRanged( #Scale intensity of all images (For images only and not labels)\n",
    "            keys=[\"image\"], a_min=-57, a_max=164,\n",
    "            b_min=0.0, b_max=1.0, clip=True,\n",
    "        ),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"), #Crop foreground of image\n",
    "        RandCropByPosNegLabeld( #Randomly crop fixed sized region\n",
    "            keys=[\"image\", \"label\"],\n",
    "            label_key=\"label\",\n",
    "            spatial_size=(96, 96, 96),\n",
    "            pos=1,\n",
    "            neg=1,\n",
    "            num_samples=4,\n",
    "            image_key=\"image\",\n",
    "            image_threshold=0,\n",
    "        ),\n",
    "        RandAffined( #Do a random affine transformation with some probability\n",
    "            keys=['image', 'label'],\n",
    "            mode=('bilinear', 'nearest'),\n",
    "            prob=0.5,\n",
    "            spatial_size=(96, 96, 96),\n",
    "            rotate_range=(np.pi/18, np.pi/18, np.pi/5),\n",
    "            scale_range=(0.05, 0.05, 0.05)\n",
    "        ),\n",
    "        EnsureTyped(keys=[\"image\", \"label\"]),\n",
    "    ]\n",
    ")\n",
    "val_transforms = Compose( #Transformations for testing dataset\n",
    "    [\n",
    "        LoadImaged(keys=[\"image\", \"label\"]),\n",
    "        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
    "        Spacingd(keys=[\"image\", \"label\"], pixdim=(\n",
    "            1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "        ScaleIntensityRanged(\n",
    "            keys=[\"image\"], a_min=-57, a_max=164,\n",
    "            b_min=0.0, b_max=1.0, clip=True,\n",
    "        ),\n",
    "        RandRotated(\n",
    "            keys=['image', 'label'],\n",
    "            mode=('bilinear', 'nearest'),\n",
    "            range_x=np.pi/18,\n",
    "            range_y=np.pi/18,\n",
    "            range_z=np.pi/5,\n",
    "            prob=1.0,\n",
    "            padding_mode=('reflection', 'reflection'),\n",
    "        ),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "        EnsureTyped(keys=[\"image\", \"label\"]),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ada5757a",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba3c7695",
   "metadata": {},
   "source": [
    "#### Visualize Image and Label (example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "689eea4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "check_ds = Dataset(data=val_files, transform=val_transforms)\n",
    "check_loader = DataLoader(check_ds, batch_size=1)\n",
    "check_data = first(check_loader)\n",
    "image, label = (check_data[\"image\"][0][0], check_data[\"label\"][0][0])\n",
    "print(f\"image shape: {image.shape}, label shape: {label.shape}\")\n",
    "# plot the slice [:, :, 80]\n",
    "plt.figure(\"check\", (12, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"image\")\n",
    "plt.imshow(image[:, :, 80], cmap=\"gray\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"label\")\n",
    "plt.imshow(label[:, :, 80])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f45ba707",
   "metadata": {},
   "source": [
    "#### Use a dataloader to load files\n",
    "    - Ability to use LMDB (Lightning Memory-Mapped Database)\n",
    "    - Here is where transforms take place and they happen on both images and labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe3285d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = LMDBDataset(data=train_files, transform=train_transforms, cache_dir=root_dir)\n",
    "# initialize cache and print meta information\n",
    "print(train_ds.info())\n",
    "\n",
    "# use batch_size=2 to load images and use RandCropByPosNegLabeld\n",
    "# to generate 2 x 4 images for network training\n",
    "train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=2)\n",
    "\n",
    "# the validation data loader will be created on the fly to ensure \n",
    "# a deterministic validation set for demo purpose.\n",
    "val_ds = LMDBDataset(data=val_files, transform=val_transforms, cache_dir=root_dir)\n",
    "# initialize cache and print meta information\n",
    "print(val_ds.info())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "455cbcdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_ds.info())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a77e7856",
   "metadata": {},
   "source": [
    "#### Now we want to download the pretrained model from NVIDIA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8539fb7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmar = {\n",
    "    RemoteMMARKeys.ID: \"clara_pt_spleen_ct_segmentation_1\",\n",
    "    RemoteMMARKeys.NAME: \"clara_pt_spleen_ct_segmentation\",\n",
    "    RemoteMMARKeys.FILE_TYPE: \"zip\",\n",
    "    RemoteMMARKeys.HASH_TYPE: \"md5\",\n",
    "    RemoteMMARKeys.HASH_VAL: None,\n",
    "    RemoteMMARKeys.MODEL_FILE: os.path.join(\"models\", \"model.pt\"),\n",
    "    RemoteMMARKeys.CONFIG_FILE: os.path.join(\"config\", \"config_train.json\"),\n",
    "    RemoteMMARKeys.VERSION: 2,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de7fb262",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmar['name']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf96f9f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\") #torch.device(\"cpu\")\n",
    "if PRETRAINED:\n",
    "    print(\"using a pretrained model.\")\n",
    "    try: #MONAI=0.8\n",
    "        unet_model = load_from_mmar(\n",
    "            item = mmar['name'], \n",
    "            mmar_dir=root_dir,\n",
    "            map_location=device,\n",
    "            version=mmar['version'],\n",
    "            pretrained=True)\n",
    "    except: #MONAI<0.8\n",
    "        unet_model = load_from_mmar(\n",
    "            mmar, \n",
    "            mmar_dir=root_dir,\n",
    "            map_location=device,\n",
    "            pretrained=True)\n",
    "    model = unet_model\n",
    "else: \n",
    "    print(\"using a randomly init. model.\")\n",
    "    model = UNet(\n",
    "        dimensions=3,\n",
    "        in_channels=1,\n",
    "        out_channels=2,\n",
    "        channels=(16, 32, 64, 128, 256),\n",
    "        strides=(2, 2, 2, 2),\n",
    "        num_res_units=2,\n",
    "        norm=Norm.BATCH,\n",
    "    )\n",
    "\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39910557",
   "metadata": {},
   "source": [
    "### This will be our test file we will view for reference\n",
    " - Here we see how our initial model appears to perform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4be7eb8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_file = data_dicts[20:21]\n",
    "test_ds = LMDBDataset(data=test_file, transform=None, cache_dir=root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2544a774",
   "metadata": {},
   "source": [
    "#### We use a sliding window technique to search the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16fd4e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes=2\n",
    "post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=num_classes)])\n",
    "post_label = Compose([EnsureType(), AsDiscrete(to_onehot=num_classes)])\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for data in DataLoader(test_ds, batch_size=1, num_workers=2):\n",
    "        test_inputs, test_labels = (\n",
    "                        data[\"image\"].to(device),\n",
    "                        data[\"label\"].to(device),\n",
    "                    )\n",
    "        roi_size = (160, 160, 160)\n",
    "        sw_batch_size = 4\n",
    "        test_outputs = sliding_window_inference(\n",
    "            test_inputs, roi_size, sw_batch_size, model, overlap=0.5)\n",
    "        test_outputspre = [post_pred(i) for i in decollate_batch(test_outputs)] # Decollate our results\n",
    "        test_labelspre = [post_label(i) for i in decollate_batch(test_labels)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9782ec96",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Actual Spleen')\n",
    "plt.imshow(test_labelspre[0].cpu().numpy()[1][:,:,200], cmap='Greys_r') #Actual spleen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76cd38e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Pretrained CalculatedSpleen')\n",
    "plt.imshow(test_outputspre[0].cpu().numpy()[1][:,:,200], cmap='Greys_r') #Pretrained model spleen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65c68242",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Differences Between Actual and Model')\n",
    "pretraineddif = test_labelspre[0].cpu().numpy()[1][:,:,200] - test_outputspre[0].cpu().numpy()[1][:,:,200]\n",
    "plt.imshow(pretraineddif, cmap='Greys_r') #Differences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f60e5b5",
   "metadata": {},
   "source": [
    "#### Using just the pretrained model, it appears we are performing pretty well\n",
    "    - We can now continue to train with our data using the NVIDIA models initial weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3e40010",
   "metadata": {},
   "source": [
    "## Training\n",
    "#### Without a GPU, training can take a while\n",
    "#### Recommend skipping next three cells and load in model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8ad6aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_function = DiceFocalLoss(to_onehot_y=True, softmax=True)\n",
    "optimizer = torch.optim.Adam(model.parameters(), 5e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91d340c",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epochs = 25\n",
    "val_interval = 2\n",
    "num_classes = 2\n",
    "best_metric = -1\n",
    "best_metric_epoch = -1\n",
    "epoch_loss_values = []\n",
    "metric_values = []\n",
    "post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=num_classes)])\n",
    "post_label = Compose([EnsureType(), AsDiscrete(to_onehot=num_classes)])\n",
    "dice_metric = DiceMetric(include_background=False, reduction=\"mean\", get_not_nans=False)\n",
    "\n",
    "for epoch in range(max_epochs):\n",
    "    print(\"-\" * 10)\n",
    "    print(f\"epoch {epoch + 1}/{max_epochs}\")\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    step = 0\n",
    "    set_determinism(seed=42)\n",
    "    for batch_data in train_loader:\n",
    "        step += 1\n",
    "        inputs, labels = (\n",
    "            batch_data[\"image\"].to(device),\n",
    "            batch_data[\"label\"].to(device),\n",
    "        )\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = loss_function(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        epoch_loss += loss.item()\n",
    "        print(\n",
    "            f\"{step}/{len(train_ds) // train_loader.batch_size}, \"\n",
    "            f\"train_loss: {loss.item():.4f}\")\n",
    "    epoch_loss /= step\n",
    "    epoch_loss_values.append(epoch_loss)\n",
    "    print(f\"epoch {epoch + 1} average loss: {epoch_loss:.4f}\")\n",
    "\n",
    "    if (epoch + 1) % val_interval == 0:\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            set_determinism(seed=42)\n",
    "            for val_data in DataLoader(val_ds, batch_size=1, num_workers=2):\n",
    "                val_inputs, val_labels = (\n",
    "                    val_data[\"image\"].to(device),\n",
    "                    val_data[\"label\"].to(device),\n",
    "                )\n",
    "                roi_size = (160, 160, 160)\n",
    "                sw_batch_size = 4\n",
    "                val_outputs = sliding_window_inference(\n",
    "                    val_inputs, roi_size, sw_batch_size, model, overlap=0.5)\n",
    "                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]\n",
    "                val_labels = [post_label(i) for i in decollate_batch(val_labels)]\n",
    "                dice_metric(y_pred=val_outputs, y=val_labels)\n",
    "            metric = dice_metric.aggregate().item()\n",
    "            dice_metric.reset()\n",
    "            metric_values.append(metric)\n",
    "            if metric > best_metric:\n",
    "                best_metric = metric\n",
    "                best_metric_epoch = epoch + 1\n",
    "                torch.save(model.state_dict(), os.path.join(\n",
    "                    root_dir, \"Spleen_best_metric_model_pretrained.pth\"))\n",
    "                print(\"saved new best metric model\")\n",
    "            print(\n",
    "                f\"current epoch: {epoch + 1} current mean dice: {metric:.4f}\"\n",
    "                f\"\\nbest mean dice: {best_metric:.4f} \"\n",
    "                f\"at epoch: {best_metric_epoch}\"\n",
    "            )\n",
    "print(\n",
    "    f\"train completed, best_metric: {best_metric:.4f} \"\n",
    "    f\"at epoch: {best_metric_epoch}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cf1fd04",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(\"train\", (12, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"Epoch Average Loss\")\n",
    "x = [i + 1 for i in range(len(epoch_loss_values))]\n",
    "y = epoch_loss_values\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.ylim([0.1, 0.7])\n",
    "plt.plot(x, y)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Val Mean Dice\")\n",
    "x = [val_interval * (i + 1) for i in range(len(metric_values))]\n",
    "y = metric_values\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.ylim([0, 1.0])\n",
    "plt.plot(x, y)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ff0035d",
   "metadata": {},
   "source": [
    "#### The model shows that it has improved fairly quickly over just 25 epochs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0499fa93",
   "metadata": {},
   "source": [
    "## Inference\n",
    "#### Without GPU skip to here to load previously trained best model (without a gpu the training will take a while)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29441405",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load('monai_data/best_metric_model_pretrained.pth'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fab5b4b9",
   "metadata": {},
   "source": [
    "#### With the model loaded let's see if much has changed for our example image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94615f38",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 2\n",
    "post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=num_classes)])\n",
    "post_label = Compose([EnsureType(), AsDiscrete(to_onehot=num_classes)])\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for data in DataLoader(test_ds, batch_size=1, num_workers=2):\n",
    "        test_inputs, test_labels = (\n",
    "                        data[\"image\"].to(device),\n",
    "                        data[\"label\"].to(device),\n",
    "                    )\n",
    "        roi_size = (160, 160, 160)\n",
    "        sw_batch_size = 4\n",
    "        test_outputs = sliding_window_inference(\n",
    "            test_inputs, roi_size, sw_batch_size, model, overlap=0.5)\n",
    "        test_outputsSpl = [post_pred(i) for i in decollate_batch(test_outputs)]\n",
    "        test_labelsSpl = [post_label(i) for i in decollate_batch(test_labels)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3f78dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Trained Calculated Spleen')\n",
    "plt.imshow(test_outputsSpl[0].cpu().numpy()[1][:,:,200], cmap='Greys_r') #Pretrained model spleen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a67f89f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Differences Between Actual and Model')\n",
    "traineddif = test_labelsSpl[0].cpu().numpy()[1][:,:,200] - test_outputsSpl[0].cpu().numpy()[1][:,:,200]\n",
    "plt.imshow(traineddif, cmap='Greys_r') #Differences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "382c7285",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Differences Between The Models')\n",
    "modelsdif = test_outputspre[0].cpu().numpy()[1][:,:,200] - test_outputsSpl[0].cpu().numpy()[1][:,:,200]\n",
    "plt.imshow(traineddif, cmap='Greys_r') #Differences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6606bce2",
   "metadata": {},
   "source": [
    "#### We see not much has changed, which is a good sign for how well the NVIDIA model performs out of the box."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cfd20c6",
   "metadata": {},
   "source": [
    "#### Here is the final image of our Spleen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91e83d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "maskedspleen = np.ma.masked_where(test_outputsSpl[0].cpu().numpy()[1][:,:,200] == 0, test_outputsSpl[0].cpu().numpy()[1][:,:,200])\n",
    "fig = plt.figure(frameon=False, figsize=(10,10))\n",
    "plt.imshow(np.rot90(test_ds[0]['image'][0][:,:,200]), cmap='Greys_r')\n",
    "plt.imshow(np.rot90(maskedspleen), cmap='viridis', alpha=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6030d210",
   "metadata": {},
   "source": [
    "#### Feel free to play around in this notebook or download it and use it where a GPU is accessible"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "896388a1",
   "metadata": {},
   "source": [
    "## Additional Exercise: Use liver segmentation in addition to spleen\n",
    "    - Just need to load liver segmentation from NVIDIA\n",
    "    - While we can't train this model, since we don't have training data, we can use it as a rough estimate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "657e44a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmarliver = {\n",
    "    RemoteMMARKeys.ID: \"clara_pt_liver_and_tumor_ct_segmentation_1\",\n",
    "    RemoteMMARKeys.NAME: \"clara_pt_liver_and_tumor_ct_segmentation\",\n",
    "    RemoteMMARKeys.FILE_TYPE: \"zip\",\n",
    "    RemoteMMARKeys.HASH_TYPE: \"md5\",\n",
    "    RemoteMMARKeys.HASH_VAL: None,\n",
    "    RemoteMMARKeys.MODEL_FILE: os.path.join(\"models\", \"model.pt\"),\n",
    "    RemoteMMARKeys.CONFIG_FILE: os.path.join(\"config\", \"config_train.json\"),\n",
    "    RemoteMMARKeys.VERSION: 1,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6fb0da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "    try: #MONAI=0.8\n",
    "        unet_model = load_from_mmar(\n",
    "            item = mmarliver['name'], \n",
    "            mmar_dir=root_dir,\n",
    "            map_location=device,\n",
    "            version=mmarliver['version'],\n",
    "            pretrained=True)\n",
    "    except: #MONAI<0.8\n",
    "        unet_model = load_from_mmar(\n",
    "            mmarliver, \n",
    "            mmar_dir=root_dir,\n",
    "            map_location=device,\n",
    "            pretrained=True)\n",
    "    model = unet_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55034354",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "print(\"using a pretrained model.\")\n",
    "try: #MONAI=0.8\n",
    "    unet_model = load_from_mmar(\n",
    "        item = mmarliver['name'], \n",
    "        mmar_dir=root_dir,\n",
    "        map_location=device,\n",
    "        version=mmarliver['version'],\n",
    "        pretrained=True)\n",
    "except: #MONAI<0.8\n",
    "    unet_model = load_from_mmar(\n",
    "        mmarliver, \n",
    "        mmar_dir=root_dir,\n",
    "        map_location=device,\n",
    "        pretrained=True)\n",
    "model = unet_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79c1731",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classesP=3\n",
    "num_classesL=2\n",
    "post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=num_classesP)])\n",
    "post_label = Compose([EnsureType(), AsDiscrete(to_onehot=num_classesL)])\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for data in DataLoader(test_ds, batch_size=1, num_workers=2):\n",
    "        test_inputs, test_labels = (\n",
    "                        data[\"image\"].to(device),\n",
    "                        data[\"label\"].to(device),\n",
    "                    )\n",
    "        roi_size = (160, 160, 160)\n",
    "        sw_batch_size = 4\n",
    "        test_outputs = sliding_window_inference(\n",
    "            test_inputs, roi_size, sw_batch_size, model, overlap=0.5)\n",
    "        test_outputsliv = [post_pred(i) for i in decollate_batch(test_outputs)] # Decollate our results\n",
    "        test_labelsliv = [post_label(i) for i in decollate_batch(test_labels)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0956706",
   "metadata": {},
   "outputs": [],
   "source": [
    "sliceval = 215\n",
    "maskedliv = np.ma.masked_where(test_outputsliv[0].cpu().numpy()[1][:,:,sliceval] == 0, test_outputsliv[0].cpu().numpy()[1][:,:,sliceval])\n",
    "maskedspleen = np.ma.masked_where(test_outputsSpl[0].cpu().numpy()[1][:,:,sliceval] == 0, test_outputsSpl[0].cpu().numpy()[1][:,:,sliceval])\n",
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Pretrained Calculated Liver and spleen')\n",
    "plt.imshow(np.rot90(test_ds[0]['image'][0][:,:,sliceval]), cmap='Greys_r')\n",
    "plt.imshow(np.rot90(maskedliv), cmap='cividis', alpha=0.75)\n",
    "plt.imshow(np.rot90(maskedspleen), cmap='viridis', alpha=0.75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bdfdbe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sliceval = 110\n",
    "maskedliv = np.ma.masked_where(test_outputsliv[0].cpu().numpy()[1][:,sliceval,:] == 0, test_outputsliv[0].cpu().numpy()[1][:,sliceval,:])\n",
    "maskedspleen = np.ma.masked_where(test_outputsSpl[0].cpu().numpy()[1][:,sliceval,:] == 0, test_outputsSpl[0].cpu().numpy()[1][:,sliceval,:])\n",
    "fig = plt.figure(frameon=False, figsize=(7,7))\n",
    "plt.title('Pretrained Calculated Liver and Spleen')\n",
    "plt.imshow(np.rot90(test_ds[0]['image'][0][:,sliceval,:]), cmap='Greys_r')\n",
    "plt.imshow(np.rot90(maskedliv), cmap='cividis', alpha=0.75)\n",
    "plt.imshow(np.rot90(maskedspleen), cmap='viridis', alpha=0.75)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51e2dc2b",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "Here you learned how to segment medical images using an NVIDIA pretrained model. To further explore NVIDIA models search the [NGC Catalog](https://catalog.ngc.nvidia.com/models) for 'CT'."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af3e7ddb",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "Make sure you shut down this VM, or delete it if you don't plan to use if further.\n",
    "\n",
    "You can also [delete the buckets](https://docs.aws.amazon.com/AmazonS3/latest/userguide/delete-bucket.html) if you don't want to pay for the data: `aws s3 rb s3://bucket-name --force`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b85a97",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}