nnU-Net_v1_example.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Background on nnU-Net\n",
        "\n",
        "Here are some useful resources about the nnU-Net system.\n",
        "* [nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation, Isensee et al., 2021](https://www.nature.com/articles/s41592-020-01008-z)\n",
        "* [30-minute overview video of how nnU-Net works](https://youtu.be/K3qRY9Q-BFo)\n",
        "* [nnU-Net Revisited: A Call for Rigorous Validation in 3D Medical Image Segmentation, Isensee et al., 2024](https://arxiv.org/abs/2404.09556)\n",
        "\n",
        "## System Overview\n",
        "<img src='https://media.springernature.com/full/springer-static/image/art%3A10.1038%2Fs41592-020-01008-z/MediaObjects/41592_2020_1008_Fig2_HTML.png?as=webp' alt='nnU-Net System Overview' width='800'/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1sIQWGU-cL7n"
      },
      "source": [
        "# Install nnU-Netv1\n",
        "\n",
        "There will be an issue with argparse. Just click `Restart Session` and then you can go on with the next section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qmi1rZJRYHSX"
      },
      "outputs": [],
      "source": [
        "%pip install git+https://github.com/MIC-DKFZ/nnUNet.git@nnunetv1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9okN5QNwgKZy"
      },
      "source": [
        "# Check available models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gO7cr0omY0fz"
      },
      "outputs": [],
      "source": [
        "!nnUNet_print_available_pretrained_models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FVN6MHwcgNVw"
      },
      "source": [
        "# Set environment variable"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XrMkL86dcB0O"
      },
      "outputs": [],
      "source": [
        "%env RESULTS_FOLDER=RESULTS_FOLDER\n",
        "!mkdir RESULTS_FOLDER\n",
        "!mkdir RESULTS_FOLDER/nnUNet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SLw58aPugWk_"
      },
      "source": [
        "# Download pretrained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9FpHD_gnc7fS"
      },
      "outputs": [],
      "source": [
        "# This is how you would normally do it but this does not work :(\n",
        "# !nnUNet_download_pretrained_model Task004_Hippocampus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A7CkJnDne9CS"
      },
      "outputs": [],
      "source": [
        "!wget -O ./RESULTS_FOLDER/nnUNet/Task004_Hippocampus.zip https://zenodo.org/record/4003545/files/Task004_Hippocampus.zip?download=1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "moosZ2g4fOuY"
      },
      "outputs": [],
      "source": [
        "!unzip ./RESULTS_FOLDER/nnUNet/Task004_Hippocampus.zip -d ./RESULTS_FOLDER/nnUNet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UlHWFYCchF0L"
      },
      "source": [
        "# Download data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dbmBYnSbfUUV"
      },
      "outputs": [],
      "source": [
        "!wget -O Task04_Hippocampus.tar https://drive.google.com/uc?id=1RzPB1_bqzQhlWvU-YGvZzhx2omcDh38C&export=download"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HcBiGl9chaXd"
      },
      "outputs": [],
      "source": [
        "!tar -xvf Task04_Hippocampus.tar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ApRipReh0Bg"
      },
      "source": [
        "# View data sample"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kh1Tbh2LxeKq"
      },
      "outputs": [],
      "source": [
        "import nibabel as nib\n",
        "import matplotlib.pyplot as plt\n",
        "%pip install celluloid\n",
        "from celluloid import Camera\n",
        "from IPython.display import HTML"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uX1NkgL4lTEM"
      },
      "outputs": [],
      "source": [
        "sample_number = '001'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qtnB_e9Th5Ge"
      },
      "outputs": [],
      "source": [
        "image_1 = nib.load(f'./Task04_Hippocampus/imagesTr/hippocampus_{sample_number}.nii.gz')\n",
        "mask_1 = nib.load(f'./Task04_Hippocampus/labelsTr/hippocampus_{sample_number}.nii.gz')\n",
        "\n",
        "image_1_arr = image_1.get_fdata()\n",
        "mask_1_arr = mask_1.get_fdata()\n",
        "\n",
        "print('Image shape:', image_1_arr.shape)\n",
        "print('Mask shape:', mask_1_arr.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wRglaUWWiCwO"
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(7, 5))\n",
        "camera = Camera(fig)\n",
        "axs = axs.ravel()\n",
        "\n",
        "for slice_number in range(image_1_arr.shape[1]):\n",
        "    axs[0].imshow(image_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[0].set_title('Input Image')\n",
        "    axs[0].axis('off')\n",
        "\n",
        "    axs[1].imshow(mask_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[1].set_title('True Mask')\n",
        "    axs[1].axis('off')\n",
        "\n",
        "    axs[2].imshow(image_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[2].imshow(mask_1_arr[:, slice_number, :], cmap='gray', alpha=0.4)\n",
        "    axs[2].set_title('Overlaid Mask')\n",
        "    axs[2].axis('off')\n",
        "\n",
        "    plt.ioff()\n",
        "    camera.snap()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TqEKLnRj4c0B"
      },
      "outputs": [],
      "source": [
        "animation_before = camera.animate()\n",
        "HTML(animation_before.to_html5_video())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W7RzipPPiqIE"
      },
      "outputs": [],
      "source": [
        "!mkdir ./Task04_Hippocampus/SUBSET_IN\n",
        "!mkdir ./Task04_Hippocampus/SUBSET_OUT\n",
        "!cp ./Task04_Hippocampus/imagesTr/hippocampus_{sample_number}.nii.gz ./Task04_Hippocampus/SUBSET_IN\n",
        "!mv ./Task04_Hippocampus/SUBSET_IN/hippocampus_{sample_number}.nii.gz ./Task04_Hippocampus/SUBSET_IN/hippocampus_{sample_number}_0000.nii.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "727QWIffjaNP"
      },
      "source": [
        "# Run segmentation model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0MnM5V6LjXxW"
      },
      "outputs": [],
      "source": [
        "!nnUNet_predict -i ./Task04_Hippocampus/SUBSET_IN/ -o ./Task04_Hippocampus/SUBSET_OUT -t Task004_Hippocampus -m 3d_fullres"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ooHih1NPj5N3"
      },
      "source": [
        "# View results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_W2J-3w1jVAQ"
      },
      "outputs": [],
      "source": [
        "image_1 = nib.load(f'./Task04_Hippocampus/imagesTr/hippocampus_{sample_number}.nii.gz')\n",
        "mask_1 = nib.load(f'./Task04_Hippocampus/labelsTr/hippocampus_{sample_number}.nii.gz')\n",
        "pred_1 = nib.load(f'./Task04_Hippocampus/SUBSET_OUT/hippocampus_{sample_number}.nii.gz')\n",
        "\n",
        "image_1_arr = image_1.get_fdata()\n",
        "mask_1_arr = mask_1.get_fdata()\n",
        "pred_1_arr = pred_1.get_fdata()\n",
        "\n",
        "print('Image shape:', image_1_arr.shape)\n",
        "print('Mask shape:', mask_1_arr.shape)\n",
        "print('Predicted mask shape:', pred_1_arr.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gYYzqDlFkRSP"
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(2, 3, figsize=(7, 4.5))\n",
        "camera = Camera(fig)\n",
        "axs = axs.ravel()\n",
        "\n",
        "for slice_number in range(image_1_arr.shape[1]):\n",
        "    axs[0].imshow(image_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[0].set_title('Input Image')\n",
        "    axs[0].axis('off')\n",
        "\n",
        "    axs[1].imshow(mask_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[1].set_title('True Mask')\n",
        "    axs[1].axis('off')\n",
        "\n",
        "    axs[2].imshow(image_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[2].imshow(mask_1_arr[:, slice_number, :], cmap='gray', alpha=0.4)\n",
        "    axs[2].set_title('Overlaid True Mask')\n",
        "    axs[2].axis('off')\n",
        "\n",
        "    axs[3].axis('off')\n",
        "\n",
        "    axs[4].imshow(pred_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[4].set_title('Predicted Mask')\n",
        "    axs[4].axis('off')\n",
        "\n",
        "    axs[5].imshow(image_1_arr[:, slice_number, :], cmap='gray')\n",
        "    axs[5].imshow(pred_1_arr[:, slice_number, :], cmap='gray', alpha=0.4)\n",
        "    axs[5].set_title('Overlaid Predicted Mask')\n",
        "    axs[5].axis('off')\n",
        "\n",
        "    plt.ioff()\n",
        "    camera.snap()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZNwFEnO-ldSb"
      },
      "outputs": [],
      "source": [
        "animation_after = camera.animate()\n",
        "HTML(animation_after.to_html5_video())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5aALdpuI6n-K"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}