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 problem 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, please follow the instructions in the README. \n",
    "\n",
    "Before running this notebook, you must download the Kits19 dataset, with code from https://github.com/neheller/kits19.\n",
    "\n",
    "**This code will take a long time to run. The downloaded data takes up around 21GB of space, 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",
    "To do this, first clone the repository and install the requirements. It is recommend to install the requirements in the `openvino_env` virtual environment. In short:\n",
    "\n",
    "```\n",
    "    1. git clone https://github.com/neheller/kits19\n",
    "    2. cd kits19\n",
    "    3. pip install -r requirements.txt\n",
    "    4. python -m starter_code.get_imaging\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 subdirectories `case_00000` to `case_00299` containing _.nii.gz_ data. FRAMES_DIR 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 NIFTI_PATH to directory that contains 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 that the NIFTI_PATH directory exists and contains 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",
    "Let's 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",
    "# nii_mask and nii_image are large objects and 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 this in 2D, we can 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, let's 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",
    "We find frames with pixels that are annotated as kidney, and show the kidney from all three sides"
   ]
  },
  {
   "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, indicating a kidney and a tumor were annotated in this slice."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ab37f1-700a-4398-bf9e-fa27a09c4e4e",
   "metadata": {},
   "source": [
    "## Apply Window-Level to Increase Contrast\n",
    "\n",
    "CT-scan data can contain a large range of pixel values. This means that the contrast in the slices shown above is low. We show histograms to visualize the distribution of the pixel values.  We then apply a soft tissue window level to increase the contrast for soft tissue in the visualization. See [Radiopaedia](https://radiopaedia.org/articles/windowing-ct) for information on windowing CT-scan data."
   ]
  },
  {
   "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 the 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)."
   ]
  },
  {
   "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": [
    "Running the next cell will convert all Nifty files in NIFTI_PATH to images that are saved in FRAMES_DIR. A soft tissue window level of (-125,225) is appplied and the segmentation labels are converted to binary kidney segmentations.\n",
    "\n",
    "Running this cell 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
}