WSI-Graph Construction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "### System\n",
    "import os, sys\n",
    "from os.path import join\n",
    "import h5py\n",
    "import math\n",
    "from math import floor\n",
    "import pdb\n",
    "from time import time\n",
    "from tqdm import tqdm\n",
    "\n",
    "### Numerical Packages\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "from scipy.stats import percentileofscore\n",
    "\n",
    "### Graph Network Packages\n",
    "import nmslib\n",
    "import networkx as nx\n",
    "\n",
    "### PyTorch / PyG\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric\n",
    "from torch_geometric.utils import convert\n",
    "\n",
    "### CLAM Path\n",
    "clam_path = 'path/to/Patch-GCN-dir/'\n",
    "sys.path.append(clam_path)\n",
    "from models.resnet_custom import resnet50_baseline\n",
    "from wsi_core.WholeSlideImage import WholeSlideImage\n",
    "from utils.utils import *\n",
    "from datasets.wsi_dataset import Wsi_Region\n",
    "\n",
    "class Hnsw:\n",
    "    def __init__(self, space='cosinesimil', index_params=None,\n",
    "                 query_params=None, print_progress=True):\n",
    "        self.space = space\n",
    "        self.index_params = index_params\n",
    "        self.query_params = query_params\n",
    "        self.print_progress = print_progress\n",
    "\n",
    "    def fit(self, X):\n",
    "        index_params = self.index_params\n",
    "        if index_params is None:\n",
    "            index_params = {'M': 16, 'post': 0, 'efConstruction': 400}\n",
    "\n",
    "        query_params = self.query_params\n",
    "        if query_params is None:\n",
    "            query_params = {'ef': 90}\n",
    "\n",
    "        # this is the actual nmslib part, hopefully the syntax should\n",
    "        # be pretty readable, the documentation also has a more verbiage\n",
    "        # introduction: https://nmslib.github.io/nmslib/quickstart.html\n",
    "        index = nmslib.init(space=self.space, method='hnsw')\n",
    "        index.addDataPointBatch(X)\n",
    "        index.createIndex(index_params, print_progress=self.print_progress)\n",
    "        index.setQueryTimeParams(query_params)\n",
    "\n",
    "        self.index_ = index\n",
    "        self.index_params_ = index_params\n",
    "        self.query_params_ = query_params\n",
    "        return self\n",
    "\n",
    "    def query(self, vector, topn):\n",
    "        # the knnQuery returns indices and corresponding distance\n",
    "        # we will throw the distance away for now\n",
    "        indices, dist = self.index_.knnQuery(vector, k=topn)\n",
    "        return indices\n",
    "\n",
    "def pt2graph(wsi_h5, radius=9):\n",
    "    from torch_geometric.data import Data as geomData\n",
    "    from itertools import chain\n",
    "    coords, features = np.array(wsi_h5['coords']), np.array(wsi_h5['features'])\n",
    "    assert coords.shape[0] == features.shape[0]\n",
    "    num_patches = coords.shape[0]\n",
    "    \n",
    "    model = Hnsw(space='l2')\n",
    "    model.fit(coords)\n",
    "    a = np.repeat(range(num_patches), radius-1)\n",
    "    b = np.fromiter(chain(*[model.query(coords[v_idx], topn=radius)[1:] for v_idx in range(num_patches)]),dtype=int)\n",
    "    edge_spatial = torch.Tensor(np.stack([a,b])).type(torch.LongTensor)\n",
    "    \n",
    "    model = Hnsw(space='l2')\n",
    "    model.fit(features)\n",
    "    a = np.repeat(range(num_patches), radius-1)\n",
    "    b = np.fromiter(chain(*[model.query(coords[v_idx], topn=radius)[1:] for v_idx in range(num_patches)]),dtype=int)\n",
    "    edge_latent = torch.Tensor(np.stack([a,b])).type(torch.LongTensor)\n",
    "\n",
    "    G = geomData(x = torch.Tensor(features),\n",
    "                 edge_index = edge_spatial,\n",
    "                 edge_latent = edge_latent,\n",
    "                 centroid = torch.Tensor(coords))\n",
    "    return G"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# H5 to PyG Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def createDir_h5toPyG(h5_path, save_path):\n",
    "    pbar = tqdm(os.listdir(h5_path))\n",
    "    for h5_fname in pbar:\n",
    "        pbar.set_description('%s - Creating Graph' % (h5_fname[:12]))\n",
    "\n",
    "        try:\n",
    "            wsi_h5 = h5py.File(os.path.join(h5_path, h5_fname), \"r\")\n",
    "            G = pt2graph(wsi_h5)\n",
    "            torch.save(G, os.path.join(save_path, h5_fname[:-3]+'.pt'))\n",
    "            wsi_h5.close()\n",
    "        except OSError:\n",
    "            pbar.set_description('%s - Broken H5' % (h5_fname[:12]))\n",
    "            print(h5_fname, 'Broken')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "h5_path = 'path/to/data_root_dir/tcga_blca_20x_features/h5_files/'\n",
    "save_path = 'path/to/data_root_dir/tcga_blca_20x_features/graph_files/'\n",
    "os.makedirs(save_path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TCGA-G2-A3IB - Creating Graph:   0%|          | 0/457 [00:02<?, ?it/s]\n"
     ]
    }
   ],
   "source": [
    "createDir_h5toPyG(h5_path, save_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}