{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fcfb563d-3f9c-4731-be1f-9c9c4b2c9dfd",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Medical Imaging Classification (Colorectal histology) using TensorFlow and the Intel® Transfer Learning Tool API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "353f2782-043f-4afd-9858-cc773275e6c5",
   "metadata": {},
   "source": [
    "This notebook facilitates implementation of medical imaging classification using Transfer Learning Toolkit. It performs Multi-class texture analysis in colorectal cancer histology dataset. The workflow uses pretrained SOTA models ( RESNET V1.5) from TF hub and transfers the knowledge from a pretrained domain to a different custom domain achieving required accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8325300b-4a75-42fb-aa14-2b089b4edea5",
   "metadata": {},
   "source": [
    "## 1. Import dependencies and setup parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04e3f113-4887-49c9-aa8b-a3c2058157a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "import tensorflow as tf\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "#tlt imports\n",
    "from tlt.datasets import dataset_factory\n",
    "from tlt.models import model_factory\n",
    "from tlt.utils.types import FrameworkType, UseCaseType\n",
    "\n",
    "from notebooks.plot_utils import plot_curves\n",
    "\n",
    "# Specify a directory for the dataset to be downloaded\n",
    "dataset_dir = os.environ[\"DATASET_DIR\"] if \"DATASET_DIR\" in os.environ else \\\n",
    "    os.path.join(os.environ[\"HOME\"], \"dataset\")\n",
    "\n",
    "# Specify a directory for output\n",
    "output_dir = os.environ[\"OUTPUT_DIR\"] if \"OUTPUT_DIR\" in os.environ else \\\n",
    "    os.path.join(os.environ[\"HOME\"], \"output\")\n",
    "\n",
    "print(\"Dataset directory:\", dataset_dir)\n",
    "print(\"Output directory:\", output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a47ea534-0661-4cbe-97ed-ca288b6203e5",
   "metadata": {},
   "source": [
    "## 2. Get the model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c9548167-cdfd-44ad-be60-cc3d131850a3",
   "metadata": {},
   "source": [
    "In this step, we call the Intel Transfer Learning Tool model factory to list supported TensorFlow image classification models. This is a list of pretrained models from TFHub that we tested with our API. Optionally, the verbose=True argument can be added to the print_supported_models function call to get more information about each model (such as the link to TFHub, image size, the original dataset, etc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bae74045-f024-498f-9ea7-3ee8ac2ea5d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See a list of available models\n",
    "model_factory.print_supported_models(use_case='image_classification', framework='tensorflow')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c06b359-1054-45c5-9223-47d4e1b7772d",
   "metadata": {},
   "source": [
    "#### Option A: Load a model\n",
    "\n",
    "Next, use the model factory to get one of the models listed in the previous cell. The `get_model` function returns a model object that will later be used for training. By default, resnet_v1_50 is used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf842217-d992-4b6f-99bf-aef475fa13dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the model\n",
    "model = model_factory.get_model(model_name=\"resnet_v1_50\", framework=\"tensorflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa83899-0a99-4c9e-9e3d-00b040fe3e1f",
   "metadata": {},
   "source": [
    "#### Option B: Load a pretrained checkpoint\n",
    "\n",
    "Optionally, to continue training using a pretrained checkpoint, the user can specify the path to folder containing __saved_model.pb__. The user can specify the path in __model__ parameter.\n",
    "\n",
    "_Note: The path is same as saved_model_dir_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b49423b-e57d-478c-87b0-7daec6401b76",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load a pretrained checkpoint\n",
    "model = model_factory.load_model(model_name='resnet_v1_50', \n",
    "                                 model='/home/intel/output/resnet_v1_50/1', \n",
    "                                 framework='tensorflow', use_case='image_classification')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba2e4d10-458b-4479-8d1f-8cc67304d5c2",
   "metadata": {},
   "source": [
    "## 3. Get the dataset\n",
    "Use dataset __colorectal_histology__ from the TensorFlow Datasets catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cddcfdcf-ba16-4faa-a708-61f6aa01b3b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset_factory.get_dataset(dataset_dir=dataset_dir,\n",
    "                                      use_case='image_classification', \n",
    "                                      framework='tensorflow',\n",
    "                                      dataset_name='colorectal_histology',\n",
    "                                      dataset_catalog='tf_datasets')\n",
    "\n",
    "print(dataset.info)\n",
    "\n",
    "print(\"\\nClass names:\", str(dataset.class_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ae3d375-554c-487b-8907-e895396d20e7",
   "metadata": {},
   "source": [
    "## 4. Prepare the dataset\n",
    "\n",
    "Once you have your dataset from Option A or Option B above, use the following cells to preprocess the dataset. We resize the images to match the selected models and batch the images, then split them into training and validation subsets.  Data augmentation can be applied by specifying the augmentations to be applied in __add_aug__ parameter. Supported augmentations are \n",
    "1. hvflip - RandomHorizontalandVerticalFlip\n",
    "2. hflip - RandomHorizontalFlip\n",
    "3. vflip - RandomVerticalFlip\n",
    "4. rotate - RandomRotate\n",
    "5. zoom - RandomZoom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d529b1-fe40-482a-b3f0-3875b13ba534",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess the dataset with an image size that matches the model and a batch size of 32\n",
    "batch_size = 32\n",
    "dataset.preprocess(model.image_size, batch_size=batch_size, add_aug=['hvflip', 'rotate'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c4af8c6-c418-45d1-b0bc-dd8c7e2625c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the dataset into training, validation and test subsets\n",
    "dataset.shuffle_split(train_pct=.80, val_pct=.10, test_pct=0.10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab02c7d0-5b63-4be3-b6a6-cdfc03e48995",
   "metadata": {},
   "source": [
    "## 5. Transfer Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feeb23c9-66a0-4670-8fd0-f7950dd2832e",
   "metadata": {},
   "source": [
    "This step calls the model's train function with the dataset that was just prepared. The training function will get the TFHub feature vector and add on a dense layer based on the number of classes in the dataset. The model is then compiled and trained based on the number of epochs specified in the argument. With the do_eval paramter set to True by default, this step will also show how the model can be evaluated and will return a list of metrics calculated from the dataset's validation subset.\n",
    "### Arguments\n",
    "#### Required\n",
    "-  **dataset** (ImageClassificationDataset, required): Dataset to use when training the model\n",
    "-  **output_dir** (str): Path to a writeable directory for checkpoint files\n",
    "-  **epochs** (int): Number of epochs to train the model (default: 1)\n",
    "#### Optional\n",
    "-  **initial_checkpoints** (str): Path to checkpoint weights to load. If the path provided is a directory, the latest checkpoint will be used.\n",
    "-  **early_stopping** (bool): Enable early stopping if convergence is reached while training at the end of each epoch. (default: False)\n",
    "-  **lr_decay** (bool): If lr_decay is True and do_eval is True, learning rate decay on the validation loss is applied at the end of each epoch.\n",
    "-  **enable_auto_mixed_precision** (bool or None): Enable auto mixed precision for training. Mixed precision uses both 16-bit and 32-bit floating point types to make training run faster and use less memory. It is recommended to enable auto mixed precision training when running on platforms that support bfloat16 (Intel third or fourth generation Xeon processors). If it is enabled on a platform that does not support bfloat16, it can be detrimental to the training performance. If enable_auto_mixed_precision is set to None, auto mixed precision will be automatically enabled when running with Intel fourth generation Xeon processors, and disabled for other platforms.\n",
    "-  **extra_layers** (list[int]): Optionally insert additional dense layers between the base model and output layer. This can help increase accuracy when fine-tuning a TFHub model. The input should be a list of integers representing the number and size of the layers, for example [1024, 512] will insert two dense layers, the first with 1024 neurons and the second with 512 neurons.\n",
    "\n",
    "Note: refer to release documentation for an up-to-date list of train arguments and their current descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa477dbb-58a9-4e06-9933-047699c35797",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mixed precision uses both 16-bit and 32-bit floating point types to make training run faster and use less memory.\n",
    "# It is recommended to enable auto mixed precision training when running on platforms that support\n",
    "# bfloat16 (Intel third or fourth generation Xeon processors). If it is enabled on a platform that\n",
    "# does not support bfloat16, it can be detrimental to the training performance.\n",
    "# If enable_auto_mixed_precision is set to None, auto mixed precision will be automatically enabled when\n",
    "# running with Intel fourth generation Xeon processors, and disabled for other platforms.\n",
    "enable_auto_mixed_precision = None\n",
    "\n",
    "# Train the model using the dataset\n",
    "history = model.train(dataset, output_dir=output_dir, epochs=50, \n",
    "                      enable_auto_mixed_precision=None, early_stopping=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40e11437-5666-49c0-9547-22e7696c0ea0",
   "metadata": {},
   "source": [
    "## 6. Evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d413b5d-ed24-4756-928d-dbb8e7bf1af2",
   "metadata": {},
   "source": [
    "The next step shows how the model can be evaluated. The model's evaluate function returns a list of metrics calculated from the dataset's validation subset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5bfc8e8-f8c6-4b56-bfc8-51b88288fcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate model on validation subset\n",
    "val_loss, val_acc = model.evaluate(dataset)\n",
    "print(\"Validation Accuracy :\", val_acc)\n",
    "print(\"Validation Loss :\", val_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c478191f-2c0c-49d6-a295-0a3acfd183a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_curves(history, os.path.join(output_dir, \"{}_checkpoints\".format(model.model_name)))\n",
    "pickle.dump(history, open(os.path.join(output_dir, \"{}_checkpoints\".format(model.model_name), 'hist.pkl'), 'wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c6fd051-97af-4142-b039-482cb742d988",
   "metadata": {},
   "source": [
    "## 7. Export"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9225921d-5340-4865-ad4b-a2c19a5210ed",
   "metadata": {},
   "source": [
    "Next, we can call the model export function to generate a saved_model.pb. The model is saved in a format that is ready to use with TensorFlow Serving. Each time the model is exported, a new numbered directory is created, which allows serving to pick up the latest model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c4fc5f8-cb72-4ab5-989b-43d7ad315e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "saved_model_dir = model.export(output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d79f06c-39cf-4c05-93bf-b35377ec06b1",
   "metadata": {},
   "source": [
    "## 8. Inference\n",
    "To perform only Inference using a saved model, follow the steps below\n",
    "1. Execute Step 2(b) to load a pretrained checkpoint with the appropriate model name.\n",
    "2. Execute Steps 3 and 4 to load and prepare the dataset.\n",
    "3. Continue with the steps below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb2182c-6302-474b-abab-7cf16c6b0966",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = pickle.load(open(os.path.join(output_dir, \"{}_checkpoints\".format(model.model_name), 'hist.pkl'), 'rb'))\n",
    "plot_curves(history, os.path.join(output_dir, \"{}_checkpoints\".format(model.model_name)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5da7be1-76d9-4554-ace6-7bf6ef0f8c21",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss, accuracy = model.evaluate(dataset, use_test_set=True)\n",
    "print('Test accuracy :', accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb585ceb-d120-4ded-a3b2-e78a5bd3dd7b",
   "metadata": {},
   "source": [
    "We get the test subset from our dataset, and use that to call predict on our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1e0450-f6b7-4abd-990c-b65b8d77e599",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual_labels = np.concatenate([y for x, y in dataset._test_subset], axis=0)\n",
    "predicted_labels = model.predict(dataset._test_subset)\n",
    "report = classification_report(actual_labels, predicted_labels)\n",
    "print(\"Classification report\")\n",
    "print(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7863e336-7996-4518-9ada-2b7fdd9672b5",
   "metadata": {},
   "source": [
    "## Dataset Citations\n",
    "\n",
    "@article{kather2016multi,<br>\n",
    "      title={Multi-class texture analysis in colorectal cancer histology},<br>\n",
    "      author={Kather, Jakob Nikolas and Weis, Cleo-Aron and Bianconi, Francesco and Melchers, Susanne M and Schad, Lothar R and Gaiser, Timo and Marx, Alexander and Z{\"o}llner, Frank Gerrit},<br>\n",
    "      journal={Scientific reports},<br>\n",
    "      volume={6},<br>\n",
    "      pages={27988},<br>\n",
    "      year={2016},<br>\n",
    "      publisher={Nature Publishing Group}<br>\n",
    "    }\n",
    "    \n",
    "Kather, J. N., Zöllner, F. G., Bianconi, F., Melchers, S. M., Schad, L. R., Gaiser, T., Marx, A., & Weis, C.-A. (2016). Collection of textures in colorectal cancer histology [Data set]. Zenodo. https://doi.org/10.5281/zenodo.53169"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}