notebooks/tile_level_datamodule.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e931c96-1fe9-4807-93c2-e134ece52cfc",
   "metadata": {},
   "source": [
    "How to use the TileLevelDataModule\n",
    "----------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "64a76b07-1d53-406d-802a-a722af4e5221",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import time\n",
    "import os\n",
    "from torch import nn\n",
    "import cv2\n",
    "import albumentations as A\n",
    "from torchvision import transforms as T\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import slide_tools\n",
    "\n",
    "cv2.setNumThreads(1)  # Otherwise cv2 will use all threads with every worker -> bad bad bad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b22f48a7-9a45-420f-94e8-5e518b6727ee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>SlideNr</th>\n",
       "      <th>HRD(BRCA1)</th>\n",
       "      <th>label</th>\n",
       "      <th>slide</th>\n",
       "      <th>annotation</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10</td>\n",
       "      <td>0</td>\n",
       "      <td>ago-tr1/labels/10.json</td>\n",
       "      <td>ago-tr1/slides/10.svs</td>\n",
       "      <td>ago-tr1/annotations/10.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>100</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/100.json</td>\n",
       "      <td>ago-tr1/slides/100.svs</td>\n",
       "      <td>ago-tr1/annotations/100.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>101</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/101.json</td>\n",
       "      <td>ago-tr1/slides/101.svs</td>\n",
       "      <td>ago-tr1/annotations/101.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>103</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/103.json</td>\n",
       "      <td>ago-tr1/slides/103.svs</td>\n",
       "      <td>ago-tr1/annotations/103.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>106</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/106.json</td>\n",
       "      <td>ago-tr1/slides/106.svs</td>\n",
       "      <td>ago-tr1/annotations/106.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>127</th>\n",
       "      <td>84</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/84.json</td>\n",
       "      <td>ago-tr1/slides/84.svs</td>\n",
       "      <td>ago-tr1/annotations/84.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>128</th>\n",
       "      <td>85</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/85.json</td>\n",
       "      <td>ago-tr1/slides/85.svs</td>\n",
       "      <td>ago-tr1/annotations/85.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>129</th>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>ago-tr1/labels/9.json</td>\n",
       "      <td>ago-tr1/slides/9.svs</td>\n",
       "      <td>ago-tr1/annotations/9.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>130</th>\n",
       "      <td>92</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/92.json</td>\n",
       "      <td>ago-tr1/slides/92.svs</td>\n",
       "      <td>ago-tr1/annotations/92.geojson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>131</th>\n",
       "      <td>99</td>\n",
       "      <td>1</td>\n",
       "      <td>ago-tr1/labels/99.json</td>\n",
       "      <td>ago-tr1/slides/99.svs</td>\n",
       "      <td>ago-tr1/annotations/99.geojson</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>132 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    SlideNr  HRD(BRCA1)                    label                   slide  \\\n",
       "0        10           0   ago-tr1/labels/10.json   ago-tr1/slides/10.svs   \n",
       "1       100           1  ago-tr1/labels/100.json  ago-tr1/slides/100.svs   \n",
       "2       101           1  ago-tr1/labels/101.json  ago-tr1/slides/101.svs   \n",
       "3       103           1  ago-tr1/labels/103.json  ago-tr1/slides/103.svs   \n",
       "4       106           1  ago-tr1/labels/106.json  ago-tr1/slides/106.svs   \n",
       "..      ...         ...                      ...                     ...   \n",
       "127      84           1   ago-tr1/labels/84.json   ago-tr1/slides/84.svs   \n",
       "128      85           1   ago-tr1/labels/85.json   ago-tr1/slides/85.svs   \n",
       "129       9           0    ago-tr1/labels/9.json    ago-tr1/slides/9.svs   \n",
       "130      92           1   ago-tr1/labels/92.json   ago-tr1/slides/92.svs   \n",
       "131      99           1   ago-tr1/labels/99.json   ago-tr1/slides/99.svs   \n",
       "\n",
       "                          annotation  \n",
       "0     ago-tr1/annotations/10.geojson  \n",
       "1    ago-tr1/annotations/100.geojson  \n",
       "2    ago-tr1/annotations/101.geojson  \n",
       "3    ago-tr1/annotations/103.geojson  \n",
       "4    ago-tr1/annotations/106.geojson  \n",
       "..                               ...  \n",
       "127   ago-tr1/annotations/84.geojson  \n",
       "128   ago-tr1/annotations/85.geojson  \n",
       "129    ago-tr1/annotations/9.geojson  \n",
       "130   ago-tr1/annotations/92.geojson  \n",
       "131   ago-tr1/annotations/99.geojson  \n",
       "\n",
       "[132 rows x 5 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root = \"/mnt/data/Lennard/gyn\"\n",
    "csv_train = \"ago-tr1/csv/finetune_train_1.csv\"\n",
    "frame = pd.read_csv(os.path.join(root, csv_train))\n",
    "frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edf8213d-1539-4874-99b3-efe8f6829d09",
   "metadata": {},
   "source": [
    "### Create a Pytorch Lightning DataModule from slides, annotations, tile labels and slide labels\n",
    "Note that the arguments have prefixes to make their influence more visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ffac9b5a-052e-4665-885a-fe7144a78de8",
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = { \n",
    "  'batch_size': 128,\n",
    "  'column_annotation': 'annotation',\n",
    "  'column_label': 'label',\n",
    "  'column_slide': 'slide',\n",
    "  'columns_global_label': ['HRD(BRCA1)'],\n",
    "  'csv_test': 'ago-tr1/csv/test.csv',\n",
    "  'csv_train': 'ago-tr1/csv/finetune_train_1.csv',\n",
    "  'csv_valid': 'ago-tr1/csv/finetune_eval_1.csv',\n",
    "  'epoch_balance_label_bins': 2,\n",
    "  'epoch_balance_label_key': 'HRD(BRCA1)',\n",
    "  'epoch_balance_size_by': 'median',\n",
    "  'epoch_shuffle': True,\n",
    "  'epoch_shuffle_chunk_size': 16,\n",
    "  'image_size': 240,\n",
    "  'norm_mean': [0.5, 0.5, 0.5],  # for tfms\n",
    "  'norm_std': [0.5, 0.5, 0.5],  # for tfms\n",
    "  'num_workers': 8,\n",
    "  'pin_memory': False,\n",
    "  'regions_annotation_align': False,\n",
    "  'regions_centroid_in_annotation': True,\n",
    "  'regions_level': 0,\n",
    "  'regions_region_overlap': 0.0,\n",
    "  'regions_return_index': False,\n",
    "  'regions_return_labels': 'HRD(BRCA1)',\n",
    "  'regions_size': None,\n",
    "  'regions_unit': 'pixel',\n",
    "  'regions_with_labels': True,\n",
    "  'root': '/mnt/data/Lennard/gyn',\n",
    "  'slide_interpolation': 'linear',\n",
    "  'slide_linear_fill_value': 0.0,\n",
    "  'slide_load_keys': None,\n",
    "  'slide_simplify_tolerance': 100,\n",
    "  'verbose': False\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "79625f70-012e-44fd-bd0b-193e599e5bcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Albumentations has the fastest agumentations acting on numpy arrays\n",
    "# only batched transforms on GPU are faster.\n",
    "\n",
    "class AlbumentationWrapper(nn.Module):\n",
    "    \"\"\"Wrapper around albumentations.\"\"\"\n",
    "\n",
    "    def __init__(self, tfms):\n",
    "        super().__init__()\n",
    "        self.tfms = A.Compose(tfms)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def forward(self, image):\n",
    "        return self.tfms(image=image)[\"image\"]\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"{self.__class__.__name__}({self.tfms.__repr__()})\"\n",
    "\n",
    "\n",
    "# TrivialAugment uses only one augmentation per sample and is therefore very efficient\n",
    "# convince yourself of the results here: https://arxiv.org/pdf/2103.10158.pdf\n",
    "trivial_augment = A.OneOf([\n",
    "    A.Flip(),\n",
    "    A.GaussNoise(),\n",
    "    A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45),\n",
    "    A.ElasticTransform(alpha=100, sigma=100 * 0.1, alpha_affine=100 * 0.03),\n",
    "    A.Blur(p=0.5),\n",
    "    A.ColorJitter(brightness=0.25, contrast=0.75, saturation=0.25, hue=0.5),\n",
    "    A.RandomBrightnessContrast(),\n",
    "    A.CLAHE(),\n",
    "    A.Equalize(),\n",
    "    A.Solarize(),\n",
    "    A.Sharpen(),\n",
    "    A.Posterize(),\n",
    "    A.RandomGamma(),\n",
    "    A.CoarseDropout(),\n",
    "])\n",
    "\n",
    "tfms_train = T.Compose([\n",
    "    AlbumentationWrapper(\n",
    "        tfms=[A.RandomResizedCrop(hparams[\"image_size\"], hparams[\"image_size\"], scale=(0.2, 1.0),\n",
    "                                interpolation=cv2.INTER_CUBIC),\n",
    "              trivial_augment,\n",
    "              A.Normalize(hparams[\"norm_mean\"], hparams[\"norm_std\"])]\n",
    "    ),\n",
    "    T.ToTensor(),\n",
    "])\n",
    "\n",
    "tfms_valid = T.Compose([\n",
    "    AlbumentationWrapper(\n",
    "        tfms=[A.Resize(hparams[\"image_size\"], hparams[\"image_size\"], interpolation=cv2.INTER_CUBIC),\n",
    "              A.Normalize(hparams[\"norm_mean\"], hparams[\"norm_std\"])]\n",
    "    ),\n",
    "    T.ToTensor(),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b4bac66e-fadf-4628-bf26-98ab0a5fb7e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Plugin: cucim.kit.cuslide] Loading the dynamic library from: /home/caduser/anaconda3/envs/tmmae/lib/python3.8/site-packages/cucim/clara/cucim.kit.cuslide@22.02.00.so\n",
      "Initializing plugin: cucim.kit.cuslide (interfaces: [cucim::io::IImageFormat v0.1]) (impl: cucim.kit.cuslide)\n",
      "[Plugin: cucim.kit.cumed] Loading the dynamic library from: /home/caduser/anaconda3/envs/tmmae/lib/python3.8/site-packages/cucim/clara/cucim.kit.cumed@22.02.00.so\n",
      "Initializing plugin: cucim.kit.cumed (interfaces: [cucim::io::IImageFormat v0.1]) (impl: cucim.kit.cumed)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5min 2s, sys: 3.71 s, total: 5min 6s\n",
      "Wall time: 5min 6s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "dm = slide_tools.tile_level.TileLevelDataModule(\n",
    "    **hparams,\n",
    "    tfms_train=tfms_train,\n",
    "    tfms_valid=tfms_valid,\n",
    "    tfms_test=tfms_valid,\n",
    ")\n",
    "dm.setup(\"fit\")  # will prepare train_dataloader and valid_dataloader\n",
    "\n",
    "# The runtime is mostly dominated by reading the json labels.\n",
    "# Saving/loading of labels will likely change in the future to a more performant variant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "56142ade-03d5-42e8-a5ee-510385b00b6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 8s, sys: 524 ms, total: 1min 9s\n",
      "Wall time: 1min 9s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "dm.setup(\"test\") # will prepare test_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "47cc8c42-b7e9-4332-8164-a73a80273908",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "538 tiles per second with num_workers=8 at train time (random tiles)\n"
     ]
    }
   ],
   "source": [
    "batch_size = 128\n",
    "num_workers = 8\n",
    "dl_train = dm.train_dataloader(batch_size=batch_size, num_workers=num_workers)\n",
    "\n",
    "N = 100\n",
    "for i, batch in enumerate(dl_train):\n",
    "    if i == 20:\n",
    "        t0 = time.time()\n",
    "    if i == (N + 20):\n",
    "        t = time.time() - t0\n",
    "        break\n",
    "print(f\"{batch_size*N/t:.0f} tiles per second with {num_workers=} at train time (random tiles)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "25d0154f-cdef-48e0-a86e-d9c94d572861",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "581 tiles per second with num_workers=8 at test time (sequential tiles)\n"
     ]
    }
   ],
   "source": [
    "batch_size = 128\n",
    "num_workers = 8\n",
    "dl_valid = dm.val_dataloader(batch_size=batch_size, num_workers=num_workers)\n",
    "\n",
    "N = 100\n",
    "for i, batch in enumerate(dl_valid):\n",
    "    if i == 20:\n",
    "        t0 = time.time()\n",
    "    if i == (N + 20):\n",
    "        t = time.time() - t0\n",
    "        break\n",
    "print(f\"{batch_size*N/t:.0f} tiles per second with {num_workers=} at test time (sequential tiles)\")"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}