class_specific_CAMs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rushing/miniconda3/envs/GMT/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os, sys\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"3\"\n",
    "\n",
    "import torch\n",
    "from torch_geometric.loader import DataLoader\n",
    "import torch_geometric.transforms as T\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "import PIL\n",
    "PIL.Image.MAX_IMAGE_PIXELS = 933120000\n",
    "\n",
    "from gnn import GNN\n",
    "from visualize import generate_relevance, fetch_slide_image\n",
    "from util import read_file, find_dataset_using_name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    device = 0\n",
    "    gnn = 'gin'\n",
    "    num_layer = 3\n",
    "    emb_dim = 64\n",
    "    drop_ratio = 0\n",
    "    jk = 'sum'\n",
    "    graph_pooling = 'gmt'\n",
    "\n",
    "    seed = 1078\n",
    "    batch_size = 1\n",
    "    num_workers = 0\n",
    "    phase = 'cams'\n",
    "    n_classes = 3\n",
    "    data_config = 'ctranspath_files'\n",
    "    fdim = 768\n",
    "    patch_size = 256\n",
    "    output = 'logs'\n",
    "\n",
    "    run_name = \"Graph-Perciever_September-17\"\n",
    "    fold_idx = 2\n",
    "    dataset = ['cptac']\n",
    "    index = None # 0 - Normal | 1 - LUSC | 2 - LUAD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import zscore\n",
    "from sklearn.preprocessing import MinMaxScaler, StandardScaler\n",
    "def plot_heat_maps(graph, scores, prob, slide_root, patch_size=256, overlay=True, clamp=0.05, norm=True, colormap='RdBu_r', save_path=None):\n",
    "    # Fetch patch coords & slide path for the tissue\n",
    "    slide_path = graph.slide_path[0]\n",
    "    coords = graph.node_coords\n",
    "    # coords = [(int(x), int(y)) for x,y in coords]\n",
    "\n",
    "    # fetch tissue image at specific downsample\n",
    "    downsample_factor = 16.0\n",
    "    image = fetch_slide_image(slide_path, slide_root, patch_size, downsample_factor=downsample_factor, gt = graph.y, save_path=save_path)\n",
    "    image = np.asarray(image.convert(\"RGB\"))\n",
    "    image = (image - image.min()) / (image.max() - image.min())\n",
    "\n",
    "    y_min, y_max, x_min, x_max = 0, image.shape[0], 0, image.shape[1]\n",
    "\n",
    "    mask = np.zeros((image.shape[0], image.shape[1]), dtype=bool) # this is the cam mask\n",
    "    heatmap = -np.ones(image.shape[:2], dtype=np.float32)\n",
    "\n",
    "    offset = patch_size + 2 # 2 is for overlap\n",
    "    d = downsample_factor\n",
    "    scores = scores.numpy()\n",
    "    scores = zscore(scores)\n",
    "\n",
    "    if clamp:\n",
    "        q05, q95 = np.quantile(scores, clamp), np.quantile(scores, 1-clamp)\n",
    "        scores = np.clip(scores, a_min=q05, a_max=q95)\n",
    "\n",
    "    # check if all values in scores are 0s\n",
    "    scores = np.nan_to_num(scores, nan=0)\n",
    "    if not np.all(scores == -1):\n",
    "        scores = MinMaxScaler(feature_range=(-1, 1)).fit_transform(scores.reshape(-1,1))\n",
    "    \n",
    "    for (x, y), s in zip(coords, scores):\n",
    "\n",
    "        # x, y = int(x)*512, int(y)*512\n",
    "        x, y = x*patch_size, y*patch_size\n",
    "        \n",
    "        mask[round(y.item()/d):round((y.item()+offset)/d), round(x.item()/d):round((x.item()+offset)/d)] = True\n",
    "        heatmap[round(y.item()/d):round((y.item()+offset)/d), round(x.item()/d):round((x.item()+offset)/d)] = s\n",
    "\n",
    "    plt.figure(figsize=(30, 30))\n",
    "    a = 1.\n",
    "    if overlay:\n",
    "        plt.imshow(image, alpha=1, cmap='gray')\n",
    "        a = 0.7\n",
    "\n",
    "    plt.imshow(heatmap, alpha=0.5*mask, cmap=colormap, interpolation='nearest')\n",
    "    cbar = plt.colorbar(location='right', orientation='vertical')\n",
    "    cbar.ax.tick_params(labelsize=40)\n",
    "    plt.axis('off')\n",
    "\n",
    "    return plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_cams(args, model, device, multiple_loaders, index=None):\n",
    "\n",
    "    # print('Evaluating...', args.dataset.upper())\n",
    "    # print(\"Dataset length\", len(loader))\n",
    "    model.eval()\n",
    "\n",
    "    y_true = []\n",
    "    y_pred = []\n",
    "\n",
    "    os.makedirs(args.output_folder, exist_ok=True)\n",
    "\n",
    "    for loader in multiple_loaders:\n",
    "\n",
    "        true_labels = list(loader.dataset.classdict.keys())\n",
    "        to_be_predicted_classes = list(loader.dataset.to_be_predicted_classes.keys())\n",
    "\n",
    "        for step, graph in enumerate(tqdm(loader, desc=\"Iteration\")):\n",
    "\n",
    "            graph = graph.to(device)\n",
    "            slide_name = graph.slide_path[0]\n",
    "            print(slide_name)\n",
    "            # print(graph.node_coords)\n",
    "\n",
    "\n",
    "            if graph.x.shape[0] == 1:\n",
    "                pass\n",
    "            else:\n",
    "\n",
    "                # GENERATE VISUALIZATION :\n",
    "                transformer_attribution, output, y_pred = generate_relevance(model, graph, index=index)\n",
    "                if index is not None:\n",
    "                    y_pred = index\n",
    "\n",
    "                print(\"logits: \", output)\n",
    "\n",
    "                prob = F.softmax(output, dim=1)\n",
    "                prob = prob.squeeze()\n",
    "\n",
    "                print(\"Slide: {}, True Class: {}, Predicted Class: {}(p={:.3f})\".format(slide_name, true_labels[graph.y], to_be_predicted_classes[y_pred], prob[y_pred].item()))\n",
    "                \n",
    "                del output\n",
    "\n",
    "                ########################################\n",
    "                slide_root = os.path.join('/SeaExp/Rushin/datasets/', args.dataset_name.upper(), 'WSIs')\n",
    "                plt = plot_heat_maps(graph, scores=transformer_attribution, prob=prob[index], slide_root=slide_root, clamp=0.05, save_path=args.output_folder, overlay=True) # attention_blend = \n",
    "                # Use numpy to save attention_blend image to a file\n",
    "                # attention_blend = Image.fromarray(attention_blend)\n",
    "                ########################################\n",
    "                \n",
    "                # plt.axis('off')\n",
    "                plt.savefig(os.path.join(args.output_folder, \"{}_{}(fold{})_cam.png\".format(slide_name, to_be_predicted_classes[y_pred], args.fold_idx)))\n",
    "                plt.close()\n",
    "                # attention_blend.save(os.path.join(args.output_folder, \"{}_{}(fold{})_cam.png\".format(slide_name, to_be_predicted_classes[y_pred], args.fold_idx)))\n",
    "                # break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def implement(index=None):\n",
    "\n",
    "    args = Args()\n",
    "    if index is not None:\n",
    "        args.index = index\n",
    "\n",
    "    ### set up seeds and gpu device\n",
    "    torch.manual_seed(args.seed)\n",
    "    torch.cuda.manual_seed(args.seed)\n",
    "    torch.cuda.manual_seed_all(args.seed)\n",
    "    np.random.seed(args.seed)\n",
    "    random.seed(args.seed)\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "    device = torch.device(\"cuda:\" + str(args.device)) if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "    args.slide_feats_folder = {}\n",
    "\n",
    "    fold_idx = args.fold_idx\n",
    "    test_loaders = []\n",
    "\n",
    "    for item in args.dataset:\n",
    "\n",
    "        args.dataset_name = item\n",
    "        # args.slide_feats_folder[item] = os.path.join('/SeaExp/Rushin/datasets', item.upper(), 'slide_features')\n",
    "        # os.makedirs(args.slide_feats_folder[item], exist_ok=True)\n",
    "\n",
    "        dataset_class = find_dataset_using_name(item)\n",
    "        print(dataset_class)\n",
    "\n",
    "        ### automatic dataloading and splitting\n",
    "        root = os.path.join('/SeaExp/Rushin/datasets', item.upper(), args.data_config)\n",
    "        wsi_file = os.path.join('/SeaExp/Rushin/datasets', item.upper(), '%s_%s.txt' % (item.upper(), args.phase))\n",
    "        wsi_ids = read_file(wsi_file)\n",
    "\n",
    "        dataset = dataset_class(root, wsi_ids, args.fdim, n_classes=args.n_classes, isTrain=False, transform=T.ToSparseTensor(remove_edge_index=False))\n",
    "        test_loaders.append(DataLoader(dataset, batch_size=1, shuffle=False, num_workers=args.num_workers))\n",
    "\n",
    "    log_path = os.path.join('logs', \"{}_fold_{}\".format(args.run_name, fold_idx))\n",
    "\n",
    "    model = GNN(gnn_type = args.gnn, num_class = dataset.num_classes, num_layer = args.num_layer, input_dim = args.fdim, emb_dim = args.emb_dim, drop_ratio = args.drop_ratio, JK = args.jk, graph_pooling = args.graph_pooling).to(device)\n",
    "\n",
    "    # model.load_state_dict(torch.load(os.path.join(log_path, \"epoch_30_model_{}_fold_{}.pth\".format(args.run_name, fold_idx))))\n",
    "    model.load_state_dict(torch.load(os.path.join(log_path, \"final_model_{}_fold_{}.pth\".format(args.run_name, fold_idx))))\n",
    "    model = model.to(device)\n",
    "    print(\"model weights loaded successfully\")\n",
    "\n",
    "    total_params = sum(p.numel() for p in model.parameters())\n",
    "    print('Total params:', total_params)\n",
    "\n",
    "\n",
    "    # args.output_folder = os.path.join(args.output, args.run_name+\"_{}_{}\".format(args.patch_size, args.phase), args.dataset_name+f'_final_epoch')\n",
    "\n",
    "    # generate_cams(args, model, device, test_loaders, args.index)\n",
    "\n",
    "    # del model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dataloaders.cptac.CptacDataset'>\n",
      "model weights loaded successfully\n",
      "Total params: 202292\n"
     ]
    }
   ],
   "source": [
    "# implement(index=0)\n",
    "# implement(index=1)\n",
    "implement(index=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GMT",
   "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.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}