notebooks/segmentation_2d_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Segmentation in 2D using U-Nets with Delira - A very short introduction\n",
    "\n",
    "*Author: Justus Schock, Alexander Moriz* \n",
    "\n",
    "*Date: 17.12.2018*\n",
    " \n",
    "This Example shows how use the U-Net implementation in Delira with PyTorch.\n",
    "\n",
    "Let's first setup the essential hyperparameters. We will use `delira`'s `Parameters`-class for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "logger = None\n",
    "import torch\n",
    "from delira.training import Parameters\n",
    "params = Parameters(fixed_params={\n",
    "    \"model\": {\n",
    "        \"in_channels\": 1, \n",
    "        \"num_classes\": 4\n",
    "    },\n",
    "    \"training\": {\n",
    "        \"batch_size\": 64, # batchsize to use\n",
    "        \"num_epochs\": 10, # number of epochs to train\n",
    "        \"optimizer_cls\": torch.optim.Adam, # optimization algorithm to use\n",
    "        \"optimizer_params\": {'lr': 1e-3}, # initialization parameters for this algorithm\n",
    "        \"losses\": {\"CE\": torch.nn.CrossEntropyLoss()}, # the loss function\n",
    "        \"lr_sched_cls\": None,  # the learning rate scheduling algorithm to use\n",
    "        \"lr_sched_params\": {}, # the corresponding initialization parameters\n",
    "        \"metrics\": {} # and some evaluation metrics\n",
    "    }\n",
    "}) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "Since we did not specify any metric, only the `CrossEntropyLoss` will be calculated for each batch. Since we have a classification task, this should be sufficient. We will train our network with a batchsize of 64 by using `Adam` as optimizer of choice.\n",
    "\n",
    "## Logging and Visualization\n",
    "To get a visualization of our results, we should monitor them somehow. For logging we will use `Visdom`. To start a visdom server you need to execute the following command inside an environment which has visdom installed: \n",
    "```shell\n",
    "visdom -port=9999\n",
    "```\n",
    "This will start a visdom server on port 9999 of your machine and now we can start to configure our logging environment. To view your results you can open [http://localhost:9999](http://localhost:9999) in your browser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from trixi.logger import PytorchVisdomLogger\n",
    "from delira.logging import TrixiHandler\n",
    "import logging\n",
    "\n",
    "logger_kwargs = {\n",
    "    'name': 'ClassificationExampleLogger', # name of our logging environment\n",
    "    'port': 9999 # port on which our visdom server is alive\n",
    "}\n",
    "\n",
    "logger_cls = PytorchVisdomLogger\n",
    "\n",
    "# configure logging module (and root logger)\n",
    "logging.basicConfig(level=logging.INFO,\n",
    "                    handlers=[TrixiHandler(logger_cls, **logger_kwargs)])\n",
    "\n",
    "\n",
    "# derive logger from root logger\n",
    "# (don't do `logger = logging.Logger(\"...\")` since this will create a new\n",
    "# logger which is unrelated to the root logger\n",
    "logger = logging.getLogger(\"Test Logger\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "Since a single visdom server can run multiple environments, we need to specify a (unique) name for our environment and need to tell the logger, on which port it can find the visdom server.\n",
    "\n",
    "## Data Praparation\n",
    "### Loading\n",
    "Next we will create a small train and validation set (in this case they will be the same to show the overfitting capability of the UNet).\n",
    "\n",
    "Our data is a brain MR-image thankfully provided by the [FSL](https://fsl.fmrib.ox.ac.uk/fsl/fslwiki) in their [introduction](http://www.fmrib.ox.ac.uk/primers/intro_primer/ExBox3/IntroBox3.html).\n",
    "\n",
    "We first download the data and extract the T1 image and the corresponding segmentation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from io import BytesIO\n",
    "from zipfile import ZipFile\n",
    "from urllib.request import urlopen\n",
    "\n",
    "resp = urlopen(\"http://www.fmrib.ox.ac.uk/primers/intro_primer/ExBox3/ExBox3.zip\")\n",
    "zipfile = ZipFile(BytesIO(resp.read()))\n",
    "#zipfile_list = zipfile.namelist()\n",
    "#print(zipfile_list)\n",
    "img_file = zipfile.extract(\"ExBox3/T1_brain.nii.gz\")\n",
    "mask_file = zipfile.extract(\"ExBox3/T1_brain_seg.nii.gz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "Now, we load the image and the mask (they are both 3D), convert them to a 32-bit floating point numpy array and ensure, they have the same shape (i.e. that for each voxel in the image, there is a voxel in the mask):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "\n",
    "# load image and mask\n",
    "img = sitk.GetArrayFromImage(sitk.ReadImage(img_file))\n",
    "img = img.astype(np.float32)\n",
    "mask = mask = sitk.GetArrayFromImage(sitk.ReadImage(mask_file))\n",
    "mask = mask.astype(np.float32)\n",
    "\n",
    "assert mask.shape == img.shape\n",
    "print(img.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "By querying the unique values in the mask, we get the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "np.unique(mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "This means, there are 4 classes (background and 3 types of tissue) in our sample.\n",
    "\n",
    "Since we want to do a 2D segmentation, we extract a single slice out of the image and the mask (we choose slice 100 here) and plot it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# load single slice\n",
    "img_slice = img[:, :, 100]\n",
    "mask_slice = mask[:, :, 100]\n",
    "\n",
    "# plot slices\n",
    "plt.figure(1, figsize=(15,10))\n",
    "plt.subplot(121)\n",
    "plt.imshow(img_slice, cmap=\"gray\")\n",
    "plt.colorbar(fraction=0.046, pad=0.04)\n",
    "plt.subplot(122)\n",
    "plt.imshow(mask_slice, cmap=\"gray\")\n",
    "plt.colorbar(fraction=0.046, pad=0.04)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "To load the data, we have to use a `Dataset`. The following defines a very simple dataset, accepting an image slice, a mask slice and the number of samples. It always returns the same sample until `num_samples` samples have been returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from delira.data_loading import AbstractDataset\n",
    "\n",
    "class CustomDataset(AbstractDataset):\n",
    "    def __init__(self, img, mask, num_samples=1000):\n",
    "        super().__init__(None, None, None, None)\n",
    "        self.data = {\"data\": img.reshape(1, *img.shape), \"label\": mask.reshape(1, *mask.shape)}\n",
    "        self.num_samples = num_samples\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        return self.data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.num_samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "Now, we can finally instantiate our datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "dataset_train = CustomDataset(img_slice, mask_slice, num_samples=10000)\n",
    "dataset_val = CustomDataset(img_slice, mask_slice, num_samples=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "### Augmentation\n",
    "For Data-Augmentation we will apply a few transformations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from batchgenerators.transforms import RandomCropTransform, \\\n",
    "                                        ContrastAugmentationTransform, Compose\n",
    "from batchgenerators.transforms.spatial_transforms import ResizeTransform\n",
    "from batchgenerators.transforms.sample_normalization_transforms import MeanStdNormalizationTransform\n",
    "\n",
    "transforms = Compose([\n",
    "    RandomCropTransform(150, label_key=\"label\"), # Perform Random Crops of Size 150 x 150 pixels\n",
    "    ResizeTransform(224, label_key=\"label\"), # Resample these crops back to 224 x 224 pixels\n",
    "    ContrastAugmentationTransform(), # randomly adjust contrast\n",
    "    MeanStdNormalizationTransform(mean=[img_slice.mean()], std=[img_slice.std()])]) # use concrete values since we only have one sample (have to estimate it over whole dataset otherwise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "With these transformations we can now wrap our datasets into datamanagers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from delira.data_loading import DataManager, SequentialSampler, RandomSampler\n",
    "\n",
    "manager_train = DataManager(dataset_train, params.nested_get(\"batch_size\"),\n",
    "                                transforms=transforms,\n",
    "                                sampler_cls=RandomSampler,\n",
    "                                n_process_augmentation=4)\n",
    "\n",
    "manager_val = DataManager(dataset_val, params.nested_get(\"batch_size\"),\n",
    "                              transforms=transforms,\n",
    "                              sampler_cls=SequentialSampler,\n",
    "                              n_process_augmentation=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "## Training\n",
    "\n",
    "After we have done that, we can finally specify our experiment and run it. We will therfore use the already implemented `UNet2dPytorch`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.simplefilter(\"ignore\", UserWarning) # ignore UserWarnings raised by dependency code\n",
    "warnings.simplefilter(\"ignore\", FutureWarning) # ignore FutureWarnings raised by dependency code\n",
    "\n",
    "\n",
    "from delira.training import PyTorchExperiment\n",
    "from delira.training.train_utils import create_optims_default_pytorch\n",
    "from delira.models.segmentation import UNet2dPyTorch\n",
    "\n",
    "if logger is not None:\n",
    "    logger.info(\"Init Experiment\")\n",
    "experiment = PyTorchExperiment(params, UNet2dPyTorch,\n",
    "                               name=\"Segmentation2dExample\",\n",
    "                               save_path=\"./tmp/delira_Experiments\",\n",
    "                               optim_builder=create_optims_default_pytorch,\n",
    "                               gpu_ids=[0], mixed_precision=True)\n",
    "experiment.save()\n",
    "\n",
    "model = experiment.run(manager_train, manager_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "## See Also\n",
    "For a more detailed explanation have a look at \n",
    "* [the introduction tutorial](tutorial_delira.ipynb, \"Introduction\")\n",
    "* [the classification example](classification_pytorch.ipynb, \"Classification\")\n",
    "* [the 3d segmentation example](segmentation_3d_pytorch.ipynb, \"Segmentation 3D\")\n",
    "* [the generative adversarial example](gan_pytorch.ipynb, \"GAN\")"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}