Code/Retrieval-Methods-and-Evaluation/Method-1-Slice-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 sklearn.metrics import average_precision_score\n",
        "import math\n",
        "import csv\n",
        " \n",
        "####################################################\n",
        "## Select one of the organs: liver |colon | pancreas | lung\n",
        "organ=\"liver\" \n",
        "## Select one of the ranking methods: freq | max | sum  \n",
        "rankingMethod = \"freq\"\n",
        "method = 'biomedclip_2d_search_'+rankingMethod\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",
        "# 2D Embeddings \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",
        "ids_organ_path = \"BiomedCLIP/MSD-Final/2d-index/\" + organ + \"-train-2d-index/faiss_ids_embd_train.tsv\"\n",
        "index_organ_path = \"BiomedCLIP/MSD-Final/2d-index/\" + organ + \"-train-2d-index/faiss_index_embd_train.index\"\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700272107500
        }
      }
    },
    {
      "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(self, xq, k=5):\n",
        "        batches = self._split_array_into_batches(xq, 16)\n",
        "        output_D, output_I = [], []\n",
        "        for batch in tqdm(batches, total=len(batches)):\n",
        "            D, I = self.index.search(batch, k)\n",
        "            I = self.get_str_ids_from_I(I)\n",
        "            output_D.append(D)\n",
        "            output_I.append(I)\n",
        "        output_D = np.concatenate(output_D, axis=0)\n",
        "        output_I = np.concatenate(output_I, axis=0)\n",
        "        return output_D, output_I\n",
        "\n",
        "index = FAISSRetrievalEngine(\n",
        "    ids_path=ids_organ_path,\n",
        "    index_path=index_organ_path\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700272110980
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Running the Image Search for all test queries  \n",
        "**"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from collections import Counter    \n",
        "  \n",
        "def sort_by_frequency(lst):  \n",
        "    freq_dict = Counter(lst)   \n",
        "    return sorted(lst, key=lambda x: (-freq_dict[x], x))  \n",
        "\n",
        "def sort_by_sum_scores(results):  \n",
        "    sorted_ids = sorted(results.keys(), key=lambda x: -np.array(results[x]).sum())\n",
        "    return sorted_ids \n",
        "\n",
        "def sort_by_max_scores(results): \n",
        "    sorted_ids = sorted(results.keys(), key=lambda x: -np.array(results[x]).max())\n",
        "    return sorted_ids \n",
        "    \n",
        "def get_3D_results_w_freq(slice_id_arr):\n",
        "    \"\"\"Baseline to get 3D img results from 2D slice results.\n",
        "    Uses 2D slice count of each 3D image in the result list.\n",
        "\n",
        "    input:\n",
        "     [\n",
        "        [slice_1_1, slice 1_2, ..], #results for query slice 1\n",
        "        [slice_2_1, slice 2_2, ..], #results for query slice 2\n",
        "        ...\n",
        "     ]\n",
        "\n",
        "     output:\n",
        "     [ img_id_1, img_id_2, .., img_id_n] # initial 3D images ranked by slice count\n",
        "\n",
        "    \"\"\"\n",
        "    results = []\n",
        "\n",
        "    for i, slice_ids in enumerate(slice_id_arr):\n",
        "        for slice_id in slice_ids:\n",
        "            slice_str = str(slice_id)\n",
        "            img_id = slice_str[:slice_str.rindex('_')]\n",
        "            results.append(img_id)\n",
        "\n",
        "    results = sort_by_frequency(results)\n",
        "\n",
        "    result_set = set()\n",
        "    final_results = []\n",
        "\n",
        "    for r in results:\n",
        "        if r not in result_set:\n",
        "            final_results.append(r)\n",
        "            result_set.add(r)\n",
        "\n",
        "    return final_results\n",
        "\n",
        "def get_3D_results_w_scores(slice_id_arr, slice_score_arr):\n",
        "    results = {}\n",
        "    for i, slice_ids in enumerate(slice_id_arr):\n",
        "        for j, slice_id in enumerate(slice_ids):\n",
        "            slice_str = str(slice_id)\n",
        "            img_id = slice_str[:slice_str.rindex('_')]\n",
        "            if img_id not in results:\n",
        "                results[img_id] = []\n",
        "            \n",
        "            results[img_id].append(slice_score_arr[i][j])\n",
        "\n",
        "    if rankingMethod == \"max\":  \n",
        "        results = sort_by_max_scores(results)\n",
        "    else: \n",
        "        if rankingMethod == \"sum\":  \n",
        "            results = sort_by_sum_scores(results)\n",
        "\n",
        "    result_set = set()\n",
        "    final_results = []\n",
        "\n",
        "    for r in results:\n",
        "        if r not in result_set:\n",
        "            final_results.append(r)\n",
        "            result_set.add(r)\n",
        "\n",
        "    return final_results"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700272111138
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Slice-based Search\n",
        "\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",
        "\n",
        "    search_results['query_img'].append(query_img_id)\n",
        "    with open(test_volume_emb, \"rb\") as f:\n",
        "        embd = pickle.load(f)\n",
        "    \n",
        "    ## retrieving 20 slices for eval\n",
        "    results = {}\n",
        "    \n",
        "    D, I = index.retrieve_images(embd,20)\n",
        "    if rankingMethod == \"freq\":  \n",
        "        results = get_3D_results_w_freq(I)\n",
        "    else: \n",
        "        if rankingMethod in (\"sum\", \"max\"):\n",
        "            results = get_3D_results_w_scores(I, D)\n",
        "\n",
        "    search_results['results'].append(' '.join(results))\n",
        "\n",
        "df = pd.DataFrame(search_results)\n",
        "\n",
        "df.to_csv(results_dir + '/' + method + '_results.csv', index=False)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700272121284
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "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": 1700272121461
        }
      }
    },
    {
      "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",
        "\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.split(' ')\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 + '_tumor_stage_eval.csv', index=False) \n",
        "tumor_flag_df = pd.DataFrame(tumor_flag_eval)\n",
        "tumor_flag_df.to_csv(results_dir + '/' + method + '_tumor_flag_eval.csv', index=False) \n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1700272121610
        },
        "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": 1700272121765
        }
      }
    },
    {
      "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.split(' ')\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"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700272121913
        }
      }
    },
    {
      "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",
        "ap = 0\n",
        "\n",
        "for i, q in enumerate(ref_tumor_flag):\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)))"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700272122061
        }
      }
    },
    {
      "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",
        "    try:\n",
        "        stage_avg_precision =  average_precision_score(all_res_tumor_stage[i], list(range(len(all_res_tumor_flag[i])))[::-1])\n",
        "        if math.isnan(stage_avg_precision):\n",
        "            stage_avg_precision=0\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)))"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1700272122219
        }
      }
    }
  ],
  "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
}