segmentation_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep learning volume segmentation\n",
    "\n",
    "Authors: Alessia Saccardo (s212246@dtu.dk) & Felipe Delestro (fima@dtu.dk)\n",
    "\n",
    "This notebook aims to demonstrate the feasibility of implementing a comprehensive deep learning segmentation pipeline solely leveraging the capabilities offered by the qim3d library. Specifically, it will highlight the utilization of the annotation tool and walk through the process of creating and training a Unet model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qim3d\n",
    "import numpy as np \n",
    "import os "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data\n",
    "Qim3d library contains a set of example volumes which can be easily loaded using `qim3d.examples.{volume_name}`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = qim3d.examples.bone_128x128x128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To easily have an insight of how the volume looks like we can interact with it using the `slicer` function from `qim3d`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qim3d.viz.slicer(vol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate dataset for training\n",
    "\n",
    "In order to train the classification model, we need to create a dataset from the volume.\n",
    "\n",
    "This means that we'll need a few slices to be used for `training` and at least one for the `test`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset for training the model is managed by `qim3d.utils.prepare_datasets` and it expects files to follow this structure:\n",
    "\n",
    "<pre>\n",
    "dataset\n",
    "├── test\n",
    "│   ├── images\n",
    "│   │   └── FileA.png\n",
    "│   └── labels\n",
    "│       └── FileA.png\n",
    "└── train\n",
    "    ├── images\n",
    "    │   ├── FileB.png\n",
    "    │   └── FileC.png\n",
    "    └── labels\n",
    "        ├── FileB.png\n",
    "        └── FileC.png\n",
    "\n",
    "\n",
    "</pre>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of slices that will be used\n",
    "ntraining = 4\n",
    "ntest = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell, we get the slice indices, making sure that we're not using the same indices for training and test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a set with all the indices\n",
    "all_idxs = set(range(vol.shape[0]))\n",
    "\n",
    "# Get indices for training data\n",
    "training_idxs = list(np.random.choice(list(all_idxs), size=ntraining))\n",
    "print(f\"Slices for training data...: {training_idxs}\")\n",
    "\n",
    "# Get indices for test data\n",
    "test_idxs = list(np.random.choice(list(all_idxs - set(training_idxs)), size=ntest, replace=False))\n",
    "print(f\"Slices for test data.......: {test_idxs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create folder structure\n",
    "Here we create the necessary directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Base path for the training data\n",
    "base_path = os.path.expanduser(\"~/dataset\")\n",
    "\n",
    "# Create directories\n",
    "print(\"Creating directories:\")\n",
    "for folder_split in [\"train\", \"test\"]:\n",
    "    for folder_type in [\"images\", \"labels\"]:\n",
    "        path = os.path.join(base_path, folder_split, folder_type)\n",
    "        os.makedirs(path, exist_ok=True)\n",
    "        print(path)\n",
    "\n",
    "# Here we have the option to remove any previous files\n",
    "clean_files = True\n",
    "if clean_files:\n",
    "    for root, dirs, files in os.walk(base_path):\n",
    "        for file in files:\n",
    "            file_path = os.path.join(root, file)\n",
    "            os.remove(file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Annotate data\n",
    "\n",
    "The following cell will generate an annotation tool for each slice that was requested. \n",
    "\n",
    "You should use the tool to drawn a mask over the structures you're willing to detect, and press the button `Update` so that the mask is saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation_tools = {}\n",
    "for idx in training_idxs + test_idxs:\n",
    "    if idx in training_idxs:\n",
    "        subset = \"training\"\n",
    "    elif idx in test_idxs:\n",
    "        subset = \"test\"\n",
    "    annotation_tools[idx] = qim3d.gui.annotation_tool.Interface()\n",
    "    annotation_tools[idx].name_suffix = f\"_{idx}\"  \n",
    "    print(f\"Annotation for slice {idx} ({subset})\")\n",
    "    annotation_tools[idx].launch(vol[idx])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting masks from the annotation tool\n",
    "The masks are stored in the annotation tool when the button `Update` is pressed\n",
    "\n",
    "Here we extract the masks and save them to disk, followign the standard needed for the DL model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Saving images and masks to disk\")\n",
    "for idx in training_idxs + test_idxs:\n",
    "    \n",
    "    if idx in training_idxs:\n",
    "        folder_split = \"train\"\n",
    "\n",
    "    elif idx in test_idxs:\n",
    "        folder_split = \"test\"\n",
    "\n",
    "    print (f\"- slice {idx} ({folder_split})\")\n",
    "    mask_dict = annotation_tools[idx].get_result()\n",
    "    mask = list(mask_dict.values())[0]\n",
    "\n",
    "    # Save image\n",
    "    qim3d.io.save(os.path.join(base_path,folder_split,\"images\",f\"{idx}.png\"), vol[idx], replace=True)\n",
    "    # Save label\n",
    "    qim3d.io.save(os.path.join(base_path,folder_split,\"labels\",f\"{idx}.png\"), mask, replace=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Unet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Building and training the Unet model is straightforward using `qim3d`. \n",
    "\n",
    "We first need to instantiate the model by defying its size, which can be either *small*, *medium* or *large*. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining model\n",
    "model = qim3d.models.UNet(size = 'small', dropout = 0.25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we need to decide which type of augumentation to apply to the data. \n",
    "\n",
    "The `qim3d.utils.Augmentation` allows to specify how the images should be reshaped to the appropriate size and the level of transformation to apply respectively to train, test and validation sets. \n",
    "\n",
    "The resize must be choosen between [*crop*, *reshape*, *padding*] and the level of transformation must be chosse between [*None*, *light*, *moderate*, *heavy*]. The user can also specify the mean and standard deviation values for normalizing pixel intensities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining augmentation\n",
    "aug = qim3d.utils.Augmentation(resize = 'crop', transform_train = 'light')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then the datasets and dataloaders are instantiated "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# datasets and dataloaders\n",
    "train_set, val_set, test_set = qim3d.utils.prepare_datasets(path = base_path,\n",
    "                                                            val_fraction = 0.5,\n",
    "                                                            model = model,\n",
    "                                                            augmentation = aug)\n",
    "\n",
    "\n",
    "train_loader, val_loader, test_loader = qim3d.utils.prepare_dataloaders(train_set, \n",
    "                                                                        val_set,\n",
    "                                                                        test_set,\n",
    "                                                                        batch_size = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The hyperparameters are defined using the function `qim3d.models.Hyperparameters` and the model can be easily trained by running the function `qim3d.utils.train_model` which returns also a plot of the losses at the end of the training if the option is selected by the user "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model hyperparameters\n",
    "hyperparameters = qim3d.models.Hyperparameters(model, n_epochs=10, \n",
    "                                               learning_rate = 5e-3, loss_function='DiceCE',\n",
    "                                               weight_decay=1e-3)\n",
    "\n",
    "# training model\n",
    "qim3d.utils.train_model(model, hyperparameters, train_loader, val_loader, plot=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check results\n",
    "\n",
    "To compute the inference step it is just needed to run `qim3d.utils.inference`.\n",
    "\n",
    "The results can be visualize with the function `qim3d.viz.grid_pred` that shows the predicted segmentation along with a comparison between the ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_targ_preds_test = qim3d.utils.inference(test_set, model)\n",
    "qim3d.viz.grid_pred(in_targ_preds_test,alpha=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compute inference on entire volume\n",
    "\n",
    "Given that the input is a volume, the goal is to perform inference on the entire volume rather than individual slices.\n",
    "\n",
    "By using the function `qim3d.utils.volume_inference` it is possible to obtain the segmentation volume output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_vol = qim3d.utils.models.volume_inference(vol, model)\n",
    "qim3d.viz.slicer(inference_vol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize the created mask together with the original volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vol_masked = qim3d.viz.vol_masked(vol, inference_vol, viz_delta=128)\n",
    "qim3d.viz.slicer(vol_masked, cmap=\"PiYG\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}