notebooks/110-ct-segmentation-quantize/data-preparation-ct-scan.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eb3ca42b-49d0-4979-a305-1b047cbb4cd8",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Data Preparation for 2D Medical Imaging\n",
    "\n",
    "## Kidney Segmentation with PyTorch Lightning and OpenVINO™ - Part 1\n",
    "\n",
    "This tutorial is part of a series on how to train, optimize, quantize and show live inference on a medical segmentation model. The goal is to accelerate inference on a kidney segmentation model. The [UNet](https://arxiv.org/abs/1505.04597) model is trained from scratch; the data is from [KiTS19](https://github.com/neheller/kits19).\n",
    "\n",
    " The KiTS19 Nifty images are 3D files. Kidney segmentation is a relatively simple case for neural networks - it is expected that a 2D neural network should work quite well. 2D networks are smaller, and easier to work with than 3D networks, and image data is easier to work with than Nifty files. \n",
    "\n",
    "This first tutorial in the series shows how to:\n",
    " \n",
    "- load Nifty images and get the data as array,\n",
    "- apply windowing to a CT scan to increase contrast,\n",
    "- convert Nifty data to 8-bit images.\n",
    "\n",
    "> Note: This will not result in the best kidney segmentation model. Optimizing the kidney segmentation model is outside the scope of this tutorial. The goal is to have a small model that works reasonably well, as a starting point.\n",
    "\n",
    "\n",
    "All notebooks in this series:\n",
    "\n",
    "- Data Preparation for 2D Segmentation of 3D Medical Data (this notebook)\n",
    "- [Train a 2D-UNet Medical Imaging Model with PyTorch Lightning](pytorch-monai-training.ipynb)\n",
    "- [Convert and Quantize a UNet Model and Show Live Inference](110-ct-segmentation-quantize.ipynb)\n",
    "- [Live Inference and Benchmark CT-scan data](../210-ct-scan-live-inference/210-ct-scan-live-inference.ipynb) \n",
    "\n",
    "\n",
    "## Instructions\n",
    "\n",
    "To install the requirements for running this notebook, follow the instructions in the [README](README.md). \n",
    "\n",
    "Before running this notebook, you must download the dataset from [KiTS19](https://github.com/neheller/kits19) with the code below. First, clone the repository and install the requirements. It is recommended to install the requirements in the `openvino_env` virtual environment.\n",
    "\n",
    "```\n",
    "   git clone https://github.com/neheller/kits19\n",
    "   cd kits19\n",
    "   pip install -r requirements.txt\n",
    "   python -m starter_code.get_imaging\n",
    "```\n",
    "\n",
    "**The code above will take a long time to run. The downloaded content takes up around 21GB of data, and the converted images around 3.5GB**. Downloading the full dataset is only required if you want to train the model yourself. To show quantization on a downloadable subset of the dataset, see the [Convert and Quantize a UNet Model and Show Live Inference](../110-ct-segmentation-quantize/110-ct-segmentation-quantize.ipynb) tutorial.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3651b8da-406e-472f-b16c-020031ce7e34",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e38ce8e-c726-425a-9502-6f8feb889829",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from pathlib import Path\n",
    "from typing import Optional, Tuple\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import nibabel as nib\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c57a3bf-cc43-4463-81de-deea10bb8fa9",
   "metadata": {},
   "source": [
    "## Settings\n",
    "\n",
    "Set `NIFTI_PATH` to the root directory of the Nifty files. This is the directory that contains the `case_00000` to `case_00299` subdirectories that contain `.nii.gz` data. The `FRAMES_DIR` variable should point to the directory to save the frames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd015d1-0c99-45b6-ad74-6f2bc730e9c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Adjust the NIFTI_PATH to a directory that contains the case_00000 to case_00299 files with .nii.gz data.\n",
    "NIFTI_PATH = Path(\"~/kits19/data\").expanduser()\n",
    "# The path to store the extracted frames. By default, this is the kits19_frames folder\n",
    "# at the same level as the data folder in the kits19 root directory.\n",
    "FRAMES_DIR = NIFTI_PATH.with_name(\"kits19_frames\")\n",
    "print(f\"Kits19 frames will be saved to {FRAMES_DIR}\")\n",
    "# Checks if the NIFTI_PATH directory exists and contains the Nifty files.\n",
    "assert NIFTI_PATH.exists(), f\"NIFTI_PATH {NIFTI_PATH} does not exist\"\n",
    "assert len(list(NIFTI_PATH.glob(\"**/*nii.gz\"))) > 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa2578e8-10a4-4eb2-83fc-295f62ce672d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-24T16:49:00.398710Z",
     "iopub.status.busy": "2021-10-24T16:49:00.398558Z",
     "iopub.status.idle": "2021-10-24T16:49:00.400472Z",
     "shell.execute_reply": "2021-10-24T16:49:00.400179Z",
     "shell.execute_reply.started": "2021-10-24T16:49:00.398695Z"
    }
   },
   "source": [
    "## Show One CT-scan\n",
    "\n",
    "Load one CT-scan and visualize the scan and the label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "279bfad2-7f9c-4803-bc7d-c7ce1462cea6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mask_path = NIFTI_PATH / \"case_00002/segmentation.nii.gz\"\n",
    "image_path = mask_path.with_name(\"imaging.nii.gz\")\n",
    "nii_mask = nib.load(mask_path)\n",
    "nii_image = nib.load(image_path)\n",
    "\n",
    "mask_data = nii_mask.get_fdata()\n",
    "image_data = nii_image.get_fdata()\n",
    "\n",
    "# The nii_mask and nii_image are large objects and are no longer required.\n",
    "del nii_mask\n",
    "del nii_image\n",
    "print(image_data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e5293c7-0c96-4006-9cb2-98a625ce85d3",
   "metadata": {},
   "source": [
    "A CT-scan is a 3D image. To visualize it in 2D, create slices or frames. This can be done in three [anatomical planes](https://en.wikipedia.org/wiki/Anatomical_plane): from the front (coronal) , from the side (sagittal), or from the top (axial).\n",
    "\n",
    "Since a kidney is relatively small, most pixels do not contain kidney data. For an indication, check the fraction of pixels that contain kidney data by dividing the number of non-zero pixels by the total number of pixels in the scan. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "774bae8e-7c79-4af3-8017-70bc842b8052",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.count_nonzero(mask_data) / np.size(mask_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d8d2169-f811-4ff2-a220-e1effb25c38c",
   "metadata": {},
   "source": [
    "This number shows that in this particular scan, less than one percent of all pixels in the scan belongs to a kidney. \n",
    "\n",
    "To find the frames with pixels that are annotated as kidney, and show the kidney from all three sides, use the code below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04d55f74-bdb6-4a5c-8b73-a6a407ae12ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.argmax([np.count_nonzero(item) for item in mask_data])\n",
    "x = np.argmax([np.count_nonzero(item) for item in np.transpose(mask_data, (1, 2, 0))])\n",
    "y = np.argmax([np.count_nonzero(item) for item in np.transpose(mask_data, (2, 1, 0))])\n",
    "print(z, x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b54677e-c538-49a6-bec5-c79b164fd140",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def show_slices(z: int, x: int, y: int):\n",
    "    fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(12, 6))\n",
    "    ax[0, 0].imshow(image_data[z], cmap=\"gray\")\n",
    "    ax[1, 0].imshow(mask_data[z], cmap=\"gray\", vmin=0, vmax=2)\n",
    "    ax[0, 1].imshow(image_data[:, x, :], cmap=\"gray\")\n",
    "    ax[1, 1].imshow(mask_data[:, x, :], cmap=\"gray\", vmin=0, vmax=2)\n",
    "    ax[0, 2].imshow(image_data[:, :, y], cmap=\"gray\")\n",
    "    ax[1, 2].imshow(mask_data[:, :, y], cmap=\"gray\", vmin=0, vmax=2);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c980a63-5410-4713-b046-24118853736b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "show_slices(z, x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2907d30b-a713-4a57-a2b7-496de54ff88d",
   "metadata": {},
   "source": [
    "The image above shows three slices, from three different perspectives, in different places in the body. The middle slices shows two colors that indicate a kidney and a tumor, which are annotated in this slice."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ab37f1-700a-4398-bf9e-fa27a09c4e4e",
   "metadata": {},
   "source": [
    "## Apply Window-Level to Increase Contrast\n",
    "\n",
    "A CT-scan data can contain a large range of pixel values. This means that the contrast in the slices shown above is low. While histograms are used to visualize the distribution of the pixel values, a soft tissue window level is applied to increase the contrast for soft tissue in the visualization. For more information on windowing a CT-scan data, see the [Radiopaedia](https://radiopaedia.org/articles/windowing-ct)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0f3c262-3e74-4b33-a123-9228f00dd59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nrows=1, ncols=3, figsize=(15, 4))\n",
    "axs[0].hist(image_data[z, ::])\n",
    "axs[1].hist(image_data[:, x, :])\n",
    "axs[2].hist(image_data[:, :, y]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c886ab-2cbb-461b-91f9-fdeb2abc3c9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (-125,225) is a suitable level for visualizing soft tissue\n",
    "window_start = -125\n",
    "window_end = 225\n",
    "image_data[image_data < window_start] = window_start\n",
    "image_data[image_data > window_end] = window_end\n",
    "show_slices(z, x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c69405ff-a2ae-4c53-abb7-e173170c349b",
   "metadata": {},
   "source": [
    "## Extract Slices from Nifty Data\n",
    "\n",
    "The `save_kits19_frames` function has `mask_path` of one `nii.gz` segmentation mask as argument, and converts the mask and corresponding image to a series of images that are saved as `.jpg` (for images) and `.png` (for masks) files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d03d2021-40c5-47a7-a22e-25158ca81f7c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def save_kits19_frames(\n",
    "    mask_path: Path,\n",
    "    root_dir: os.PathLike,\n",
    "    window_level: Optional[Tuple] = None,\n",
    "    make_binary: bool = True,\n",
    "):\n",
    "    \"\"\"\n",
    "    Save KiTS19 CT-scans to image files, optionally applying a window level.\n",
    "    Images and masks are saved in a subdirectory of root_dir: case_XXXXX.\n",
    "    Images are saved in imaging_frames, masks in segmentation frames, which are\n",
    "    both subdirectories of the case directory.\n",
    "    Frames are taken in the axial direction.\n",
    "\n",
    "    :param mask_path: Path to segmentation.nii.gz file. The corresponding imaging.nii.gz\n",
    "                      file should be in the same directory.\n",
    "    :param root_dir: Root directory to save the generated image files. Will be generated\n",
    "                     if it does not exist\n",
    "    :param window_level: Window level top apply to the data before saving\n",
    "    :param make_binary: If true, create a binary mask where all non-zero pixels are\n",
    "                        considered to be \"foreground\" pixels and get pixel value 1.\n",
    "    \"\"\"\n",
    "    start_time = time.time()\n",
    "\n",
    "    Path(root_dir).mkdir(exist_ok=True)\n",
    "    image_path = mask_path.with_name(\"imaging.nii.gz\")\n",
    "\n",
    "    assert mask_path.exists(), f\"mask_path {mask_path} does not exist!\"\n",
    "    assert image_path.exists(), f\"image_path {image_path} does not exist!\"\n",
    "\n",
    "    nii_mask = nib.load(mask_path)\n",
    "    nii_image = nib.load(image_path)\n",
    "\n",
    "    mask_data = nii_mask.get_fdata()\n",
    "    image_data = nii_image.get_fdata()\n",
    "    assert (\n",
    "        mask_data.shape == image_data.shape\n",
    "    ), f\"Mask and image shape of {mask_path} are not equal\"\n",
    "\n",
    "    del nii_mask\n",
    "    del nii_image\n",
    "\n",
    "    if make_binary:\n",
    "        mask_data[mask_data > 0] = 1\n",
    "\n",
    "    if window_level is not None:\n",
    "        window_start, window_end = window_level\n",
    "        image_data[image_data < window_start] = window_start\n",
    "        image_data[image_data > window_end] = window_end\n",
    "\n",
    "    image_directory = Path(root_dir) / mask_path.parent.name / \"imaging_frames\"\n",
    "    mask_directory = Path(root_dir) / mask_path.parent.name / \"segmentation_frames\"\n",
    "    image_directory.parent.mkdir(exist_ok=True)\n",
    "    image_directory.mkdir(exist_ok=True)\n",
    "    mask_directory.mkdir(exist_ok=True)\n",
    "\n",
    "    for i, (mask_frame, image_frame) in enumerate(zip(mask_data, image_data)):\n",
    "        image_frame = (image_frame - image_frame.min()) / (\n",
    "            image_frame.max() - image_frame.min()\n",
    "        )\n",
    "        image_frame = image_frame * 255\n",
    "        image_frame = image_frame.astype(np.uint8)\n",
    "\n",
    "        new_image_path = str(image_directory / f\"{mask_path.parent.name}_{i:04d}.jpg\")\n",
    "        new_mask_path = str(mask_directory / f\"{mask_path.parent.name}_{i:04d}.png\")\n",
    "        cv2.imwrite(new_image_path, image_frame)\n",
    "        cv2.imwrite(new_mask_path, mask_frame)\n",
    "\n",
    "    end_time = time.time()\n",
    "    print(\n",
    "        f\"Saved {mask_path.parent.name} with {mask_data.shape[0]} frames \"\n",
    "        f\"in {end_time-start_time:.2f} seconds\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94085493-e3ec-427f-8b93-6c3893224485",
   "metadata": {},
   "source": [
    "Run the code below to convert all Nifty files in `NIFTI_PATH` to images that are saved in `FRAMES_DIR`. A soft tissue window level of `(-125,225)` is applied and the segmentation labels are converted to binary kidney segmentations.\n",
    "\n",
    "> Note: Running the following code will take quite a long time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6e9367b-93cd-4ce4-aa9d-51cac71a6b9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_paths = sorted(NIFTI_PATH.glob(\"case_*/segmentation.nii.gz\"))\n",
    "\n",
    "for mask_path in mask_paths:\n",
    "    save_kits19_frames(\n",
    "        mask_path=mask_path,\n",
    "        root_dir=FRAMES_DIR,\n",
    "        window_level=(-125, 225),\n",
    "        make_binary=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d5bc30-f703-4d3e-bf8b-8f7f26093069",
   "metadata": {
    "tags": []
   },
   "source": [
    "## References\n",
    "\n",
    "- [KiTS19 Challenge Homepage](https://kits19.grand-challenge.org/)\n",
    "- [KiTS19 Github Repository](https://github.com/neheller/kits19)\n",
    "- [The KiTS19 Challenge Data: 300 Kidney Tumor Cases with Clinical Context, CT Semantic Segmentations, and Surgical Outcomes](https://arxiv.org/abs/1904.00445)\n",
    "- [The state of the art in kidney and kidney tumor segmentation in contrast-enhanced CT imaging: Results of the KiTS19 challenge](https://www.sciencedirect.com/science/article/pii/S1361841520301857)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}