04-Train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "corresponding-replication",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "Finally we are going to train our tumor segmentation network. <br />\n",
    "Here we apply some small changes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sacred-bahamas",
   "metadata": {},
   "source": [
    "## Imports:\n",
    "\n",
    "* Pathlib for easy path handling\n",
    "* torch for tensor handling\n",
    "* pytorch lightning for efficient and easy training implementation\n",
    "* ModelCheckpoint and TensorboardLogger for checkpoint saving and logging\n",
    "* imgaug for Data Augmentation\n",
    "* numpy for file loading and array ops\n",
    "* matplotlib for visualizing some images\n",
    "* tqdm for progress par when validating the model\n",
    "* celluloid for easy video generation\n",
    "* Our dataset and model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pregnant-density",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "import imgaug.augmenters as iaa\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "from celluloid import Camera\n",
    "\n",
    "from dataset import LungDataset\n",
    "from model import UNet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brutal-implementation",
   "metadata": {},
   "source": [
    "## Dataset Creation\n",
    "Here we create the train and validation dataset. <br />\n",
    "Additionally we define our data augmentation pipeline.\n",
    "Subsequently the two dataloaders are created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "broadband-arcade",
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = iaa.Sequential([\n",
    "    iaa.Affine(translate_percent=(0.15), \n",
    "               scale=(0.85, 1.15), # zoom in or out\n",
    "               rotate=(-45, 45)#\n",
    "               ),  # rotate up to 45 degrees\n",
    "    iaa.ElasticTransformation()  # Elastic Transformations\n",
    "                ])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "laughing-rehabilitation",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the dataset objects\n",
    "train_path = Path(\"Task06_Lung/Preprocessed/train/\")\n",
    "val_path = Path(\"Task06_Lung/Preprocessed/val/\")\n",
    "\n",
    "train_dataset = LungDataset(train_path, seq)\n",
    "val_dataset = LungDataset(val_path, None)\n",
    "\n",
    "print(f\"There are {len(train_dataset)} train images and {len(val_dataset)} val images\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rubber-theater",
   "metadata": {},
   "source": [
    "## Oversampling to tackle strong class imbalance\n",
    "Lung tumors are often very small, thus we need to make sure that our model does not learn a trivial solution which simply outputs 0 for all voxels.<br />\n",
    "In this notebook we will use oversampling to sample slices which contain a tumor more often.\n",
    "\n",
    "To do so we can use the **WeightedRandomSampler** provided by pytorch which needs a weight for each sample in the dataset.\n",
    "Typically you have one weight for each class, which means that we need to calculate two weights, one for slices without tumors and one for slices with a tumor and create list that assigns each sample from the dataset the corresponding weight"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "amino-nashville",
   "metadata": {},
   "source": [
    "To do so, we at first need to create a list containing only the class labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "comic-fence",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_list = []\n",
    "for _, label in tqdm(train_dataset):\n",
    "    # Check if mask contains a tumorous pixel:\n",
    "    if np.any(label):\n",
    "        target_list.append(1)\n",
    "    else:\n",
    "        target_list.append(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "commercial-antenna",
   "metadata": {},
   "source": [
    "Then we can calculate the weight for each class:\n",
    "To do so, we can simply compute the fraction between the classes and then create the weight list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thirty-connectivity",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniques = np.unique(target_list, return_counts=True)\n",
    "uniques"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bridal-advantage",
   "metadata": {},
   "outputs": [],
   "source": [
    "fraction = uniques[1][0] / uniques[1][1]\n",
    "fraction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "laden-brunswick",
   "metadata": {},
   "source": [
    "Subsequently we assign the weight 1 to each slice without a tumor and ~9 to each slice with a tumor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "planned-stanford",
   "metadata": {},
   "outputs": [],
   "source": [
    "weight_list = []\n",
    "for target in target_list:\n",
    "    if target == 0:\n",
    "        weight_list.append(1)\n",
    "    else:\n",
    "        weight_list.append(fraction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "charitable-detroit",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "weight_list[:50]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "entire-tampa",
   "metadata": {},
   "source": [
    "Finally we create the sampler which we can pass to the DataLoader.\n",
    "**Important:** Only use a sampler for the train loader! We don't want to change the validation data to get a real validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "backed-acoustic",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = torch.utils.data.sampler.WeightedRandomSampler(weight_list, len(weight_list))                     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "final-assets",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8#TODO\n",
    "num_workers = 4# TODO\n",
    "\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,\n",
    "                                           num_workers=num_workers, sampler=sampler)\n",
    "val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spread-growth",
   "metadata": {},
   "source": [
    "We can verify that our sampler works by taking a batch from the train loader and count how many labels are larger than zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unauthorized-wesley",
   "metadata": {},
   "outputs": [],
   "source": [
    "verify_sampler = next(iter(train_loader))  # Take one batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bright-glasgow",
   "metadata": {},
   "outputs": [],
   "source": [
    "(verify_sampler[1][:,0]).sum([1, 2]) > 0  # ~ half the batch size "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mechanical-startup",
   "metadata": {},
   "source": [
    "## Loss\n",
    "\n",
    "As this is a harder task to train you might try different loss functions:\n",
    "We achieved best results by using the Binary Cross Entropy instead of the Dice Loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prime-radio",
   "metadata": {},
   "source": [
    "## Full Segmentation Model\n",
    "\n",
    "We now combine everything into the full pytorch lightning model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "little-dover",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TumorSegmentation(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.model = UNet()\n",
    "        \n",
    "        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)\n",
    "        self.loss_fn = torch.nn.BCEWithLogitsLoss()\n",
    "    \n",
    "    def forward(self, data):\n",
    "        pred = self.model(data)\n",
    "        return pred\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        ct, mask = batch\n",
    "        mask = mask.float()\n",
    "        \n",
    "        pred = self(ct)\n",
    "        loss = self.loss_fn(pred, mask)\n",
    "        \n",
    "        # Logs\n",
    "        self.log(\"Train Dice\", loss)\n",
    "        if batch_idx % 50 == 0:\n",
    "            self.log_images(ct.cpu(), pred.cpu(), mask.cpu(), \"Train\")\n",
    "        return loss\n",
    "    \n",
    "        \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        ct, mask = batch\n",
    "        mask = mask.float()\n",
    "\n",
    "        pred = self(ct)\n",
    "        loss = self.loss_fn(pred, mask)\n",
    "        \n",
    "        # Logs\n",
    "        self.log(\"Val Dice\", loss)\n",
    "        if batch_idx % 50 == 0:\n",
    "            self.log_images(ct.cpu(), pred.cpu(), mask.cpu(), \"Val\")\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    \n",
    "    def log_images(self, ct, pred, mask, name):\n",
    "        \n",
    "        results = []\n",
    "        \n",
    "        pred = pred > 0.5 # As we use the sigomid activation function, we threshold at 0.5\n",
    "        \n",
    "        \n",
    "        fig, axis = plt.subplots(1, 2)\n",
    "        axis[0].imshow(ct[0][0], cmap=\"bone\")\n",
    "        mask_ = np.ma.masked_where(mask[0][0]==0, mask[0][0])\n",
    "        axis[0].imshow(mask_, alpha=0.6)\n",
    "        axis[0].set_title(\"Ground Truth\")\n",
    "        \n",
    "        axis[1].imshow(ct[0][0], cmap=\"bone\")\n",
    "        mask_ = np.ma.masked_where(pred[0][0]==0, pred[0][0])\n",
    "        axis[1].imshow(mask_, alpha=0.6, cmap=\"autumn\")\n",
    "        axis[1].set_title(\"Pred\")\n",
    "\n",
    "        self.logger.experiment.add_figure(f\"{name} Prediction vs Label\", fig, self.global_step)\n",
    "\n",
    "            \n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        #Caution! You always need to return a list here (just pack your optimizer into one :))\n",
    "        return [self.optimizer]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "featured-nature",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instanciate the model\n",
    "model = TumorSegmentation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "diagnostic-particle",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the checkpoint callback\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    monitor='Val Dice',\n",
    "    save_top_k=30,\n",
    "    mode='min')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pacific-inventory",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the trainer\n",
    "# Change the gpus parameter to the number of available gpus in your computer. Use 0 for CPU training\n",
    "\n",
    "gpus = 1 #TODO\n",
    "trainer = pl.Trainer(gpus=gpus, logger=TensorBoardLogger(save_dir=\"./logs\"), log_every_n_steps=1,\n",
    "                     callbacks=checkpoint_callback,\n",
    "                     max_epochs=30)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "handy-testing",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "illegal-harvest",
   "metadata": {},
   "source": [
    "## Evaluation:\n",
    "Let's evaluate the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "challenging-color",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiceScore(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    class to compute the Dice Loss\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, pred, mask):\n",
    "                \n",
    "        #flatten label and prediction tensors\n",
    "        pred = torch.flatten(pred)\n",
    "        mask = torch.flatten(mask)\n",
    "        \n",
    "        counter = (pred * mask).sum()  # Counter       \n",
    "        denum = pred.sum() + mask.sum()  # denominator\n",
    "        dice = (2*counter)/denum\n",
    "        \n",
    "        return dice\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial-strain",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = TumorSegmentation.load_from_checkpoint(\"../checkpoints/epoch=29-step=53759.ckpt\")\n",
    "model.eval();\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imported-neighborhood",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "preds = []\n",
    "labels = []\n",
    "\n",
    "for slice, label in tqdm(val_dataset):\n",
    "    slice = torch.tensor(slice).float().to(device).unsqueeze(0)\n",
    "    with torch.no_grad():\n",
    "        pred = torch.sigmoid(model(slice))\n",
    "    preds.append(pred.cpu().numpy())\n",
    "    labels.append(label)\n",
    "    \n",
    "preds = np.array(preds)\n",
    "labels = np.array(labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "digital-fourth",
   "metadata": {},
   "source": [
    "Compute overall Dice Score: This is not a bad result!\n",
    "Those tumors are extremely small and already some wrongly segmented pixels strongly reduce the Dice Score.\n",
    "The Visualization below demonstrates that!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "breathing-oregon",
   "metadata": {},
   "outputs": [],
   "source": [
    "dice_score = DiceScore()(torch.from_numpy(preds), torch.from_numpy(labels).unsqueeze(0).float())\n",
    "print(f\"The Val Dice Score is: {dice_score}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dense-extent",
   "metadata": {},
   "source": [
    "## Visualization\n",
    "\n",
    "Feel free to play around with the threshold.\n",
    "\n",
    "What happens if you decrease it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "central-check",
   "metadata": {},
   "outputs": [],
   "source": [
    "THRESHOLD = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wound-terrorism",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nibabel as nib\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prime-regard",
   "metadata": {},
   "outputs": [],
   "source": [
    "subject = Path(\"/path/to/Task06_Lung/imagesTs/lung_013.nii.gz\")\n",
    "ct = nib.load(subject).get_fdata() / 3071  # standardize\n",
    "ct = ct[:,:,30:]  # crop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unnecessary-accreditation",
   "metadata": {},
   "outputs": [],
   "source": [
    "segmentation = []\n",
    "label = []\n",
    "scan = []\n",
    "\n",
    "for i in range(ct.shape[-1]):\n",
    "    slice = ct[:,:,i]\n",
    "    slice = cv2.resize(slice, (256, 256))\n",
    "    slice = torch.tensor(slice)\n",
    "    scan.append(slice)\n",
    "    slice = slice.unsqueeze(0).unsqueeze(0).float().to(device)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        pred = model(slice)[0][0].cpu()\n",
    "    pred = pred > THRESHOLD\n",
    "    segmentation.append(pred)\n",
    "    label.append(segmentation)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "covered-elimination",
   "metadata": {},
   "source": [
    "Plotting the predicted segmentation (red)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "desperate-scholar",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "camera = Camera(fig)  # create the camera object from celluloid\n",
    "\n",
    "for i in range(0, len(scan), 2):  # Sagital view. Skip every second slice to reduce the video length\n",
    "    plt.imshow(scan[i], cmap=\"bone\")\n",
    "    mask = np.ma.masked_where(segmentation[i]==0, segmentation[i])\n",
    "    plt.imshow(mask, alpha=0.5, cmap=\"autumn\")  # Use autumn colormap to get red segmentation \n",
    "    \n",
    "    plt.axis(\"off\")\n",
    "    camera.snap()  # Store the current slice\n",
    "animation = camera.animate()  # create the animation\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "periodic-occasion",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "HTML(animation.to_html5_video())  # convert the animation to a video\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "supposed-function",
   "metadata": {},
   "source": [
    "Congratulations! You created a lung cancer segmentation model!"
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}