notebooks/W&B_SegViz_3d_LSPK_combined_anon.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XVHtEqy6Pg0j"
      },
      "source": [
        "## Setup environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B3eIjeNzPg0j",
        "tags": []
      },
      "outputs": [],
      "source": [
        "from monai.utils import first, set_determinism\n",
        "from monai.transforms import (\n",
        "    AsDiscrete,\n",
        "    AsDiscreted,\n",
        "    EnsureChannelFirstd,\n",
        "    Compose,\n",
        "    CropForegroundd,\n",
        "    LoadImaged,\n",
        "    Orientationd,\n",
        "    RandCropByPosNegLabeld,\n",
        "    SaveImaged,\n",
        "    RandAffined,\n",
        "    ScaleIntensityRanged,\n",
        "    Spacingd,\n",
        "    Invertd,\n",
        "    Resized,\n",
        ")\n",
        "from monai.handlers.utils import from_engine\n",
        "from monai.networks.nets import UNet, BasicUNet\n",
        "from monai.networks.layers import Norm\n",
        "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
        "\n",
        "from monai.metrics import DiceMetric, ROCAUCMetric, MSEMetric\n",
        "from monai.networks.utils import copy_model_state\n",
        "from monai.optimizers import generate_param_groups\n",
        "from monai.losses import DiceLoss\n",
        "from monai.inferers import sliding_window_inference\n",
        "from monai.data import CacheDataset, DataLoader, Dataset, decollate_batch\n",
        "from monai.config import print_config\n",
        "from monai.apps import download_and_extract\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "import tempfile\n",
        "import shutil\n",
        "import os\n",
        "import random\n",
        "import glob\n",
        "import numpy as np\n",
        "import wandb\n",
        "import copy\n",
        "import nibabel as nib\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QovUUOzklUx1"
      },
      "outputs": [],
      "source": [
        "torch.backends.cudnn.benchmark = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Ciz6cVnWYWl"
      },
      "outputs": [],
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G8kPOLVIPg0k"
      },
      "source": [
        "## Setup imports"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r7hO81pBPg0l"
      },
      "source": [
        "## Setup data directory\n",
        "\n",
        "You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.  \n",
        "This allows you to save results and reuse downloads.  \n",
        "If not specified a temporary directory will be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8h5uHSK9Pg0l",
        "tags": []
      },
      "outputs": [],
      "source": [
        "directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n",
        "root_dir = '/content'\n",
        "print(root_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a540_rw-Pg0l"
      },
      "source": [
        "## Download dataset\n",
        "\n",
        "Downloads and extracts the dataset.  \n",
        "The dataset comes from http://medicaldecathlon.com/."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AI6oo1sBPg0m"
      },
      "source": [
        "## Set MSD Spleen dataset path"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "data_dir_spleen = '/content/LSPK_combined'"
      ],
      "metadata": {
        "id": "sbtYof-UFO4G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LGNMKPSzPg0m"
      },
      "outputs": [],
      "source": [
        "train_images = sorted(\n",
        "    glob.glob(os.path.join(data_dir_spleen, \"imagesTr\", \"*.nii.gz\")))\n",
        "train_labels = sorted(\n",
        "    glob.glob(os.path.join(data_dir_spleen, \"labelsTr\", \"*.nii.gz\")))\n",
        "data_dicts_spleen = [\n",
        "    {\"image\": image_name, \"label\": label_name}\n",
        "    for image_name, label_name in zip(train_images, train_labels)\n",
        "]\n",
        "\n",
        "random.seed(108)\n",
        "random.shuffle(data_dicts_spleen)\n",
        "\n",
        "train_files_spleen, val_files_spleen = data_dicts_spleen[:-132], data_dicts_spleen[-132:]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yMK0hwYxPg0m"
      },
      "source": [
        "## Set deterministic training for reproducibility"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_VfXGpEAPg0m"
      },
      "outputs": [],
      "source": [
        "set_determinism(seed=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_y_VXPUFPg0m"
      },
      "source": [
        "## Setup transforms for training and validation\n",
        "\n",
        "Here we use several transforms to augment the dataset:\n",
        "1. `LoadImaged` loads the spleen CT images and labels from NIfTI format files.\n",
        "1. `EnsureChannelFirstd` ensures the original data to construct \"channel first\" shape.\n",
        "1. `Orientationd` unifies the data orientation based on the affine matrix.\n",
        "1. `Spacingd` adjusts the spacing by `pixdim=(1.5, 1.5, 2.)` based on the affine matrix.\n",
        "1. `ScaleIntensityRanged` extracts intensity range [-57, 164] and scales to [0, 1].\n",
        "1. `CropForegroundd` removes all zero borders to focus on the valid body area of the images and labels.\n",
        "1. `RandCropByPosNegLabeld` randomly crop patch samples from big image based on pos / neg ratio.  \n",
        "The image centers of negative samples must be in valid body area.\n",
        "1. `RandAffined` efficiently performs `rotate`, `scale`, `shear`, `translate`, etc. together based on PyTorch affine transform."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SUNPr9J5Pg0m"
      },
      "outputs": [],
      "source": [
        "train_transforms_spleen = Compose(\n",
        "    [\n",
        "        LoadImaged(keys=[\"image\", \"label\"]),\n",
        "        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
        "        ScaleIntensityRanged(\n",
        "            keys=[\"image\"], a_min=-200, a_max=304,\n",
        "            b_min=0.0, b_max=1.0, clip=True,\n",
        "        ),\n",
        "        Resized(keys=[\"image\"], spatial_size=(256,256,128)),   \n",
        "        Resized(keys=[\"label\"], spatial_size=(256,256,128), mode='nearest'),   \n",
        "        # CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
        "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
        "        Spacingd(keys=[\"image\", \"label\"], pixdim=(\n",
        "            1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n",
        "        RandCropByPosNegLabeld(\n",
        "            keys=[\"image\", \"label\"],\n",
        "            label_key=\"label\",\n",
        "            spatial_size=(128,128,32),\n",
        "            pos=1,\n",
        "            neg=1,\n",
        "            num_samples=4,\n",
        "            image_key=\"image\",\n",
        "            image_threshold=0,\n",
        "        ),\n",
        "    ]\n",
        ")\n",
        "val_transform_spleen = Compose(\n",
        "    [\n",
        "        LoadImaged(keys=[\"image\", \"label\"]),\n",
        "        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
        "        ScaleIntensityRanged(\n",
        "            keys=[\"image\"], a_min=-200, a_max=304,\n",
        "            b_min=0.0, b_max=1.0, clip=True,\n",
        "        ),\n",
        "        Resized(keys=[\"image\"], spatial_size=(256,256,128)),   \n",
        "        Resized(keys=[\"label\"], spatial_size=(256,256,128), mode='nearest'), \n",
        "        #CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
        "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
        "        Spacingd(keys=[\"image\", \"label\"], pixdim=(\n",
        "            1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n",
        "    ]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMmJmbGvPg0n"
      },
      "source": [
        "## Check transforms in DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WFK7donFPg0n",
        "tags": []
      },
      "outputs": [],
      "source": [
        "check_ds = Dataset(data=val_files_spleen, transform=val_transform_spleen)\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",
        "print(np.unique(label))\n",
        "\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[:, :, 16], cmap=\"gray\")\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.title(\"label\")\n",
        "plt.imshow(label[:, :, 16])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OTxtC2ycjSpH"
      },
      "outputs": [],
      "source": [
        "config = {\n",
        "    # data\n",
        "    \"cache_rate_spleen\": 1.0,\n",
        "    \"num_workers\": 14,\n",
        "\n",
        "\n",
        "    # train settings\n",
        "    \"train_batch_size\": 2,\n",
        "    \"val_batch_size\": 1,\n",
        "    \"learning_rate\": 1e-4,\n",
        "    \"max_epochs\": 500,\n",
        "    \"val_interval\": 2, # check validation score after n epochs\n",
        "    \"lr_scheduler\": \"cosine_decay\", # just to keep track\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    # Unet model (you can even use nested dictionary and this will be handled by W&B automatically)\n",
        "    \"model_type_spleen\": \"unet\", # just to keep track\n",
        "    \"model_params_spleen\": dict(spatial_dims=3,\n",
        "                  in_channels=1,\n",
        "                  out_channels=5,\n",
        "                  channels=(16, 32, 64, 128, 256),\n",
        "                  strides=(2, 2, 2, 2),\n",
        "                  num_res_units=2,\n",
        "                  norm=Norm.BATCH,),\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a1vDaw4aPg0n"
      },
      "source": [
        "## Define CacheDataset and DataLoader for training and validation\n",
        "\n",
        "Here we use CacheDataset to accelerate training and validation process, it's 10x faster than the regular Dataset.  \n",
        "To achieve best performance, set `cache_rate=1.0` to cache all the data, if memory is not enough, set lower value.  \n",
        "Users can also set `cache_num` instead of `cache_rate`, will use the minimum value of the 2 settings.  \n",
        "And set `num_workers` to enable multi-threads during caching.  \n",
        "If want to to try the regular Dataset, just change to use the commented code below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DiKh-dJvj0Q2"
      },
      "outputs": [],
      "source": [
        "train_ds_spleen = CacheDataset(\n",
        "    data=train_files_spleen, transform=train_transforms_spleen,\n",
        "    cache_rate=config['cache_rate_spleen'], num_workers=config['num_workers'])\n",
        "# train_ds_spleen = Dataset(data=train_files_spleen, transform=train_transforms_spleen)\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_spleen = DataLoader(train_ds_spleen, batch_size=config['train_batch_size'], shuffle=True, num_workers=config['num_workers'])\n",
        "\n",
        "val_ds_spleen = CacheDataset(\n",
        "    data=val_files_spleen, transform=val_transform_spleen, cache_rate=config['cache_rate_spleen'], num_workers=config['num_workers'])\n",
        "# val_ds = Dataset(data=val_files, transform=val_transforms)\n",
        "val_loader_spleen = DataLoader(val_ds_spleen, batch_size=config['val_batch_size'], num_workers=config['num_workers'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DzC3W1bCPg0n"
      },
      "source": [
        "## Create Model, Loss, Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XzqANOZHkzbm"
      },
      "outputs": [],
      "source": [
        "# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer\n",
        "model_spleen = UNet(**config['model_params_spleen']).to(device)\n",
        "\n",
        "loss_function_spleen = DiceLoss(to_onehot_y=True, softmax=True)\n",
        "\n",
        "\n",
        "optimizer_spleen = torch.optim.Adam(model_spleen.parameters(), lr=config['learning_rate'])\n",
        "\n",
        "\n",
        "dice_metric_spleen = DiceMetric(include_background=False, reduction=\"mean\")\n",
        "\n",
        "\n",
        "scheduler_spleen = CosineAnnealingLR(optimizer_spleen, T_max=config['max_epochs'], eta_min=1e-9)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nmOtJ4uXPg0n"
      },
      "source": [
        "## Execute a typical PyTorch training process"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qoyyI_QmPg0n",
        "scrolled": true,
        "tags": []
      },
      "outputs": [],
      "source": [
        "# 🐝 initialize a wandb run\n",
        "wandb.init(\n",
        "    project=\"SegViz_LSPK_combined\",\n",
        "    config=config\n",
        ")\n",
        "\n",
        "# 🐝 log gradients of the model to wandb\n",
        "wandb.watch(model_spleen, log_freq=100)\n",
        "\n",
        "max_epochs = 500\n",
        "val_interval = 2\n",
        "best_metric = -1\n",
        "best_metric_spleen = -1\n",
        "best_metric_liver = -1\n",
        "best_metric_pan = -1\n",
        "\n",
        "best_metric_epoch = -1\n",
        "best_metric_epoch_spleen = -1\n",
        "best_metric_epoch_liver = -1\n",
        "best_metric_epoch_pan = -1\n",
        "\n",
        "epoch_loss_values = []\n",
        "metric_values = []\n",
        "\n",
        "epoch_loss_values_spleen = []\n",
        "metric_values_spleen = []\n",
        "\n",
        "epoch_loss_values_liver = []\n",
        "metric_values_liver = []\n",
        "\n",
        "epoch_loss_values_pan = []\n",
        "metric_values_pan = []\n",
        "\n",
        "post_pred_spleen = Compose([AsDiscrete(argmax=True, to_onehot=5)])\n",
        "post_label_spleen = Compose([AsDiscrete(to_onehot=5)])\n",
        "\n",
        "for epoch in range(max_epochs):\n",
        "  epoch_loss_spleen = 0\n",
        "\n",
        "  step_0 = 0\n",
        "  step_1 = 0\n",
        "  \n",
        "  # For one epoch\n",
        "  \n",
        "  print(\"-\" * 10)\n",
        "  print(f\"epoch {epoch + 1}/{max_epochs}\")\n",
        "  \n",
        "  model_spleen.train()\n",
        "  \n",
        "  # One forward pass of the spleen data through the spleen UNet\n",
        "  for batch_data_spleen in train_loader_spleen:\n",
        "      step_0 += 1\n",
        "      inputs_spleen, labels_spleen = (\n",
        "          batch_data_spleen[\"image\"].to(device),\n",
        "          batch_data_spleen[\"label\"].to(device),\n",
        "      )\n",
        "      optimizer_spleen.zero_grad()\n",
        "      outputs_spleen = model_spleen(inputs_spleen)\n",
        "      \n",
        "      if \"liver\" in batch_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "        labels_spleen[labels_spleen != 2] = 0\n",
        "      elif \"pancreas\" in batch_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "        labels_spleen[labels_spleen != 3] = 0\n",
        "      elif \"spleen\" in batch_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "        labels_spleen[labels_spleen != 1] = 0\n",
        "      elif \"case\" in batch_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "        labels_spleen[labels_spleen != 4] = 0\n",
        "\n",
        "      loss_spleen = loss_function_spleen(outputs_spleen, labels_spleen)\n",
        "      loss_spleen.backward()\n",
        "      \n",
        "      optimizer_spleen.step()\n",
        "      epoch_loss_spleen += loss_spleen.item()\n",
        "      print(\n",
        "          f\"{step_0}/{len(train_ds_spleen) // train_loader_spleen.batch_size}, \"\n",
        "          f\"train_loss: {loss_spleen.item():.4f}\")\n",
        "      wandb.log({\"train/loss combined\": loss_spleen.item()})\n",
        "  epoch_loss_spleen /= step_0\n",
        "  epoch_loss_values_spleen.append(epoch_loss_spleen)\n",
        "  print(f\"epoch {epoch + 1} average loss combined: {epoch_loss_spleen:.4f}\")\n",
        "  \n",
        "  scheduler_spleen.step()\n",
        "\n",
        "  wandb.log({\"train/loss_epoch comb\": epoch_loss_spleen})\n",
        "    # 🐝 log learning rate after each epoch to wandb\n",
        "  wandb.log({\"learning_rate comb\": scheduler_spleen.get_last_lr()[0]})\n",
        "\n",
        "  # Validation \n",
        "  if (epoch + 1) % val_interval == 0:\n",
        "    model_spleen.eval()\n",
        "    with torch.no_grad():\n",
        "\n",
        "        # Validation forward spleen\n",
        "        for val_data_spleen in val_loader_spleen:\n",
        "            val_inputs_spleen, val_labels_spleen = (\n",
        "                val_data_spleen[\"image\"].to(device),\n",
        "                val_data_spleen[\"label\"].to(device),\n",
        "            )\n",
        "            roi_size = (160, 160, 160)\n",
        "            sw_batch_size = 4\n",
        "            \n",
        "            if \"liver\" in val_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "              val_labels_spleen[val_labels_spleen != 2] = 0\n",
        "            elif \"pancreas\" in val_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "              val_labels_spleen[val_labels_spleen != 3] = 0\n",
        "            elif \"spleen\" in val_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "              val_labels_spleen[val_labels_spleen != 1] = 0\n",
        "            elif \"case\" in val_data_spleen['image_meta_dict']['filename_or_obj'][0].split('/')[-1]:\n",
        "              val_labels_spleen[val_labels_spleen != 4] = 0\n",
        "\n",
        "            \n",
        "            val_outputs_spleen = sliding_window_inference(\n",
        "                val_inputs_spleen, roi_size, sw_batch_size, model_spleen)\n",
        "            val_outputs_spleen = [post_pred_spleen(i) for i in decollate_batch(val_outputs_spleen)]\n",
        "            val_labels_spleen = [post_label_spleen(i) for i in decollate_batch(val_labels_spleen)]\n",
        "            # compute metric for current iteration\n",
        "            dice_metric_spleen(y_pred=val_outputs_spleen, y=val_labels_spleen)\n",
        "\n",
        "        # aggregate the final mean dice result\n",
        "        metric_spleen = dice_metric_spleen.aggregate().item()\n",
        "        wandb.log({\"val/dice_metric combined\": metric_spleen})\n",
        "        scheduler_spleen.step(metric_spleen)\n",
        "        # reset the status for next validation round\n",
        "        dice_metric_spleen.reset()\n",
        "\n",
        "        metric_values_spleen.append(metric_spleen)\n",
        "        if metric_spleen > best_metric_spleen:\n",
        "            best_metric_spleen = metric_spleen\n",
        "            best_metric_epoch_spleen = epoch + 1\n",
        "            torch.save(model_spleen.state_dict(), os.path.join(\n",
        "                root_dir, \"best_metric_model_LSPK_combined_trial2.pth\"))\n",
        "            print(\"saved new best metric model for combined dataset\")\n",
        "        print(\n",
        "            f\"current epoch: {epoch + 1} current mean dice for combined: {metric_spleen:.4f}\"\n",
        "            f\"\\nbest mean dice for combined: {best_metric_spleen:.4f} \"\n",
        "            f\"at epoch: {best_metric_epoch_spleen}\"\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oEuTuXmaPg0o"
      },
      "source": [
        "## Plot the loss and metric"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csK69MIDPg0o"
      },
      "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_spleen))]\n",
        "y = epoch_loss_values_spleen\n",
        "plt.xlabel(\"epoch\")\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_spleen))]\n",
        "y = metric_values_spleen\n",
        "plt.xlabel(\"epoch\")\n",
        "plt.plot(x, y)\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "pytorch",
      "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"
    },
    "vscode": {
      "interpreter": {
        "hash": "0ca5fdb1364b7850985fadde1ce5cdfbf8a36ab9b8a2920002462139b38e4875"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}