notebook/cancer_indicator/02_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preliminaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../../')\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import Callback\n",
    "from torch.optim import Adam\n",
    "\n",
    "from src.patch_loader import OpenslidePatchLoader\n",
    "from src.cancer_indicator.patch_mask_preparer import PatchMaskPreparer\n",
    "from src.cancer_indicator.cancer_indicator_litmodule import CancerIndicatorLitModule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "img_base_dir = '../../data/panda/images/'\n",
    "mask_base_dir = '../../data/panda/masks_prep/'\n",
    "patch_size = 256\n",
    "batch_size=32\n",
    "n_patches_per_split = None\n",
    "max_epochs=1000\n",
    "samples_per_epoch=None\n",
    "use_latest_weights=True\n",
    "retrain=False\n",
    "\n",
    "# # for debugging, those values reduce training time\n",
    "n_patches_per_split = {\n",
    "    # 'train': 1000,\n",
    "    # 'val': 10,\n",
    "    'test': 500,\n",
    "}\n",
    "samples_per_epoch = 1000\n",
    "# max_epochs=5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.cancer_indicator.masked_patch_datamodule import MaskedPatchDataModule\n",
    "from src.transforms import ImageTransforms\n",
    "\n",
    "# set up the data module that works with the individual images from the PANDA dataset as well as the corresponding patches to extract individual patches with corresponding labels\n",
    "\n",
    "# The prepared masks all have:\n",
    "pmp = PatchMaskPreparer(\n",
    "    image_base_dir=img_base_dir,                # directory that contains the images (*.tiff)\n",
    "    mask_base_dir=mask_base_dir,                # directory that contains the processed masks (*.tiff)\n",
    "    mask_filename_suffix=\"_mask\",               # suffix to add to the image filename to get the mask filename\n",
    "    \n",
    "    # The preprocessed masks have the following attributed\n",
    "    mask_downsample_rate=16,                    # downsampled by a factor of 16 compared to the original image\n",
    "    fg_mask_channel=3,                          # mask channel that contains the foreground mask is at position 3\n",
    "    label_mask_channel=2,                       # mask channel that contains the label mask (cancer yes no) is at position 2\n",
    "    \n",
    "    # The following thresholds define the patch-wise label generation from the segmentation masks \n",
    "    patch_size=patch_size,                      # patch size to use\n",
    "    label_mask_threshold=0.5,                   # minimum relative amount of fg on the patch\n",
    "    fg_mask_threshold=0.1,                      # minimum amount of cancer on the tissue to be considered a cancerous patch\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign (non-meaningful) split according to the filename\n",
    "def get_split(filename):\n",
    "    if filename[2] < 'a':\n",
    "        return 'train'\n",
    "    elif (filename[2] == 'a') and (filename[3] == '9'):\n",
    "        print(filename)\n",
    "        return 'val'\n",
    "    else:\n",
    "        return 'test'\n",
    "\n",
    "experiment_df = (\n",
    "    pd.DataFrame(pmp.get_filenames(), columns=['filename'])\n",
    "    .assign(id_=lambda df_: df_['filename'].apply(lambda x: os.path.splitext(x)[0]))\n",
    "    .set_index(\"id_\")\n",
    "    .assign(split=lambda df_: df_['filename'].apply(get_split))\n",
    ")\n",
    "experiment_df\n",
    "\n",
    "# This Preparer can now be used inside the datamodule that generates (patch, label) pairs:\n",
    "mpdm = MaskedPatchDataModule(\n",
    "    image_base_dir=img_base_dir,                # directory that contains the images (*.tiff)\n",
    "    patch_size=patch_size,                      # patch size to use\n",
    "    experiment_df=experiment_df,                # dataframe that contains the experiment metadata (mainly for the data split and filenames)\n",
    "    patch_mask_preparer=pmp,                    # As previously explained\n",
    "    eval_split=['test'],                        # What split to predict on\n",
    "    \n",
    "    # The following parameters are used to control the number of patches returned by the dataloader\n",
    "    n_patches_per_split=n_patches_per_split,    # the total number of patches to return per split. If None, all patches are returned\n",
    "    samples_per_epoch=samples_per_epoch,        # the number of samples to return per epoch. If None, all patches are returned\n",
    "    batch_size=batch_size,                      # and the batch size\n",
    ")\n",
    "\n",
    "mpdm.setup('fit')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with the datamodule setup, we can now inspect the data\n",
    "\n",
    "# one batch of data is a tuple of patches, labels, and metadata\n",
    "for patches, labels, metadata in mpdm.train_dataloader():\n",
    "\n",
    "    # as an example one can visualize...\n",
    "    sample_idx = 0\n",
    "\n",
    "    # ... the first patch in the batch\n",
    "    patch = patches[sample_idx]\n",
    "    # to plot the patch, the first dimension has to be the channel dimension\n",
    "    patch = patch.permute(1, 2, 0)\n",
    "    # Since the patch is normalized, we have to rescale it between 0 and 1 to plot\n",
    "    patch_rescale = (patch - patch.min()) / (patch.max() - patch.min())\n",
    "    plt.imshow(patch_rescale)\n",
    "\n",
    "    # ... the label of the first patch in the batch\n",
    "    label = labels[sample_idx]\n",
    "    title = \"Cancerous Patch\" if label else \"Non-cancerous Patch\" + \" (normalized)\"\n",
    "    plt.title(title)\n",
    "\n",
    "    # ... and the metadata of the first patch in the batch\n",
    "    metadata_sample = pd.DataFrame(metadata['patch_info']).iloc[sample_idx]\n",
    "    display(metadata_sample.to_frame())\n",
    "\n",
    "    # Only do this once as an example so break out of the loop \n",
    "    break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# It is also possible to check the number of batches\n",
    "print(len(mpdm.train_dataloader()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.architecture import get_cancer_indicator_net\n",
    "\n",
    "# The cancer indicator network consists of an effnet backbone and a target head that predicts the probability of a patch being cancerous\n",
    "cancer_indicator_net = get_cancer_indicator_net(\n",
    "    effnet_kwargs=dict(\n",
    "        pretrained=True,\n",
    "    ),\n",
    "    target_kwargs=dict(\n",
    "        drop_rate=0.5,\n",
    "        num_classes=2,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For training, the CancerIndicatorLitModule wraps the network and the optimizer\n",
    "litmodule = CancerIndicatorLitModule(\n",
    "    net=cancer_indicator_net,\n",
    "    optimizer=Adam,\n",
    "    n_classes=2, # only predict cancerous / non-cancerous\n",
    "    use_weights=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MetricTracker(Callback):\n",
    "\n",
    "  def __init__(self):\n",
    "    self.outputs = []\n",
    "\n",
    "  def on_validation_batch_end(self, trainer, module, outputs, *args):\n",
    "    self.outputs.append(outputs)\n",
    "\n",
    "# This tracker is used for the learning curve below\n",
    "cb = MetricTracker()\n",
    "\n",
    "trainer = Trainer(\n",
    "    accelerator=\"cpu\",\n",
    "    max_epochs=max_epochs,\n",
    "    num_sanity_val_steps=0,\n",
    "    callbacks=[cb],\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_latest_version_checkpoint_path(logs_path='lightning_logs'):\n",
    "    \"\"\"Returns the path to the latest version of the checkpoint file in the given directory.\n",
    "    \"\"\"\n",
    "\n",
    "    import os\n",
    "    from glob import glob\n",
    "    import pandas as pd\n",
    "\n",
    "    paths = glob(\n",
    "        os.path.join(logs_path, '**/*.ckpt'), \n",
    "        recursive=True,\n",
    "    )\n",
    "\n",
    "    if len(paths) == 0:\n",
    "        print(f\"No checkpoint files found in {logs_path}\")\n",
    "        return None\n",
    "\n",
    "    return (\n",
    "        pd.DataFrame(paths, columns=['filepath'])\n",
    "        .assign(version=lambda x: x['filepath'].str.extract(r'version_(\\d+)').astype(float))\n",
    "        .sort_values('version')\n",
    "        .dropna()\n",
    "        .iloc[-1]\n",
    "    ).filepath\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = find_latest_version_checkpoint_path()\n",
    "\n",
    "if retrain:\n",
    "    if use_latest_weights:\n",
    "        # get newest checkpoint path from lightning_logs\n",
    "        print(f\"Loading weights from {path} to retrain.\")\n",
    "        trainer.fit(model=litmodule, datamodule=mpdm, ckpt_path=path)\n",
    "    else:\n",
    "        print(f\"Retraining from scratch.\")\n",
    "        trainer.fit(model=litmodule, datamodule=mpdm)\n",
    "else:\n",
    "    if path is None:\n",
    "        print(\"No checkpoint found. Retraining from scratch.\")\n",
    "        trainer.fit(model=litmodule, datamodule=mpdm)\n",
    "    else:\n",
    "        litmodule.load_from_checkpoint(\n",
    "            path,\n",
    "            net=cancer_indicator_net,\n",
    "        )\n",
    "        print(f\"Loading weights from {path}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(cb.outputs) > 0:\n",
    "    # If training was performed, the learning curve can be shown\n",
    "    losses = pd.Series([out['loss'] for out in cb.outputs], name='loss', dtype='float32')\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    ax.plot(\n",
    "        losses,\n",
    "        marker='x',\n",
    "        color='tab:blue'\n",
    "    )\n",
    "    ax.set_xlabel('Step')\n",
    "    ax.set_ylabel('Loss')\n",
    "\n",
    "    # only whole numbers on x axis\n",
    "    ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For prediction, prepare the dataloader and litmodule\n",
    "mpdm.setup('predict')\n",
    "litmodule.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the predictions and convert them to a pandas dataframe along with the provided metadata\n",
    "predictions = CancerIndicatorLitModule.get_predictions_df(\n",
    "    trainer.predict(model=litmodule, datamodule=mpdm)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_predictions(predictions, n, highest=True):\n",
    "    \"\"\"Extract the highest or lowest predictions from the predictions dataframe.\n",
    "    \"\"\"\n",
    "\n",
    "    # only keep 5 patches per slide\n",
    "    result = (\n",
    "        predictions\n",
    "        .groupby('filename', as_index=False)\n",
    "        .sample(5, random_state=42)\n",
    "        .sort_values('prediction', ascending=not highest)\n",
    "        .head(n)\n",
    "    )\n",
    "\n",
    "    if highest:\n",
    "        result = result.iloc[::-1]\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "# Finally, plot examples of the lowest and highest predictions indicated by the prediction color (green=low, red=high). \n",
    "# The border around each patch indicates a negative (green) or positive label (red)\n",
    "\n",
    "nrows = 4\n",
    "samples_per_group = 8\n",
    "ncols= 4 # middle column is left empty\n",
    "fig, axs = plt.subplots(nrows=nrows, ncols=ncols, dpi=300)\n",
    "\n",
    "for i, ax_col in enumerate(axs.T):\n",
    "\n",
    "    for j, ax in enumerate(ax_col):\n",
    "\n",
    "        if i < 2:\n",
    "            is_high_predictions = False\n",
    "            tmp_df = extract_predictions(predictions, samples_per_group, highest=False)\n",
    "        else:\n",
    "            is_high_predictions = True\n",
    "            tmp_df = extract_predictions(predictions, samples_per_group, highest=True)\n",
    "\n",
    "        tmp_pred = tmp_df.iloc[i%2::2].iloc[j]\n",
    "\n",
    "        OpenslidePatchLoader(\n",
    "            filepath=os.path.join(img_base_dir, tmp_pred.filename),\n",
    "            patch_size=patch_size,\n",
    "            channel_is_first_axis=False,\n",
    "        ).plot_patch(tmp_pred.row, tmp_pred.col, ax=ax)    \n",
    "\n",
    "        ax.set_title(f\"{tmp_pred.prediction:.2f}\", fontsize=8, color='tab:red' if is_high_predictions else 'tab:green')\n",
    "\n",
    "        tmp_color = 'tab:red' if tmp_pred.label == 1 else 'tab:green'\n",
    "        for spine in ax.spines.values():\n",
    "            spine.set_edgecolor(tmp_color)\n",
    "            spine.set_linewidth(2)\n",
    "    \n",
    "fig.suptitle(\"Low predictions\" + \" \" * 30 + \"High predictions\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('output/ci_prediction_examples.png', dpi=300, bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env-pcai",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}