image_retrieval_and_severity_order.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src import *\n",
    "from src.supcon import *\n",
    "\n",
    "import seaborn as sns\n",
    "import timm\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "CFG.epochs = 50\n",
    "CFG.samples_per_class = 400"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_data = pd.read_csv(os.path.join(DATA_FOLDER, 'trainLabels.csv'))\n",
    "train_data = pd.read_csv(os.path.join(DATA_FOLDER, 'trainLabels_cropped.csv'))\n",
    "train_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove all images from the csv if they are not in the folder\n",
    "lst = map(lambda x: x[:-5], os.listdir(TRAIN_DATA_FOLDER))\n",
    "train_data = train_data[train_data.image.isin(lst)].reset_index(drop=True)\n",
    "train_data = train_data.groupby('level').head(CFG.samples_per_class).reset_index(drop=True)\n",
    "len(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the transformations\n",
    "train_dataset = ImageTrainDataset(TRAIN_DATA_FOLDER, train_data, val_transforms)\n",
    "image, label = train_dataset[15]\n",
    "transformed_img_pil = func.to_pil_image(image)\n",
    "plt.imshow(transformed_img_pil)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Note that the severity levels in test / query images are 0, 1, 2, 3 ie only 4 levels.\n",
    "##### However the severity levels in train images are 0, 1, 2, 3, 4 ie 5 levels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.listdir(TEST_DATA_FOLDER )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "model = SupConModel(timm.create_model(CFG.model_name, pretrained=True, num_classes=0))\n",
    "model.load_state_dict(torch.load(OUTPUT_FOLDER + 'ckpt_epoch_8.pth'))\n",
    "dataloader = DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=CFG.workers)\n",
    "model.eval()\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_level(level):\n",
    "    \"\"\"\n",
    "    converts training level to  query level\n",
    "    \"\"\"\n",
    "    if level == 0:\n",
    "        return 0\n",
    "    else:\n",
    "        return level * 3 / 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieving Nearest Neighbours and Selecting Label Accordingly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "# make vector store\n",
    "res = faiss.StandardGpuResources()\n",
    "config = faiss.GpuIndexFlatConfig()\n",
    "config.device = 1 \n",
    "dimension = 128  # Dimension of the vectors\n",
    "\n",
    "index = faiss.GpuIndexFlatL2(res, dimension, config)\n",
    "k = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all embeddings and store them in the vector store\n",
    "model.eval()\n",
    "embeddings = []\n",
    "with torch.no_grad():\n",
    "    for image, _ in tqdm(dataloader, desc=\"Generating embeddings and indexing\"):\n",
    "        image = image.to(device)\n",
    "        output = model(image)\n",
    "        embedding = output.cpu().numpy().astype('float32')\n",
    "        index.add(embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nearest_images(query_img, k=k):\n",
    "    # query image embeddings\n",
    "    query_image = val_transforms(query_img).unsqueeze(0).to(device)\n",
    "    query_embedding = model(query_image).detach().cpu().numpy().astype('float32')\n",
    "\n",
    "    # search\n",
    "    distances, indices = index.search(query_embedding, k)\n",
    "\n",
    "    # flatten\n",
    "    indices = indices.flatten()\n",
    "    distances = distances.flatten()\n",
    "\n",
    "    print(f\"Indices of {k} nearest images:\", indices)\n",
    "    print(\"Levels of nearest images [0 - 4]:\", [train_data.loc[idx].level for idx in indices])\n",
    "    print(\"Levels in query form [0 - 3]:\", [convert_level(train_data.loc[idx].level) for idx in indices])\n",
    "    print(\"Distances:\", distances)\n",
    "    print(\"Mean distance:\", np.mean(distances))\n",
    "    print(\"Median distance:\", np.median(distances))\n",
    "\n",
    "    print(\"----------------------------------------\")\n",
    "    print(f\"Predicted Level [from 0 to 3]: {np.mean([convert_level(train_data.loc[idx].level) for idx in indices]): .1f}\")\n",
    "    \n",
    "    return indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# taking a random image from the DR 2 class\n",
    "query_img = Image.open(TEST_DATA_FOLDER + '/DR2/1ffa93c6-8d87-11e8-9daf-6045cb817f5b..JPG')\n",
    "indices = get_nearest_images(query_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the nearest images\n",
    "fig, axs = plt.subplots(1, k+1, figsize=(20, 4))\n",
    "\n",
    "# query image\n",
    "axs[0].imshow(query_img)\n",
    "axs[0].axis('off')\n",
    "axs[0].set_title(\"Query Image - level 2\")\n",
    "\n",
    "for i, idx in enumerate(indices):\n",
    "    img = Image.open(TRAIN_DATA_FOLDER + train_data.loc[idx].image + \".jpeg\")\n",
    "    axs[i+1].imshow(img)\n",
    "    axs[i+1].axis('off')\n",
    "    axs[i+1].set_title(f\"Nearest {i+1} - level {train_data.loc[idx].level}\")\n",
    "    \n",
    "plt.suptitle('Predicted Level: 2.2 - using SupCLR')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Disease Severity Order on a Continuous Spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# randomly sample 10 images from the train data folder\n",
    "\n",
    "normal_images = list(train_data[train_data.level == 0].sample(10).image)\n",
    "fig, axs = plt.subplots(2, 5, figsize=(20, 8))\n",
    "\n",
    "axs = axs.ravel()\n",
    "\n",
    "for i, img_name in enumerate(normal_images):\n",
    "    img = Image.open(TRAIN_DATA_FOLDER + img_name + \".jpeg\")\n",
    "    axs[i].imshow(img)\n",
    "    axs[i].axis('off')\n",
    "    axs[i].set_title(f\"Normal Image {i+1}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embeddings(img, model):\n",
    "    img = val_transforms(img).unsqueeze(0).to(device)\n",
    "    output = model(img)\n",
    "    return output\n",
    "\n",
    "def get_severity_order(query_img, model, base_images):\n",
    "    # median euclidean distance of the query image from the base images\n",
    "    \n",
    "    # Get embeddings for the query image and base images in a single batch\n",
    "    query_embedding = get_embeddings(query_img, model)\n",
    "    base_embeddings = torch.cat([get_embeddings(Image.open(TRAIN_DATA_FOLDER + img_name + \".jpeg\"), model)\n",
    "                                   for img_name in base_images])\n",
    "\n",
    "    # Calculate pairwise distances using vectorized operations\n",
    "    distances = torch.norm(query_embedding - base_embeddings, dim=1)\n",
    "\n",
    "    # Get the median distance\n",
    "    severity = torch.median(distances).item()\n",
    "\n",
    "    return severity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "severity = get_severity_order(query_img, model, normal_images)\n",
    "print(f'Order of severity of the disease: {severity:.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Note that this severity order is totally different from the levels we had in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# finding severity for all test images\n",
    "severities = {'0': [], '1': [], '2': [], '3': []}\n",
    "\n",
    "rel = {'0': 'Normal/', '1': 'DR1/', '2': 'DR2/', '3': 'DR3/'}\n",
    "\n",
    "for level in rel:\n",
    "    for img_name in os.listdir(TEST_DATA_FOLDER + rel[level]):\n",
    "        img = Image.open(TEST_DATA_FOLDER + rel[level] + img_name)\n",
    "        severity = get_severity_order(img, model, normal_images)\n",
    "        severities[level].append(severity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a box plot for level vs predicted severity order\n",
    "fig = plt.figure(figsize=(8, 10))\n",
    "sns.swarmplot(data=[severities['0'], severities['1'], severities['2'], severities['3']])\n",
    "plt.xlabel('Level')\n",
    "plt.ylabel('Severity Order')\n",
    "plt.title('Severity Order vs Level')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.savefig(OUTPUT_FOLDER + 'severity_order_sup_con.png', dpi=300, bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}