playing_with_pytorch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "private_outputs": true,
      "gpuType": "T4",
      "collapsed_sections": [
        "yTWBzkVr_KRO",
        "-HGBrW0YIm94",
        "Sfa4he1uic2L",
        "emU2Bj_aoOsc",
        "GMPLpB3Utsum"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "doj9EzIVIMFv"
      },
      "source": [
        "# Playing with SimpleITK and nnU-Net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qz32_aKqJmIh"
      },
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/woctezuma/playing-with-simpleitk.git\n",
        "\n",
        "%pip install -qq SimpleITK nnunetv2"
      ],
      "metadata": {
        "id": "vL5WJI3NVxIe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3AoxeZYwF46I"
      },
      "source": [
        "## Data nomenclature"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ROOT_FOLDER = '/content/'\n",
        "NNUNET_FOLDER_NAME = f\"{ROOT_FOLDER}nnUNet_base/\"\n",
        "\n",
        "ORIGINAL_DATASET_FNAME = \"Dataset000_Original\"\n",
        "ORIGINAL_DATASET_PATH = f\"{NNUNET_FOLDER_NAME}{ORIGINAL_DATASET_FNAME}/\"\n",
        "\n",
        "DATASET_FNAME = \"Dataset001_Downsampled\"\n",
        "DATASET_PATH = f\"{NNUNET_FOLDER_NAME}{DATASET_FNAME}/\"\n",
        "\n",
        "TRIMMED_DATASET_FNAME = \"Dataset002_Trimmed\"\n",
        "TRIMMED_DATASET_PATH = f\"{NNUNET_FOLDER_NAME}{TRIMMED_DATASET_FNAME}/\""
      ],
      "metadata": {
        "id": "Kri0EHwI296P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8AJwRH-8FmaD"
      },
      "source": [
        "PATIENT_INDICES = [1,2,5,6,8,10,14,16,18,19] + list(range(21, 31))\n",
        "IMAGE_TYPE = '.nii.gz'\n",
        "\n",
        "def get_image_file_name(patient_no, modality_no = 0, data_folder = DATASET_PATH):\n",
        "  image_folder = f'{data_folder}imagesTr/'\n",
        "  return f'{image_folder}patientID{patient_no}_{modality_no:04}{IMAGE_TYPE}'\n",
        "\n",
        "def get_ground_truth_file_name(patient_no, data_folder = DATASET_PATH):\n",
        "  ground_truth_folder = f'{data_folder}labelsTr/'\n",
        "  return f'{ground_truth_folder}patientID{patient_no}{IMAGE_TYPE}'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BMvKoAVsJPx5"
      },
      "source": [
        "## Import data"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can skip all the steps and download the pre-pocessed dataset straight away with:"
      ],
      "metadata": {
        "id": "2su-ec82_ayC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# %cd /content\n",
        "\n",
        "# !curl -OL https://github.com/woctezuma/playing-with-simpleitk/releases/download/original_data/Dataset000_Original.tar.gz\n",
        "# !tar xzf Dataset000_Original.tar.gz\n",
        "\n",
        "# !curl -OL https://github.com/woctezuma/playing-with-simpleitk/releases/download/data/Dataset001_Downsampled.tar.gz\n",
        "# !tar xzf Dataset001_Downsampled.tar.gz"
      ],
      "metadata": {
        "id": "fjuE9MQj_Z1a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you prefer to download the original data and run the steps by yourself, run the cells below."
      ],
      "metadata": {
        "id": "_mXfHD2W_gqN"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download data"
      ],
      "metadata": {
        "id": "yTWBzkVr_KRO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!curl -O https://zenodo.org/records/3431873/files/CHAOS_Train_Sets.zip\n",
        "!unzip -qq CHAOS_Train_Sets.zip"
      ],
      "metadata": {
        "id": "bSD1Q_6Quslt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/playing-with-simpleitk\n",
        "%mv /content/Train_Sets data/\n",
        "!python convert_to_nii.py"
      ],
      "metadata": {
        "id": "qugHxbsFvCtM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%mkdir -p {ORIGINAL_DATASET_PATH}\n",
        "%mv data/output/* {ORIGINAL_DATASET_PATH}"
      ],
      "metadata": {
        "id": "l5aUxh3p1Pg4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-HGBrW0YIm94"
      },
      "source": [
        "### Tests"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fdlLa9Rf-luy"
      },
      "source": [
        "import numpy as np\n",
        "import SimpleITK as sitk\n",
        "\n",
        "patient_no = PATIENT_INDICES[0]\n",
        "\n",
        "for image_name in [\n",
        "    get_image_file_name(patient_no, data_folder=ORIGINAL_DATASET_PATH),\n",
        "    get_ground_truth_file_name(patient_no, data_folder=ORIGINAL_DATASET_PATH),\n",
        "    ]:\n",
        "\n",
        "  image = sitk.ReadImage(image_name)\n",
        "  print(image.GetSize())\n",
        "\n",
        "  v = sitk.GetArrayViewFromImage(image)\n",
        "  print(np.unique(v))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w4Q8PKbvSA3Q"
      },
      "source": [
        "### Down-sample images"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Utils"
      ],
      "metadata": {
        "id": "CiJuvGw36-6Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import SimpleITK as sitk\n",
        "\n",
        "def resample_sitk_image(volume, new_spacing, interpolator=None):\n",
        "    if not interpolator:\n",
        "        interpolator = sitk.sitkLinear\n",
        "        pixelid = volume.GetPixelIDValue()\n",
        "\n",
        "        if pixelid not in [1, 2, 4]:\n",
        "            raise NotImplementedError(\n",
        "                'Set `interpolator` manually, '\n",
        "                'can only infer for 8-bit unsigned or 16, 32-bit signed integers')\n",
        "        if pixelid == 1: #  8-bit unsigned int\n",
        "            interpolator = sitk.sitkNearestNeighbor\n",
        "\n",
        "    # Reference: https://discourse.itk.org/t/resample-volume-to-specific-voxel-spacing-simpleitk/3531/2\n",
        "    original_spacing = volume.GetSpacing()\n",
        "    original_size = volume.GetSize()\n",
        "    new_size = [int(round(osz*ospc/nspc)) for osz,ospc,nspc in zip(original_size, original_spacing, new_spacing)]\n",
        "    return sitk.Resample(volume, new_size, sitk.Transform(), interpolator,\n",
        "                         volume.GetOrigin(), new_spacing, volume.GetDirection(), 0,\n",
        "                         volume.GetPixelID())"
      ],
      "metadata": {
        "id": "xJBpvyy2iRI6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sfa4he1uic2L"
      },
      "source": [
        "#### Down-sample for faster checks."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%mkdir -p {DATASET_PATH}imagesTr\n",
        "%mkdir -p {DATASET_PATH}labelsTr"
      ],
      "metadata": {
        "id": "A01SkrH0ff0n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Eh_I49fVTbpw"
      },
      "source": [
        "num_dim = 3\n",
        "new_spacing = [4]*num_dim\n",
        "\n",
        "for patient_no in PATIENT_INDICES:\n",
        "  for file_name in [\n",
        "      get_ground_truth_file_name(patient_no, data_folder=\"\"),\n",
        "      get_image_file_name(patient_no, data_folder=\"\"),\n",
        "\n",
        "  ]:\n",
        "    print(file_name)\n",
        "\n",
        "    original_image = sitk.ReadImage(f\"{ORIGINAL_DATASET_PATH}{file_name}\")\n",
        "\n",
        "    # Copy information w.r.t. original spacing\n",
        "    original_intensity_image = sitk.ReadImage(get_image_file_name(patient_no, data_folder=ORIGINAL_DATASET_PATH))\n",
        "    original_image.CopyInformation(original_intensity_image)\n",
        "\n",
        "    resampled_image = resample_sitk_image(original_image, new_spacing)\n",
        "\n",
        "    print(original_image.GetSize())\n",
        "    print(resampled_image.GetSize())\n",
        "\n",
        "    sitk.WriteImage(resampled_image, f\"{DATASET_PATH}{file_name}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "emU2Bj_aoOsc"
      },
      "source": [
        "### Binarize label maps\n",
        "\n",
        "nnUNet wants consecutive labels: 0, 1, etc.\n",
        "\n",
        "Originally, the label map in the CHAOS challenge contains only two labels:\n",
        "- 0 (background),\n",
        "- 255 (region of interest).\n",
        "\n",
        "We can simply binarize the label map.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HnOCG_CgoN-h"
      },
      "source": [
        "import SimpleITK as sitk\n",
        "import numpy as np\n",
        "\n",
        "write_to_disk = True\n",
        "\n",
        "for patient_no in PATIENT_INDICES:\n",
        "  print('Patient n°{}'.format(patient_no))\n",
        "\n",
        "  input_image_name = get_ground_truth_file_name(patient_no,\n",
        "                                                data_folder=DATASET_PATH)\n",
        "  input_image = sitk.ReadImage(input_image_name)\n",
        "  print('Image size: {}'.format(input_image.GetSize()))\n",
        "\n",
        "  v = sitk.GetArrayFromImage(input_image)\n",
        "\n",
        "  labels = np.unique(v)\n",
        "\n",
        "  max_val = max(labels)\n",
        "  median_val = max_val/2\n",
        "\n",
        "  # Binarize\n",
        "\n",
        "  binarized_v = np.zeros(v.shape, v.dtype)\n",
        "  binarized_v[v>median_val] = 1\n",
        "\n",
        "  print(f'Labels: {np.unique(v)} -> {median_val} -> {np.unique(binarized_v)}')\n",
        "\n",
        "  output_image = sitk.GetImageFromArray(binarized_v)\n",
        "\n",
        "  # Copy meta-data\n",
        "  output_image.CopyInformation(input_image)\n",
        "  # print('Image size: {}'.format(output_image.GetSize()))\n",
        "\n",
        "  if write_to_disk:\n",
        "    output_image_name = input_image_name\n",
        "    sitk.WriteImage(output_image, output_image_name)\n",
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J9SWg7ClGxXs"
      },
      "source": [
        "## nnUNet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z-2_JJlaLqZD"
      },
      "source": [
        "### Installation"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%env nnUNet_raw=/content/nnUNet_base/\n",
        "%env nnUNet_preprocessed=/content/nnUNet_preprocessed/\n",
        "%env nnUNet_results=/content/nnUNet_results/"
      ],
      "metadata": {
        "id": "l3OXi7a40T5b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jdfq0ZASGDX"
      },
      "source": [
        "## Edit dataset.json at the root of the dataset folder\n",
        "\n",
        "Reference: https://github.com/MIC-DKFZ/nnUNet/tree/master/nnunetv2/dataset_conversion"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xY3xmwWBTcKE"
      },
      "source": [
        "import os\n",
        "import glob\n",
        "import json\n",
        "from collections import OrderedDict\n",
        "\n",
        "def create_dataset_json(dataset_path):\n",
        "  json_dict = OrderedDict()\n",
        "\n",
        "  json_dict['channel_names'] = {\n",
        "      \"0\": \"CT\"\n",
        "  }\n",
        "\n",
        "  json_dict['labels'] = {\n",
        "      \"background\": 0,\n",
        "      \"liver\": 1,\n",
        "  }\n",
        "\n",
        "  file_names = glob.glob(f'{dataset_path}labelsTr/*{IMAGE_TYPE}')\n",
        "  num_patients = len(file_names)\n",
        "\n",
        "  json_dict['numTraining'] = num_patients\n",
        "  json_dict['file_ending'] = IMAGE_TYPE\n",
        "\n",
        "  with open(os.path.join(dataset_path, \"dataset.json\"), 'w') as f:\n",
        "      json.dump(json_dict, f, indent=4, sort_keys=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "create_dataset_json(DATASET_PATH)"
      ],
      "metadata": {
        "id": "MWKKfmqSCr7A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%mkdir -p {TRIMMED_DATASET_PATH}\n",
        "%cp -r {DATASET_PATH}* {TRIMMED_DATASET_PATH}\n",
        "\n",
        "num_patients = 5\n",
        "for i in PATIENT_INDICES[num_patients:]:\n",
        "  %rm {TRIMMED_DATASET_PATH}imagesTr/patientID{i}_0000{IMAGE_TYPE}\n",
        "  %rm {TRIMMED_DATASET_PATH}labelsTr/patientID{i}{IMAGE_TYPE}"
      ],
      "metadata": {
        "id": "Ons8wocxBuOu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "create_dataset_json(TRIMMED_DATASET_PATH)"
      ],
      "metadata": {
        "id": "ABwk-v-RpGM1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GMPLpB3Utsum"
      },
      "source": [
        "## Run pre-processing"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!nnUNetv2_plan_and_preprocess -d 001 --verify_dataset_integrity"
      ],
      "metadata": {
        "id": "JrqUAqQH3GQ8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!nnUNetv2_plan_and_preprocess -d 002 --verify_dataset_integrity"
      ],
      "metadata": {
        "id": "bIjq2vvGHXld"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qg276gkO2CVS"
      },
      "source": [
        "## Training\n",
        "\n",
        "To continue training, append `--c` to the command-line."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!nnUNetv2_train {TRIMMED_DATASET_FNAME} 2d all"
      ],
      "metadata": {
        "id": "doOwmNttnpoV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WC_sIpdZ2E8I"
      },
      "source": [
        "!nnUNetv2_train {TRIMMED_DATASET_FNAME} 3d_fullres all"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_UGqa6I-rLqm"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# If you want to check the results before the end of the training, copy the checkpoint as below.\n",
        "\n",
        "%cd /content/nnUNet_results/Dataset002_Trimmed/nnUNetTrainer__nnUNetPlans__3d_fullres/fold_all/\n",
        "%cp checkpoint_best.pth checkpoint_final.pth\n",
        "\n",
        "%cd /content/nnUNet_results/Dataset002_Trimmed/nnUNetTrainer__nnUNetPlans__2d/fold_all/\n",
        "%cp checkpoint_best.pth checkpoint_final.pth"
      ],
      "metadata": {
        "id": "0X2Ds_sDrj-O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zf9hZKqCrNph"
      },
      "source": [
        "INPUT_FOLDER = f'{DATASET_PATH}imagesTr/'\n",
        "ROOT_OUTPUT_FOLDER = '/content/sample_data/output/'\n",
        "\n",
        "for config in [\"2d\", \"3d_fullres\"]:\n",
        "  print(config)\n",
        "\n",
        "  OUTPUT_FOLDER = f\"{ROOT_OUTPUT_FOLDER}{config}/\"\n",
        "  !mkdir -p {OUTPUT_FOLDER}\n",
        "\n",
        "  !nnUNetv2_predict \\\n",
        "  -i {INPUT_FOLDER} \\\n",
        "  -o {OUTPUT_FOLDER} \\\n",
        "  -c {config} \\\n",
        "  -d {TRIMMED_DATASET_FNAME} \\\n",
        "  -f all"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GItq5R6v5S3W"
      },
      "source": [
        "## Visualize segmentation results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fKPvjGVj4aFG"
      },
      "source": [
        "import numpy as np\n",
        "import SimpleITK as sitk\n",
        "\n",
        "INPUT_FOLDER = f'{DATASET_PATH}labelsTr/'\n",
        "\n",
        "for config in [\"2d\", \"3d_fullres\"]:\n",
        "  print(config)\n",
        "\n",
        "  OUTPUT_FOLDER = f\"{ROOT_OUTPUT_FOLDER}{config}/\"\n",
        "\n",
        "  dice_scores = []\n",
        "\n",
        "  for patient_no in PATIENT_INDICES:\n",
        "    ground_truth_name = f\"{INPUT_FOLDER}patientID{patient_no}.nii.gz\"\n",
        "    prediction_name = f\"{OUTPUT_FOLDER}patientID{patient_no}.nii.gz\"\n",
        "\n",
        "    ground_truth = sitk.ReadImage(ground_truth_name)\n",
        "    prediction = sitk.ReadImage(prediction_name)\n",
        "\n",
        "    v = sitk.GetArrayViewFromImage(ground_truth)\n",
        "    w = sitk.GetArrayViewFromImage(prediction)\n",
        "\n",
        "    dice_score = 2 * np.sum((v*w) > 0) / (np.sum(v>0) + np.sum(w>0))\n",
        "    print(f'Patient n°{patient_no} ; Dice = {dice_score:.3f}')\n",
        "\n",
        "    dice_scores.append(dice_score)\n",
        "\n",
        "  print(f'\\n[training and validation dataset] Average Dice score = {np.mean(dice_scores):.3f} (#patients={len(dice_scores)})')\n",
        "  print(f'NB: training was done with the first {num_patients} patients. Look for possible over-fitting!')\n",
        "\n",
        "  dice_scores_for_training = dice_scores[:num_patients]\n",
        "  print(f'\\n[training dataset] Average Dice score = {np.mean(dice_scores_for_training):.3f} (#patients={len(dice_scores_for_training)})')\n",
        "\n",
        "  dice_scores_for_validation = dice_scores[num_patients:]\n",
        "  print(f'\\n[validation dataset] Average Dice score = {np.mean(dice_scores_for_validation):.3f} (#patients={len(dice_scores_for_validation)})')\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}