Session 3 -- Interpretability of MIL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eb7615aa",
   "metadata": {
    "id": "eb7615aa"
   },
   "source": [
    "### Interpretability of Attention-based MIL\n",
    "\n",
    "In this session, we will explore two post-hoc interpretability methods for understanding the behaviour of an Attention-based MIL model trained to subtype lung cancer into LUAD and LUSC, the two most common lung cancer subtypes.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf08e37d",
   "metadata": {
    "id": "cf08e37d"
   },
   "source": [
    "#### On the importance of interpretability in medical AI applications\n",
    "\n",
    "- **Trust & control**: Medical AI systems are often used in critical decision-making processes that directly impact patients' health and lives. Interpretability aims to provide the right level of insights and control over the system such that trust and confidence can be established between users and the system.\n",
    "\n",
    "\n",
    "- **Insights and Understanding**: Interpretable models are also valuable to provide insights into the underlying factors or features that contribute to a prediction. In medical AI, this understanding can be critical for healthcare professionals to gain insights into disease mechanisms, identify risk factors, or discover novel biomarkers. This aspect goes in the direction of biomarker discovery, where we assume that the AI system might use different features than what form the current standard (e.g., the current grading criteria in cancer).\n",
    "\n",
    "\n",
    "- **Error Analysis and Diagnosis**: Interpretability helps in error analysis, allowing the identification and understanding of the model's mistakes or mispredictions. In medical AI, where misdiagnoses can have severe consequences, interpretability enables clinicians to evaluate cases where the model failed and diagnose potential pitfalls or limitations. This feedback loop can guide improvements in the model, dataset, or feature engineering, leading to better performance and more reliable predictions.\n",
    "\n",
    "\n",
    "And the wishful thinking considerations...\n",
    "\n",
    "\n",
    "- **Legal and Ethical Considerations**: The interpretability of AI models could be used for addressing legal and ethical concerns. In healthcare, decisions made by AI systems need to be explainable to patients, healthcare professionals, regulatory bodies, and other stakeholders. By providing interpretability, AI systems can adhere to legal requirements, such as the General Data Protection Regulation (GDPR), which grants individuals the right to an explanation for automated decisions that significantly impact them.\n",
    "\n",
    "\n",
    "- **Safety and Robustness**: Deep learning models are susceptible to biases, adversarial attacks, or data distribution shifts that can lead to incorrect or unreliable predictions. Interpretability helps in detecting these issues and assessing the model's safety and robustness. By understanding the model's internal workings, it becomes possible to identify potential biases, investigate cases where the model may be overconfident or underperform, and design safeguards to mitigate risks.\n",
    "\n",
    "\n",
    "- **Regulatory Compliance**: Interpretability is increasingly becoming a regulatory requirement in various domains, including healthcare. Regulatory bodies, such as the U.S. Food and Drug Administration (FDA), often demand explanations and justifications for the decisions made by AI systems before approving their deployment. Interpretability allows the model's behavior to be audited, validated, and aligned with regulatory standards, ensuring compliance and patient safety.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "QcNSlWhRR-X3",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QcNSlWhRR-X3",
    "outputId": "1463bbdc-5441-4500-85d1-a51bedcaab3e"
   },
   "outputs": [],
   "source": [
    "# All data can be downloaded here: https://drive.google.com/drive/folders/1TmAfG7EWC1hjD7cHFGiJzUx2y3jLXdcP?usp=sharing\n",
    "# Once downloaded, transfer the data into your local cloned repo \n",
    "\n",
    "use_drive = False\n",
    "if use_drive:\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/drive')\n",
    "    !mkdir -p \"/content/drive/My Drive/ai4healthsummerschool/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b4e6950",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2b4e6950",
    "outputId": "d9d2c01a-50af-42c3-c13f-a5985d155817"
   },
   "outputs": [],
   "source": [
    "# Load and print ABMIL model previously trained\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pandas as pd\n",
    "\n",
    "class AttentionTanhSigmoidGating(nn.Module):\n",
    "    def __init__(self, D=64, L=64, dropout=0.25):\n",
    "        r\"\"\"\n",
    "        Global attention pooling layer with tanh non-linearity and sigmoid gating (Ilse et al. 2018).\n",
    "\n",
    "        Args:\n",
    "            D (int): input feature dimension.\n",
    "            L (int): hidden layer dimension. Notation changed from M from Ilse et al 2018, as M is overloaded to also describe # of patch embeddings in a WSI.\n",
    "            dropout (float): Dropout probability.\n",
    "\n",
    "        Returns:\n",
    "            A_norm (torch.Tensor): [M x 1]-dim tensor of normalized attention scores (sum to 1)\n",
    "        \"\"\"\n",
    "        super(AttentionTanhSigmoidGating, self).__init__()\n",
    "        self.tanhV = nn.Sequential(*[nn.Linear(D, L), nn.Tanh(), nn.Dropout(dropout)])\n",
    "        self.sigmU = nn.Sequential(*[nn.Linear(D, L), nn.Sigmoid(), nn.Dropout(dropout)])\n",
    "        self.w = nn.Linear(L, 1)\n",
    "\n",
    "    def forward(self, H, return_raw_attention=False):\n",
    "        A_raw = self.w(self.tanhV(H).mul(self.sigmU(H))) # exponent term\n",
    "        A_norm = F.softmax(A_raw, dim=0)                 # apply softmax to normalize weights to 1\n",
    "        assert abs(A_norm.sum() - 1) < 1e-3              # Assert statement to check sum(A) ~= 1\n",
    "        if return_raw_attention:\n",
    "            return A_norm, A_raw\n",
    "        return A_norm\n",
    "\n",
    "\n",
    "class ABMIL(nn.Module):\n",
    "    def __init__(self, input_dim=320, hidden_dim=64, dropout=0.25, n_classes=2):\n",
    "        r\"\"\"\n",
    "        Attention-Based Multiple Instance Learning (Ilse et al. 2018).\n",
    "\n",
    "        Args:\n",
    "            input_dim (int): input feature dimension.\n",
    "            hidden_dim (int): hidden layer dimension.\n",
    "            dropout (float): Dropout probability.\n",
    "            n_classes (int): Number of classes.\n",
    "        \"\"\"\n",
    "        super(ABMIL, self).__init__()\n",
    "        self.inst_level_fc = nn.Sequential(*[nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout)]) # Fully-Connected Layer, applied \"instance-wise\" to each embedding\n",
    "        self.global_attn = AttentionTanhSigmoidGating(L=hidden_dim, D=hidden_dim)                              # Attention Function\n",
    "        self.bag_level_classifier = nn.Linear(hidden_dim, n_classes)                                            # Bag-Level Classifier\n",
    "\n",
    "    def forward(self, X: torch.randn(100, 320), return_raw_attention=False):\n",
    "        r\"\"\"\n",
    "        Takes as input a [M x D]-dim bag of patch features (representing a WSI), and outputs: 1) logits for classification, 2) un-normalized attention scores.\n",
    "\n",
    "        Args:\n",
    "            X (torch.Tensor): [M x D]-dim bag of patch features (representing a WSI)\n",
    "\n",
    "        Returns:\n",
    "            logits (torch.Tensor): [1 x n_classes]-dim tensor of un-normalized logits for classification task.\n",
    "            A_norm (torch.Tensor): [M,]- or [M x 1]-dim tensor of attention scores.\n",
    "        \"\"\"\n",
    "        H_inst = self.inst_level_fc(X)         # 1. Process each feature embedding to be of size \"hidden-dim\"\n",
    "        if return_raw_attention:\n",
    "            A_norm, A_raw = self.global_attn(H_inst, return_raw_attention=True)\n",
    "        else:\n",
    "            A_norm = self.global_attn(H_inst)      # 2. Get normalized attention scores for each embedding (s.t. sum(A_norm) ~= 1)\n",
    "        z = torch.sum(A_norm * H_inst, dim=0)  # 3. Output of global attention pooling over the bag\n",
    "        logits = self.bag_level_classifier(z).unsqueeze(dim=0)   # 4. Get un-normalized logits for classification task\n",
    "        try:\n",
    "            assert logits.shape == (1,2)\n",
    "        except:\n",
    "            print(f\"Logit tensor shape is not formatted correctly. Should output [1 x 2] shape, but got {logits.shape} shape\")\n",
    "        \n",
    "        if return_raw_attention:\n",
    "            return logits, A_raw\n",
    "        return logits, A_norm\n",
    "\n",
    "    def captum(self, X: torch.randn(100, 320)):\n",
    "        r\"\"\"\n",
    "        Takes as input a [M x D]-dim bag of patch features (representing a WSI), and outputs: 1) logits for classification, 2) un-normalized attention scores.\n",
    "\n",
    "        Args:\n",
    "            X (torch.Tensor): [M x D]-dim bag of patch features (representing a WSI)\n",
    "\n",
    "        Returns:\n",
    "            logits (torch.Tensor): [1 x n_classes]-dim tensor of un-normalized logits for classification task.\n",
    "            A_norm (torch.Tensor): [M,]- or [M x 1]-dim tensor of attention scores.\n",
    "        \"\"\"\n",
    "        H_inst = self.inst_level_fc(X)         # 1. Process each feature embedding to be of size \"hidden-dim\"\n",
    "        A_norm = self.global_attn(H_inst)      # 2. Get normalized attention scores for each embedding (s.t. sum(A_norm) ~= 1)\n",
    "        z = torch.sum(A_norm * H_inst, dim=0)  # 3. Output of global attention pooling over the bag\n",
    "        logits = self.bag_level_classifier(z).unsqueeze(dim=0)   # 4. Get un-normalized logits for classification task\n",
    "        try:\n",
    "            assert logits.shape == (1,2)\n",
    "        except:\n",
    "            print(f\"Logit tensor shape is not formatted correctly. Should output [1 x 2] shape, but got {logits.shape} shape\")\n",
    "        return logits\n",
    "\n",
    "\n",
    "device = torch.device('cpu')\n",
    "model = ABMIL(input_dim=320, hidden_dim=64).to(device)\n",
    "\n",
    "if use_drive:\n",
    "    path = '/content/drive/My Drive/ai4healthsummerschool/abmil.ckpt'\n",
    "else:\n",
    "    path = os.path.join('data', 'checkpoints', 'abmil.ckpt')\n",
    "    \n",
    "model.load_state_dict(torch.load(path))\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "472c25de",
   "metadata": {
    "id": "472c25de"
   },
   "source": [
    "#### Interpreting Attention-based MIL with Attention weights\n",
    "\n",
    "Attention weights are a mechanism used in deep learning models to determine the importance or relevance of different parts of the input data. These weights can be utilized to interpret deep learning predictions by providing insights into which parts of the input contribute more strongly to the model's decision-making process."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "640c0ccf",
   "metadata": {
    "id": "640c0ccf"
   },
   "source": [
    "###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3fcb8c9",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c3fcb8c9",
    "outputId": "1a810f85-66ae-459b-a07c-7ee48c43b60e"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Load features and corresponding label\n",
    "if use_drive:\n",
    "    feats_dirpath='/content/drive/My Drive/ai4healthsummerschool/feats_pt'\n",
    "    csv_fpath='/content/drive/My Drive/ai4healthsummerschool/tcga_lung_splits.csv'\n",
    "else:\n",
    "    feats_dirpath = os.path.join('data', 'processed', 'feats_pt')\n",
    "    csv_fpath = os.path.join('data', 'processed', 'tcga_lung_splits.csv')\n",
    "index = 5  # (LUAD sample)\n",
    "\n",
    "csv = pd.read_csv(csv_fpath)\n",
    "which_labelcol = 'OncoTreeCode_Binarized'\n",
    "csv_split = csv[csv['split']=='test']\n",
    "\n",
    "features = torch.load(os.path.join(feats_dirpath, csv_split.iloc[index]['slide_id']+'.pt'))\n",
    "label = csv_split.iloc[index][which_labelcol]\n",
    "\n",
    "print('We will analyse the prediction and attention scores of:')\n",
    "print(csv_split.iloc[index]['slide_id']+'.pt')\n",
    "print('Label: ', label)\n",
    "print('Features:', features.shape)\n",
    "\n",
    "# Run inference and store attention weights\n",
    "logits, attention = model(features, return_raw_attention=True)\n",
    "logits = logits.squeeze()\n",
    "attention = attention.squeeze().detach().numpy()\n",
    "print(logits)\n",
    "print('Shape: ', attention.shape, 'Min: ', np.min(attention).item(), 'Max: ', np.max(attention).item())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7x2kFBsnwq4T",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "id": "7x2kFBsnwq4T",
    "outputId": "4b5ebc78-b641-4fc1-8a7d-3f217db0a315"
   },
   "outputs": [],
   "source": [
    "!pip install openslide-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6273b4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the attention weights on the input WSIs\n",
    "import h5py\n",
    "import numpy as np\n",
    "import cv2\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "from openslide import OpenSlide\n",
    "\n",
    "# load coords corresponding to sample of interest\n",
    "if use_drive:\n",
    "    path = '/content/drive/My Drive/ai4healthsummerschool/TCGA-35-3615-01Z-00-DX1.585128eb-6652-4b05-9a83-dc8f242904a6_patches.h5'\n",
    "else:\n",
    "    path = os.path.join('data', 'processed', 'TCGA-35-3615-01Z-00-DX1.585128eb-6652-4b05-9a83-dc8f242904a6_patches.h5')\n",
    "    \n",
    "with h5py.File(path, \"r\") as f:\n",
    "    coords = f['coords'][:]\n",
    "    print('Coords:', coords.shape)\n",
    "\n",
    "# load whole slide imahe using OpenSlide \n",
    "if use_drive:\n",
    "    slide_path = '/content/drive/My Drive/ai4healthsummerschool/TCGA-35-3615-01Z-00-DX1.585128eb-6652-4b05-9a83-dc8f242904a6.svs'\n",
    "else:\n",
    "    slide_path = os.path.join('data', 'slides', 'TCGA-35-3615-01Z-00-DX1.585128eb-6652-4b05-9a83-dc8f242904a6.tiff')\n",
    "wsi = OpenSlide(slide_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "R09MoSFKYa85",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 263
    },
    "id": "R09MoSFKYa85",
    "outputId": "0d079399-0a8e-4fbb-9ebb-94d8e63d12ca"
   },
   "outputs": [],
   "source": [
    "def draw_heatmap(scores, coords, wsi, vis_level=-1,\n",
    "                patch_size=(256, 256),\n",
    "                blank_canvas=False, canvas_color=(220, 20, 50), alpha=0.4,\n",
    "                overlap=0.0, use_holes=True,\n",
    "                convert_to_percentiles=False, thresh=0.5,\n",
    "                max_size=None, custom_downsample = 4,\n",
    "                cmap='coolwarm'):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        scores (numpy array of float): Attention scores\n",
    "        coords (numpy array of int, n_patches x 2): Corresponding coordinates (relative to lvl 0)\n",
    "        wsi (openslide): WSI opened with openslide\n",
    "        vis_level (int): WSI pyramid level to visualize\n",
    "        patch_size (tuple of int): Patch dimensions (relative to lvl 0)\n",
    "        blank_canvas (bool): Whether to use a blank canvas to draw the heatmap (vs. using the original slide)\n",
    "        canvas_color (tuple of uint8): Canvas color\n",
    "        alpha (float [0, 1]): blending coefficient for overlaying heatmap onto original slide\n",
    "        blur (bool): apply gaussian blurring\n",
    "        overlap (float [0 1]): percentage of overlap between neighboring patches (only affect radius of blurring)\n",
    "        segment (bool): whether to use tissue segmentation contour (must have already called self.segmentTissue such that\n",
    "                        self.contours_tissue and self.holes_tissue are not None\n",
    "        use_holes (bool): whether to also clip out detected tissue cavities (only in effect when segment == True)\n",
    "        convert_to_percentiles (bool): whether to convert attention scores to percentiles\n",
    "        binarize (bool): only display patches > threshold\n",
    "        threshold (float): binarization threshold\n",
    "        max_size (int): Maximum canvas size (clip if goes over)\n",
    "        custom_downsample (int): additionally downscale the heatmap by specified factor\n",
    "        cmap (str): name of matplotlib colormap to use\n",
    "    \"\"\"\n",
    "\n",
    "    downsample = (0.25, 0.25)\n",
    "    patch_size  = np.ceil(np.array(patch_size)).astype(int)\n",
    "    coords = np.ceil(coords * np.array(downsample)).astype(int)\n",
    "\n",
    "    region_size = wsi.level_dimensions[vis_level]\n",
    "    w, h = region_size\n",
    "\n",
    "    print('\\ncreating heatmap for: ')\n",
    "    print('w: {}, h: {}'.format(w, h))\n",
    "    print('scaled patch size: ', patch_size)\n",
    "\n",
    "    # heatmap overlay: tracks attention score over each pixel of heatmap\n",
    "    # overlay counter: tracks how many times attention score is accumulated over each pixel of heatmap\n",
    "\n",
    "    overlay = np.full(np.flip(region_size), 0).astype(float)\n",
    "    counter = np.full(np.flip(region_size), 0).astype(np.uint16)\n",
    "    count = 0\n",
    "    for idx in range(len(coords)):\n",
    "        score = scores[idx].item()\n",
    "        coord = coords[idx]\n",
    "        # accumulate scores\n",
    "        overlay[coord[1]:coord[1]+patch_size[1], coord[0]:coord[0]+patch_size[0]] += score\n",
    "        # accumulate counter\n",
    "        counter[coord[1]:coord[1]+patch_size[1], coord[0]:coord[0]+patch_size[0]] += 1\n",
    "\n",
    "    # fetch attended region and average accumulated attention\n",
    "    zero_mask = counter == 0\n",
    "    overlay[~zero_mask] = overlay[~zero_mask] / counter[~zero_mask]\n",
    "    del counter\n",
    "\n",
    "    img = np.array(wsi.read_region((0, 0), vis_level, region_size).convert(\"RGB\"))\n",
    "\n",
    "    print('\\ncomputing heatmap image')\n",
    "    print('total of {} patches'.format(len(coords)))\n",
    "    twenty_percent_chunk = max(1, int(len(coords) * 0.2))\n",
    "\n",
    "    if isinstance(cmap, str):\n",
    "        cmap = plt.get_cmap(cmap)\n",
    "        norm = plt.Normalize(scores.min(), scores.max())\n",
    "\n",
    "    for idx in range(len(coords)):\n",
    "        if (idx + 1) % twenty_percent_chunk == 0:\n",
    "            print('progress: {}/{}'.format(idx, len(coords)))\n",
    "\n",
    "        score = scores[idx].item()\n",
    "        coord = coords[idx]\n",
    "\n",
    "        # attention block\n",
    "        raw_block = overlay[coord[1]:coord[1]+patch_size[1], coord[0]:coord[0]+patch_size[0]]\n",
    "\n",
    "        # image block (either blank canvas or orig image)\n",
    "        img_block = img[coord[1]:coord[1]+patch_size[1], coord[0]:coord[0]+patch_size[0]].copy()\n",
    "\n",
    "        # color block (cmap applied to attention block)\n",
    "        color_block = (cmap(norm(raw_block)) * 255)[:,:,:3].astype(np.uint8)\n",
    "\n",
    "        # copy over entire color block\n",
    "        img_block = color_block\n",
    "\n",
    "        # rewrite image block\n",
    "        img[coord[1]:coord[1]+patch_size[1], coord[0]:coord[0]+patch_size[0]] = img_block.copy()\n",
    "\n",
    "    #return Image.fromarray(img) #overlay\n",
    "    print('Done')\n",
    "    del overlay\n",
    "\n",
    "    img = Image.fromarray(img)\n",
    "    w, h = img.size\n",
    "\n",
    "    if custom_downsample > 1:\n",
    "        img = img.resize((int(w/custom_downsample), int(h/custom_downsample)))\n",
    "\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54ecda1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "heatmap = draw_heatmap(\n",
    "  scores=attention,\n",
    "  coords=coords,\n",
    "  wsi=wsi,\n",
    "  use_holes=True,\n",
    "  vis_level=1,\n",
    "  blank_canvas=False,\n",
    "  convert_to_percentiles=False\n",
    ")\n",
    "\n",
    "heatmap.save(os.path.join('data', 'interpretability', 'attention_heatmap.png'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8BKWfDSQwoQ_",
   "metadata": {
    "id": "8BKWfDSQwoQ_"
   },
   "outputs": [],
   "source": [
    "# Extract the most important patches \n",
    "n_samples = 10 \n",
    "\n",
    "# sort attention weights \n",
    "to_keep = np.argsort(attention)[-n_samples:]\n",
    "scores_to_keep = attention[to_keep]\n",
    "coords_to_keep = coords[to_keep]\n",
    "\n",
    "print('Indices to keep:', to_keep)\n",
    "print('Scores:', scores_to_keep)\n",
    "\n",
    "for idx in range(n_samples):\n",
    "    patch = wsi.read_region(coords_to_keep[idx], 1, (256, 256))\n",
    "    patch.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63667a5a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "88af12fa",
   "metadata": {
    "id": "88af12fa"
   },
   "source": [
    "#### Interpretating with Integrated Gradients\n",
    "\n",
    "In this section, we will use Captum, an open-source package that provides off-the-shelf post-hoc interpretability techniques, including Integrated Gradients (IG)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hImkjH_VZB9h",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hImkjH_VZB9h",
    "outputId": "6f5e8bbe-3e21-4f82-988c-c31e5a4224e2",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install captum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "593facd1",
   "metadata": {
    "id": "593facd1"
   },
   "source": [
    "1. **Notation:**\n",
    "   - Let's consider a deep learning model with an input vector `x` and output `f(x)`.\n",
    "   - The baseline or reference point is denoted as `x'`, typically chosen as a point with low complexity (e.g., all zeros or random noise).\n",
    "   - The attribution score for each input feature `i` is denoted as `A_i`.\n",
    "\n",
    "2. **Gradient Calculation:**\n",
    "   - Compute the gradients of the model's output with respect to the input features:\n",
    "   \n",
    "     $$\\vec{\\nabla} f(x) = \\left(\\frac{\\partial f(x)}{\\partial x_1}, \\frac{\\partial f(x)}{\\partial x_2}, \\ldots, \\frac{\\partial f(x)}{\\partial x_n}\\right)$$\n",
    "\n",
    "3. **Integrated Gradients Formula:**\n",
    "   - The Integrated Gradients score for each feature `i` is calculated as follows:\n",
    "   \n",
    "     $$A_i = (x_i - x'_i) \\times \\int_{\\alpha=0}^1 \\left(\\frac{\\partial f(x'+\\alpha(x-x'))}{\\partial x_i}\\right) d\\alpha$$\n",
    "\n",
    "4. **Explanation:**\n",
    "   - Integrated Gradients computes the contribution of each feature `i` by taking into account the difference between the input `x` and the baseline `x'`.\n",
    "   - It then integrates the gradients of the model's output with respect to feature `i` along a straight path from the baseline `x'` to the input `x`.\n",
    "   - The integral is calculated over a series of steps (α) from 0 to 1, representing the interpolation between the baseline and the input.\n",
    "   - The gradients at each interpolation point measure the sensitivity of the output to changes in feature `i` as we move from the baseline to the input.\n",
    "   - The contribution of feature `i` is multiplied by the difference between the input and baseline for that feature, capturing the change in the model's output caused by that feature.\n",
    "\n",
    "5. **Implementation Steps:**\n",
    "   - Select a baseline point `x'` (all zeros, random noise, or other relevant choice).\n",
    "   - Define the number of steps or intervals for the integration.\n",
    "   - For each step α from 0 to 1, calculate `x'+α(x-x')` as the intermediate input.\n",
    "   - Compute the gradients of the model's output with respect to the intermediate input at each step.\n",
    "   - Accumulate the gradients and multiply by the difference between `x` and `x'` for each feature `i` to calculate the attribution scores `A_i`.\n",
    "\n",
    "The Integrated Gradients technique provides feature-level attribution scores, allowing us to understand the importance of each input feature in the model's predictions. By visualizing or analyzing these scores, we can gain insights into which features are influential or critical for the model's decision-making process.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "700f5efc",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "700f5efc",
    "outputId": "54cb8bcd-27d6-4202-9a58-482135ac5f11"
   },
   "outputs": [],
   "source": [
    "# use captum to get IG scores.\n",
    "from captum.attr import IntegratedGradients\n",
    "\n",
    "num_classes = 2\n",
    "\n",
    "def interpret_sample(features):\n",
    "    return model.captum(X=features)\n",
    "\n",
    "ig = IntegratedGradients(interpret_sample)\n",
    "features.requires_grad_()\n",
    "patch_preds = []\n",
    "for target in range(num_classes):\n",
    "    ig_attr = ig.attribute((features), n_steps=50, target=target)\n",
    "    ig_attr = ig_attr.squeeze().sum(dim=1).cpu().detach()\n",
    "    patch_preds.append(ig_attr)\n",
    "patch_preds = torch.stack(patch_preds, dim=0)\n",
    "\n",
    "print(patch_preds.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f383771",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the histogram of IG scores\n",
    "scores = patch_preds[0, :].detach().numpy()\n",
    "bins = np.linspace(scores.min(), scores.max(), 50)\n",
    "plt.hist(scores, bins, histtype='bar', rwidth=0.8)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "YHQLG265dbLy",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YHQLG265dbLy",
    "outputId": "8b029e9b-a215-46f8-ed02-7d7e6205ba7a"
   },
   "outputs": [],
   "source": [
    "# Heatmap with Captum\n",
    "heatmap = draw_heatmap(\n",
    "  scores=patch_preds[0, :],\n",
    "  coords=coords,\n",
    "  wsi=wsi,\n",
    "  cmap='jet',\n",
    "  alpha=1.0,\n",
    "  use_holes=True,\n",
    "  vis_level=1,\n",
    "  blank_canvas=False,\n",
    "  convert_to_percentiles=False\n",
    ")\n",
    "\n",
    "heatmap.save(os.path.join('data', 'interpretability', 'ig_heatmap.png'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5OWwQUiVdyqt",
   "metadata": {
    "id": "5OWwQUiVdyqt"
   },
   "outputs": [],
   "source": [
    "# Extract the most important patches \n",
    "n_samples = 10 \n",
    "\n",
    "# sort attention weights \n",
    "to_keep = np.argsort(patch_preds[0, :])[-n_samples:]\n",
    "scores_to_keep = attention[to_keep]\n",
    "coords_to_keep = coords[to_keep]\n",
    "\n",
    "print('Indices to keep:', to_keep)\n",
    "print('Scores:', scores_to_keep)\n",
    "\n",
    "for idx in range(n_samples):\n",
    "    patch = wsi.read_region(coords_to_keep[idx], 1, (256, 256))\n",
    "    patch.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a560042",
   "metadata": {},
   "source": [
    "The following link at [http://clam.mahmoodlab.org](http://clam.mahmoodlab.org) visualizes high-attention heatmaps for LUAD vs LUSC subtyping via CLAM (similar to `ABMIL`) and confidence scores for each slides."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "399aedae",
   "metadata": {},
   "source": [
    "#### Discussion\n",
    "\n",
    "- What is the main advantage of IG over attention for model interpretability?\n",
    "\n",
    "\n",
    "- What are the limitations of feature-attribution methods?\n",
    "\n",
    "\n",
    "- What's the difference interpretability and explaianability? Link it to the notion of control.\n",
    "\n",
    "\n",
    "- If you were a clinical pathologist looking at these visualizations, what insights or concerns would you have in letting an AI algorithm assist you medical diagnoses?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3255af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}