Code/Retrieval-Methods-and-Evaluation/Method-2-Volume-based-Retrieval.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import faiss\n",
        "import pandas as pd\n",
        "import torch\n",
        "from tqdm.notebook import tqdm\n",
        "import glob \n",
        "import pickle\n",
        "import numpy as np\n",
        "from random import sample\n",
        "import os\n",
        "from collections import OrderedDict \n",
        "import csv\n",
        "from sklearn.metrics import average_precision_score\n",
        "import math\n",
        "####################################################\n",
        "## Select one of the organs: liver |colon | pancreas | lung\n",
        "organ=\"liver\"\n",
        "## Select one of the aggregation methods: median | maxpooling | avgpooling |std\n",
        "method = \"median\"  # median / maxpooling / avgpooling / std\n",
        "####################################################\n",
        "\n",
        "## Quantification Files\n",
        "quantification_dir = \"data/msd_final/quantification_file\" \n",
        "quantification_csv = quantification_dir + \"/msd_\" + organ + \"_metrics_3D_subject_including_normal_tissue.csv\"\n",
        "\n",
        "## Data \n",
        "data_splits_folder = \"data/msd_train_test_split\"\n",
        "if (organ == \"lung\"):\n",
        "    data_splits_folder = \"data/msd_train_test_split/lung\"\n",
        "\n",
        "all_embs_folder = \"data/msd_embeddings/biomedclip_2d\"\n",
        "\n",
        "test_set_csv =  data_splits_folder + \"/\" + organ + \"/\" + organ + \"_full_test_split.csv\"\n",
        "\n",
        "## Output files \n",
        "results_dir = \"BiomedCLIP/MSD-Final/results/results_\" + organ\n",
        "results_eval_metrics = results_dir + \"/\" + organ + \"eval_metrics.csv\"\n",
        "\n",
        "## Index files\n",
        "train_3d_index = \"BiomedCLIP/MSD-Final/3d-index/\" + method + \"/\" + organ + \"-train-3d-index/faiss_index_embd_train.index\"\n",
        "train_3d_index_tsv = \"BiomedCLIP/MSD-Final/3d-index/\" + method + \"/\" + organ + \"-train-3d-index/faiss_ids_embd_train.tsv\"\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700279794722
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class FAISSRetrievalEngine:\n",
        "    def __init__(self, ids_path, index_path):\n",
        "        self.ids_path = ids_path\n",
        "        self.ids = pd.read_csv(self.ids_path, sep=\"\\t\")\n",
        "        self.ids = {row[\"ids_int\"]: row[\"ids_str\"] for i, row in self.ids.iterrows()}\n",
        "\n",
        "        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "        self.index_path = index_path\n",
        "        self.index = faiss.read_index(self.index_path)\n",
        "        if self.device == \"cuda\":\n",
        "            res = faiss.StandardGpuResources()\n",
        "            self.index = faiss.index_cpu_to_gpu(res, 0, self.index)\n",
        "    \n",
        "    def _split_array_into_batches(self, arr, batch_size):\n",
        "        num_rows = arr.shape[0]\n",
        "        num_batches = (num_rows + batch_size - 1) // batch_size  \n",
        "        \n",
        "        # Create a list of batches\n",
        "        batches = []\n",
        "        for i in range(num_batches):\n",
        "            start_idx = i * batch_size\n",
        "            end_idx = min((i + 1) * batch_size, num_rows)\n",
        "            batch = arr[start_idx:end_idx, :]\n",
        "            batches.append(batch)\n",
        "\n",
        "        return batches\n",
        "\n",
        "    def get_str_ids_from_I(self, I):\n",
        "        mapping_fn = lambda x: self.ids[x]\n",
        "        return np.vectorize(mapping_fn)(I)\n",
        "\n",
        "    def retrieve_images_with_3D_emb(self, xq, k=5):\n",
        "        output_D, output_I = [], []\n",
        "        D, I = self.index.search(xq, k)\n",
        "        I_str = self.get_str_ids_from_I(I)\n",
        "        output_D.append(D)\n",
        "        output_I.append(I_str)\n",
        "        return output_D, output_I\n",
        "\n",
        "\n",
        "index = FAISSRetrievalEngine(\n",
        "    ids_path=train_3d_index_tsv,\n",
        "    index_path=train_3d_index\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700279795447
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Running the 3D Image Search for all test queries  \n",
        "**"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "aggregated_df = pd.read_csv(quantification_csv) \n",
        "aggregated_df['name'] = aggregated_df['name'].replace('.nii.gz', '')\n",
        "for i in range(len(aggregated_df['name'])):\n",
        "    filename = aggregated_df['name'][i]\n",
        "    aggregated_df['name'][i] = filename[:filename.index('.')] \n",
        "aggregated_df.head()"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700279795914
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Search \n",
        "search_results = {'query_img': [], 'results': []}\n",
        "## Embeddings of the test set \n",
        "test_csv = pd.read_csv(test_set_csv)\n",
        "\n",
        "for img in test_csv['testing']: \n",
        "    query_img_id = img[:img.index('.')]\n",
        "    test_volume_emb = os.path.join(all_embs_folder, query_img_id + '.pkl')\n",
        "    search_results['query_img'].append(query_img_id)\n",
        "    with open(test_volume_emb, \"rb\") as f:\n",
        "        embd = pickle.load(f)\n",
        "    emb_3d = []\n",
        "    if (method == \"maxpooling\"): \n",
        "        emb_3d = np.max(embd, axis=0)\n",
        "    if (method == \"avgpooling\"): \n",
        "        emb_3d = np.mean(embd, axis=0)\n",
        "    if (method == \"std\"): \n",
        "        emb_3d = np.std(embd, axis=0)\n",
        "    if (method == \"median\"): \n",
        "        emb_3d = np.median(embd, axis=0)\n",
        "\n",
        "    D, I = index.retrieve_images_with_3D_emb(np.array([emb_3d]), 50) ## retrieving 20 volumes\n",
        "    results = [str(x) for x in I[0][0]]\n",
        "    resultswoDup = list(OrderedDict.fromkeys(results)) \n",
        "    if len(resultswoDup) > 10:\n",
        "        resultswoDup = resultswoDup[:10]\n",
        "\n",
        "    search_results['results'].append(resultswoDup)\n",
        "\n",
        "df = pd.DataFrame(search_results)\n",
        "df.to_csv(results_dir + '/' + method + '_3d_results.csv', index=False)\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700279796821
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Evaluation\n",
        "\n",
        "tumor_stage_eval  = {'query': [] } \n",
        "tumor_flag_eval = {'query': [] }\n",
        "\n",
        "for i in range(10):\n",
        "    tumor_stage_eval[f\"Top {i+1}\"] = []\n",
        "    tumor_flag_eval[f\"Top {i+1}\"] = []\n",
        "\n",
        "for query, res in zip(search_results['query_img'], search_results['results']):\n",
        "    query_metrics = aggregated_df.loc[aggregated_df['name'] == query].iloc[0]\n",
        "    \n",
        "    tumor_stage_eval['query'].append(query)\n",
        "    tumor_flag_eval['query'].append(query)\n",
        "\n",
        "    tumor_flag = organ + \"_cancer_flag\" \n",
        "    top_n = res\n",
        "    \n",
        "    for k, top_k in enumerate(top_n):\n",
        "        if k > 9:\n",
        "            break\n",
        "\n",
        "        res_metrics = aggregated_df.loc[aggregated_df['name'] == top_k].iloc[0]\n",
        "        # tumor stage\n",
        "        tumorStage=0\n",
        "        if(query_metrics['cancer_stage']-res_metrics['cancer_stage'] == 0):\n",
        "            tumorStage=1 \n",
        "\n",
        "        tumor_stage_eval[f\"Top {k+1}\"].append(tumorStage)\n",
        "        tumor_flag_eval[f\"Top {k+1}\"].append(1-abs(query_metrics[tumor_flag]-res_metrics[tumor_flag]))\n",
        "        \n",
        "    if len(top_n) < 10:\n",
        "        for k in range(len(top_n), 10):\n",
        "            tumor_stage_eval[f\"Top {k+1}\"].append('NA')\n",
        "            tumor_flag_eval[f\"Top {k+1}\"].append('NA')\n",
        "\n",
        "tumor_stage_df = pd.DataFrame(tumor_stage_eval)\n",
        "tumor_stage_df.to_csv(results_dir + '/' + method + '_3d_tumor_stage_eval.csv', index=False) \n",
        "tumor_flag_df = pd.DataFrame(tumor_flag_eval)\n",
        "tumor_flag_df.to_csv(results_dir + '/' + method + '_3d_tumor_flag_eval.csv', index=False)  "
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700279796960
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Tumor Flag\n",
        "print(\"# Tumor Flag: \")\n",
        "\n",
        "tumor_flag_df['RowMean'] = tumor_flag_df.iloc[:, 1:4].mean(axis=1)\n",
        "overall_mean = tumor_flag_df['RowMean'].mean()\n",
        "print(\"tumor_flag_p@3: \", '{:.4f}'.format(overall_mean))\n",
        "\n",
        "tumor_flag_df['RowMean'] = tumor_flag_df.iloc[:, 1:6].mean(axis=1)\n",
        "overall_mean = tumor_flag_df['RowMean'].mean()\n",
        "print(\"tumor_flag_p@5: \", '{:.4f}'.format(overall_mean))\n",
        "\n",
        "tumor_flag_df['RowMean'] = tumor_flag_df.iloc[:, 1:].mean(axis=1)\n",
        "overall_mean = tumor_flag_df['RowMean'].mean()\n",
        "print(\"tumor_flag_p@10: \", '{:.4f}'.format(overall_mean))\n",
        "\n",
        "## Tumor Stage\n",
        "print(\"# Tumor Stage: \")\n",
        "\n",
        "tumor_stage_df['RowMean'] = tumor_stage_df.iloc[:, 1:4].mean(axis=1)\n",
        "overall_mean = tumor_stage_df['RowMean'].mean()\n",
        "print(\"tumor_stage_p@3: \", '{:.4f}'.format(overall_mean))\n",
        "\n",
        "tumor_stage_df['RowMean'] = tumor_stage_df.iloc[:, 1:6].mean(axis=1)\n",
        "overall_mean = tumor_stage_df['RowMean'].mean()\n",
        "print(\"tumor_stage_p@5: \", '{:.4f}'.format(overall_mean))\n",
        "\n",
        "tumor_stage_df['RowMean'] = tumor_stage_df.iloc[:, 1:].mean(axis=1)\n",
        "overall_mean = tumor_stage_df['RowMean'].mean()\n",
        "print(\"tumor_stage_p@10: \", '{:.4f}'.format(overall_mean))"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700279797101
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Average Precision"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ref_tumor_flag = []\n",
        "all_res_tumor_flag = []\n",
        "ref_tumor_stage = []\n",
        "all_res_tumor_stage = []\n",
        "\n",
        "for query, res in zip(search_results['query_img'], search_results['results']):\n",
        "    query_metrics = aggregated_df.loc[aggregated_df['name'] == query].iloc[0]\n",
        "    tumor_flag = organ + \"_cancer_flag\" \n",
        "    ref_tumor_flag.append(query_metrics[tumor_flag])\n",
        "    ref_tumor_stage.append(query_metrics['cancer_stage'])\n",
        "    top_n = res\n",
        "    res_tumor_flag = []\n",
        "    res_tumor_stage = []\n",
        "    for k, top_k in enumerate(top_n):\n",
        "        if k > 9:\n",
        "            break\n",
        "        res_metrics = aggregated_df.loc[aggregated_df['name'] == top_k].iloc[0]\n",
        "        res_tumor_flag.append(res_metrics[tumor_flag])\n",
        "        res_tumor_stage.append(res_metrics['cancer_stage'])\n",
        "    all_res_tumor_flag.append(res_tumor_flag)\n",
        "    all_res_tumor_stage.append(res_tumor_stage)\n",
        "\n",
        "## Adaptation for Average Precision (AP) computation\n",
        "## AP - Tumor Flag\n",
        "for i, q in enumerate(ref_tumor_flag):\n",
        "    if q == 0:\n",
        "        for j, res in enumerate(all_res_tumor_flag[i]):\n",
        "            if res == 0:\n",
        "                all_res_tumor_flag[i][j] = 1\n",
        "            else:\n",
        "                all_res_tumor_flag[i][j] = 0\n",
        "\n",
        "## AP - Tumor Stage\n",
        "for i, q in enumerate(ref_tumor_stage):\n",
        "    print(q)\n",
        "    for j, res in enumerate(all_res_tumor_stage[i]):\n",
        "        if res == q:\n",
        "            all_res_tumor_stage[i][j] = 1\n",
        "        else:\n",
        "            all_res_tumor_stage[i][j] = 0\n",
        "            "
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700279797477
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "##############\n",
        "## Tumor Flag\n",
        "##############\n",
        "\n",
        "ref_tumor_flag = np.array(ref_tumor_flag)\n",
        "all_res_tumor_flag = np.array(all_res_tumor_flag)\n",
        "\n",
        "ap = 0\n",
        "for i, q in enumerate(ref_tumor_flag):\n",
        "    \n",
        "    try:\n",
        "        flag_avg_precision =  average_precision_score(all_res_tumor_flag[i], list(range(len(all_res_tumor_flag[i])))[::-1])\n",
        "        \n",
        "        if math.isnan(flag_avg_precision):\n",
        "            flag_avg_precision=0\n",
        "\n",
        "        ap = ap + flag_avg_precision\n",
        "    except:\n",
        "        print(all_res_tumor_flag[i])\n",
        "\n",
        "print(\"flag_avg_precision: \", '{:.4f}'.format(ap/len(ref_tumor_flag)))\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700279797667
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "###############\n",
        "## Tumor Stage\n",
        "###############\n",
        "\n",
        "ref_tumor_stage = np.array(ref_tumor_stage)\n",
        "all_res_tumor_stage = np.array(all_res_tumor_stage)\n",
        "ap = 0\n",
        "\n",
        "for i, q in enumerate(ref_tumor_stage):\n",
        "    \n",
        "    try:\n",
        "        stage_avg_precision =  average_precision_score(all_res_tumor_stage[i], list(range(len(all_res_tumor_flag[i])))[::-1])\n",
        "        \n",
        "        if math.isnan(stage_avg_precision):\n",
        "            stage_avg_precision=0\n",
        "        \n",
        "        ap = ap + stage_avg_precision\n",
        "    except:\n",
        "        print(all_res_tumor_stage[i])\n",
        "\n",
        "print(\"stage_avg_precision: \", '{:.4f}'.format(ap/len(ref_tumor_stage)))\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700279797899
        }
      }
    }
  ],
  "metadata": {
    "kernel_info": {
      "name": "python38-azureml"
    },
    "kernelspec": {
      "name": "python38-azureml",
      "language": "python",
      "display_name": "Python 3.8 - AzureML"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.5",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "microsoft": {
      "host": {
        "AzureML": {
          "notebookHasBeenCompleted": true
        }
      },
      "ms_spell_check": {
        "ms_spell_check_language": "en"
      }
    },
    "nteract": {
      "version": "nteract-front-end@1.0.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}