final_evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d6d5fa71-4824-423b-b5b3-6ee07f9c48b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/guest/.conda/envs/slava_env/lib/python3.10/site-packages/tqdm/auto.py:21: 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 pickle\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import evaluate\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "from transformers import CLIPModel, AutoTokenizer, AutoProcessor\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\" \n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"6\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5c07c393-76c5-4585-a47c-e313232f0748",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_captions = pickle.load(open('processed_data/stimuli_original_captions.pkl', 'rb'))\n",
    "# same for all subjects\n",
    "test_ids = pickle.load(open('processed_data/subj01/sig_test_sub1.pkl', 'rb')).keys()\n",
    "\n",
    "# select only the test captions\n",
    "gt_captions = [gt_captions[idx] for idx in test_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "91261d6c-04c1-486c-8e3f-deb5dbbb6031",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_pickle(path):\n",
    "    return pickle.load(open(path, 'rb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e666e81e-1f04-4223-b5c8-0663c4600aa8",
   "metadata": {},
   "source": [
    "# Meteor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bedd678f-946e-48dc-b2c8-cd2a3e42155b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package wordnet to /home/guest/nltk_data...\n",
      "[nltk_data]   Package wordnet is already up-to-date!\n",
      "[nltk_data] Downloading package punkt to /home/guest/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n",
      "[nltk_data] Downloading package omw-1.4 to /home/guest/nltk_data...\n",
      "[nltk_data]   Package omw-1.4 is already up-to-date!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======SUB01 ONLY========\n",
      "sub                   1.000000\n",
      "dino_vs_gt            0.343938\n",
      "linear_vs_gt          0.267075\n",
      "linear_vs_dino        0.415223\n",
      "widecnn_vs_gt         0.271113\n",
      "widecnn_vs_dino       0.457303\n",
      "shallowcnn_vs_gt      0.268387\n",
      "shallowcnn_vs_dino    0.435663\n",
      "dtype: float64\n",
      "======AVERAGED========\n",
      "sub: 3.75000 2.75379\n",
      "dino_vs_gt: 0.34394 0.00000\n",
      "linear_vs_gt: 0.26291 0.00708\n",
      "linear_vs_dino: 0.41363 0.01083\n",
      "widecnn_vs_gt: 0.27345 0.00813\n",
      "widecnn_vs_dino: 0.45694 0.01552\n",
      "shallowcnn_vs_gt: 0.26669 0.00937\n",
      "shallowcnn_vs_dino: 0.43647 0.02579\n"
     ]
    }
   ],
   "source": [
    "meteor = evaluate.load('meteor')\n",
    "\n",
    "scores = []\n",
    "\n",
    "for sub in [1, 2, 5, 7]:\n",
    "    # dinov2 captions\n",
    "    # its same for all subjects \n",
    "    dinov2_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_dinov2_captions.pkl')\n",
    "    \n",
    "    # linear model [Brain Diffuser]\n",
    "    linear_captions = load_pickle(f'results/sub0{sub}_linear_w_beam_fmri_captions.pkl')\n",
    "\n",
    "    # CNN wide\n",
    "    widecnn_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    # CNN shallow\n",
    "    shallowcnn_captions = load_pickle(f'results/sub0{sub}_shallowcnn_w_beam_fmri_captions.pkl')\n",
    "\n",
    "\n",
    "    scores.append({\n",
    "        \"sub\": sub,\n",
    "        \"dino_vs_gt\": meteor.compute(predictions=dinov2_captions, references=gt_captions)['meteor'],\n",
    "        \"linear_vs_gt\": meteor.compute(predictions=linear_captions, references=gt_captions)['meteor'],\n",
    "        \"linear_vs_dino\": meteor.compute(predictions=linear_captions, references=dinov2_captions)['meteor'],\n",
    "        \"widecnn_vs_gt\": meteor.compute(predictions=widecnn_captions, references=gt_captions)['meteor'],\n",
    "        \"widecnn_vs_dino\": meteor.compute(predictions=widecnn_captions, references=dinov2_captions)['meteor'],\n",
    "        \"shallowcnn_vs_gt\": meteor.compute(predictions=shallowcnn_captions, references=gt_captions)['meteor'],\n",
    "        \"shallowcnn_vs_dino\": meteor.compute(predictions=shallowcnn_captions, references=dinov2_captions)['meteor'],\n",
    "    })\n",
    "\n",
    "\n",
    "scores = pd.DataFrame(scores)\n",
    "print(\"======SUB01 ONLY========\")\n",
    "print(scores[scores['sub']==1].mean())\n",
    "print(\"======AVERAGED========\")\n",
    "for col in scores.columns:\n",
    "    print(f'{col}: {scores[col].mean() :.5f} {scores[col].std() :.5f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e1fcaf6-5c76-4a49-90c6-c70ec0781d7b",
   "metadata": {},
   "source": [
    "# Rouge-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "057d41c4-9df2-44ed-9d46-ced3bcc42d25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======SUB01 ONLY========\n",
      "sub                   1.000000\n",
      "dino_vs_gt            0.415151\n",
      "linear_vs_gt          0.337494\n",
      "linear_vs_dino        0.467228\n",
      "widecnn_vs_gt         0.346206\n",
      "widecnn_vs_dino       0.513230\n",
      "shallowcnn_vs_gt      0.344890\n",
      "shallowcnn_vs_dino    0.491798\n",
      "dtype: float64\n",
      "======AVERAGED========\n",
      "sub: 3.75000 2.75379\n",
      "dino_vs_gt: 0.41515 0.00000\n",
      "linear_vs_gt: 0.33197 0.00902\n",
      "linear_vs_dino: 0.46236 0.01065\n",
      "widecnn_vs_gt: 0.34682 0.00818\n",
      "widecnn_vs_dino: 0.50934 0.01245\n",
      "shallowcnn_vs_gt: 0.34008 0.00893\n",
      "shallowcnn_vs_dino: 0.49201 0.02399\n"
     ]
    }
   ],
   "source": [
    "rouge = evaluate.load('rouge')\n",
    "\n",
    "scores = []\n",
    "\n",
    "for sub in [1, 2, 5, 7]:\n",
    "    # dinov2 captions\n",
    "    # its same for all subjects \n",
    "    dinov2_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_dinov2_captions.pkl')\n",
    "    \n",
    "    # linear model [Brain Diffuser]\n",
    "    linear_captions = load_pickle(f'results/sub0{sub}_linear_w_beam_fmri_captions.pkl')\n",
    "\n",
    "    # CNN wide\n",
    "    widecnn_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    # CNN shallow\n",
    "    shallowcnn_captions = load_pickle(f'results/sub0{sub}_shallowcnn_w_beam_fmri_captions.pkl')\n",
    "\n",
    "\n",
    "    scores.append({\n",
    "        \"sub\": sub,\n",
    "        \"dino_vs_gt\": rouge.compute(predictions=dinov2_captions, references=gt_captions)['rouge1'],\n",
    "        \"linear_vs_gt\": rouge.compute(predictions=linear_captions, references=gt_captions)['rouge1'],\n",
    "        \"linear_vs_dino\": rouge.compute(predictions=linear_captions, references=dinov2_captions)['rouge1'],\n",
    "        \"widecnn_vs_gt\": rouge.compute(predictions=widecnn_captions, references=gt_captions)['rouge1'],\n",
    "        \"widecnn_vs_dino\": rouge.compute(predictions=widecnn_captions, references=dinov2_captions)['rouge1'],\n",
    "        \"shallowcnn_vs_gt\": rouge.compute(predictions=shallowcnn_captions, references=gt_captions)['rouge1'],\n",
    "        \"shallowcnn_vs_dino\": rouge.compute(predictions=shallowcnn_captions, references=dinov2_captions)['rouge1'],\n",
    "    })\n",
    "    \n",
    "    \n",
    "scores = pd.DataFrame(scores)\n",
    "print(\"======SUB01 ONLY========\")\n",
    "print(scores[scores['sub']==1].mean())\n",
    "print(\"======AVERAGED========\")\n",
    "for col in scores.columns:\n",
    "    print(f'{col}: {scores[col].mean() :.5f} {scores[col].std() :.5f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19ca3e58-72ef-4614-b4b7-41f4e3fc91ba",
   "metadata": {},
   "source": [
    "# Rouge-L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6b784e53-55d6-411e-a044-ef3d37dff2d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======SUB01 ONLY========\n",
      "sub                   1.000000\n",
      "dino_vs_gt            0.374883\n",
      "linear_vs_gt          0.306171\n",
      "linear_vs_dino        0.440812\n",
      "widecnn_vs_gt         0.315513\n",
      "widecnn_vs_dino       0.491127\n",
      "shallowcnn_vs_gt      0.316554\n",
      "shallowcnn_vs_dino    0.469370\n",
      "dtype: float64\n",
      "======AVERAGED========\n",
      "sub: 3.75000 2.75379\n",
      "dino_vs_gt: 0.37488 0.00000\n",
      "linear_vs_gt: 0.30107 0.00748\n",
      "linear_vs_dino: 0.43779 0.00965\n",
      "widecnn_vs_gt: 0.31729 0.00724\n",
      "widecnn_vs_dino: 0.48818 0.01206\n",
      "shallowcnn_vs_gt: 0.31216 0.00870\n",
      "shallowcnn_vs_dino: 0.46933 0.02624\n"
     ]
    }
   ],
   "source": [
    "rouge = evaluate.load('rouge')\n",
    "\n",
    "scores = []\n",
    "\n",
    "for sub in [1, 2, 5, 7]:\n",
    "    # dinov2 captions\n",
    "    # its same for all subjects \n",
    "    dinov2_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_dinov2_captions.pkl')\n",
    "    \n",
    "    # linear model [Brain Diffuser]\n",
    "    linear_captions = load_pickle(f'results/sub0{sub}_linear_w_beam_fmri_captions.pkl')\n",
    "\n",
    "    # CNN wide\n",
    "    widecnn_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    # CNN shallow\n",
    "    shallowcnn_captions = load_pickle(f'results/sub0{sub}_shallowcnn_w_beam_fmri_captions.pkl')\n",
    "\n",
    "\n",
    "    scores.append({\n",
    "        \"sub\": sub,\n",
    "        \"dino_vs_gt\": rouge.compute(predictions=dinov2_captions, references=gt_captions)['rougeL'],\n",
    "        \"linear_vs_gt\": rouge.compute(predictions=linear_captions, references=gt_captions)['rougeL'],\n",
    "        \"linear_vs_dino\": rouge.compute(predictions=linear_captions, references=dinov2_captions)['rougeL'],\n",
    "        \"widecnn_vs_gt\": rouge.compute(predictions=widecnn_captions, references=gt_captions)['rougeL'],\n",
    "        \"widecnn_vs_dino\": rouge.compute(predictions=widecnn_captions, references=dinov2_captions)['rougeL'],\n",
    "        \"shallowcnn_vs_gt\": rouge.compute(predictions=shallowcnn_captions, references=gt_captions)['rougeL'],\n",
    "        \"shallowcnn_vs_dino\": rouge.compute(predictions=shallowcnn_captions, references=dinov2_captions)['rougeL'],\n",
    "    })\n",
    "    \n",
    "    \n",
    "scores = pd.DataFrame(scores)\n",
    "print(\"======SUB01 ONLY========\")\n",
    "print(scores[scores['sub']==1].mean())\n",
    "print(\"======AVERAGED========\")\n",
    "for col in scores.columns:\n",
    "    print(f'{col}: {scores[col].mean() :.5f} {scores[col].std() :.5f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fd0cfb8-0717-4e57-8fbe-07c94e8db691",
   "metadata": {},
   "source": [
    "# Sentence transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d93cb2d0-d9a1-464a-a4ba-6e4b8a657ee8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======SUB01 ONLY========\n",
      "sub                   1.000000\n",
      "dino_vs_gt            0.577794\n",
      "linear_vs_gt          0.358235\n",
      "linear_vs_dino        0.428488\n",
      "widecnn_vs_gt         0.396886\n",
      "widecnn_vs_dino       0.498551\n",
      "shallowcnn_vs_gt      0.372515\n",
      "shallowcnn_vs_dino    0.463180\n",
      "dtype: float64\n",
      "======AVERAGED========\n",
      "sub: 3.75000 2.75379\n",
      "dino_vs_gt: 0.57779 0.00000\n",
      "linear_vs_gt: 0.34921 0.01518\n",
      "linear_vs_dino: 0.42038 0.01865\n",
      "widecnn_vs_gt: 0.38912 0.02196\n",
      "widecnn_vs_dino: 0.48418 0.02767\n",
      "shallowcnn_vs_gt: 0.36710 0.02537\n",
      "shallowcnn_vs_dino: 0.45750 0.03300\n"
     ]
    }
   ],
   "source": [
    "scores = []\n",
    "sentence_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
    "\n",
    "for sub in [1, 2, 5, 7]:\n",
    "    # dinov2 captions\n",
    "    # its same for all subjects \n",
    "    dinov2_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_dinov2_captions.pkl')\n",
    "    \n",
    "    # linear model [Brain Diffuser]\n",
    "    linear_captions = load_pickle(f'results/sub0{sub}_linear_w_beam_fmri_captions.pkl')\n",
    "\n",
    "    # CNN wide\n",
    "    widecnn_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    # CNN shallow\n",
    "    shallowcnn_captions = load_pickle(f'results/sub0{sub}_shallowcnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        # select first caption\n",
    "        embedding_gt = sentence_model.encode([i[0] for i in gt_captions], convert_to_tensor=True)\n",
    "        \n",
    "        embedding_dinov2 = sentence_model.encode(dinov2_captions, convert_to_tensor=True)\n",
    "        embedding_linear = sentence_model.encode(linear_captions, convert_to_tensor=True)\n",
    "        embedding_widecnn = sentence_model.encode(widecnn_captions, convert_to_tensor=True)\n",
    "        embedding_shallowcnn = sentence_model.encode(shallowcnn_captions, convert_to_tensor=True)\n",
    "\n",
    "        scores.append({\n",
    "            \"sub\": sub,\n",
    "            \"dino_vs_gt\": util.pytorch_cos_sim(embedding_dinov2, embedding_gt).diag().mean().item(),\n",
    "            \"linear_vs_gt\": util.pytorch_cos_sim(embedding_linear, embedding_gt).diag().mean().item(),\n",
    "            \"linear_vs_dino\": util.pytorch_cos_sim(embedding_linear, embedding_dinov2).diag().mean().item(),\n",
    "            \"widecnn_vs_gt\": util.pytorch_cos_sim(embedding_widecnn, embedding_gt).diag().mean().item(),\n",
    "            \"widecnn_vs_dino\": util.pytorch_cos_sim(embedding_widecnn, embedding_dinov2).diag().mean().item(),\n",
    "            \"shallowcnn_vs_gt\": util.pytorch_cos_sim(embedding_shallowcnn, embedding_gt).diag().mean().item(),\n",
    "            \"shallowcnn_vs_dino\": util.pytorch_cos_sim(embedding_shallowcnn, embedding_dinov2).diag().mean().item(),\n",
    "        })\n",
    "\n",
    "scores = pd.DataFrame(scores)\n",
    "print(\"======SUB01 ONLY========\")\n",
    "print(scores[scores['sub']==1].mean())\n",
    "print(\"======AVERAGED========\")\n",
    "for col in scores.columns:\n",
    "    print(f'{col}: {scores[col].mean() :.5f} {scores[col].std() :.5f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03a67162-072e-4bfb-9f69-9d69d5e388ac",
   "metadata": {},
   "source": [
    "# CLIP-B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dd0fe865-0b0d-437a-a233-37dafa50ee93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======SUB01 ONLY========\n",
      "sub                   1.000000\n",
      "dino_vs_gt            0.773245\n",
      "linear_vs_gt          0.671410\n",
      "linear_vs_dino        0.714764\n",
      "widecnn_vs_gt         0.687283\n",
      "widecnn_vs_dino       0.746361\n",
      "shallowcnn_vs_gt      0.675333\n",
      "shallowcnn_vs_dino    0.731151\n",
      "dtype: float64\n",
      "======AVERAGED========\n",
      "sub: 3.75000 2.75379\n",
      "dino_vs_gt: 0.77325 0.00000\n",
      "linear_vs_gt: 0.66732 0.00605\n",
      "linear_vs_dino: 0.71140 0.00667\n",
      "widecnn_vs_gt: 0.67788 0.01158\n",
      "widecnn_vs_dino: 0.73604 0.01386\n",
      "shallowcnn_vs_gt: 0.67217 0.01318\n",
      "shallowcnn_vs_dino: 0.72788 0.01758\n"
     ]
    }
   ],
   "source": [
    "scores = []\n",
    "\n",
    "model_clip = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "tokenizer =  AutoTokenizer.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "\n",
    "for sub in [1, 2, 5, 7]:\n",
    "    # dinov2 captions\n",
    "    # its same for all subjects \n",
    "    dinov2_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_dinov2_captions.pkl')\n",
    "    \n",
    "    # linear model [Brain Diffuser]\n",
    "    linear_captions = load_pickle(f'results/sub0{sub}_linear_w_beam_fmri_captions.pkl')\n",
    "\n",
    "    # CNN wide\n",
    "    widecnn_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    # CNN shallow\n",
    "    shallowcnn_captions = load_pickle(f'results/sub0{sub}_shallowcnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        # select first caption\n",
    "        embedding_gt = model_clip.get_text_features(**tokenizer([i[0] for i in gt_captions],return_tensors=\"pt\",padding=True))\n",
    "        \n",
    "        embedding_dinov2 = model_clip.get_text_features(**tokenizer(dinov2_captions,return_tensors=\"pt\",padding=True))\n",
    "        embedding_linear = model_clip.get_text_features(**tokenizer(linear_captions,return_tensors=\"pt\",padding=True))\n",
    "        embedding_widecnn = model_clip.get_text_features(**tokenizer(widecnn_captions,return_tensors=\"pt\",padding=True))\n",
    "        embedding_shallowcnn = model_clip.get_text_features(**tokenizer(shallowcnn_captions,return_tensors=\"pt\",padding=True))\n",
    "        \n",
    "        scores.append({\n",
    "            \"sub\": sub,\n",
    "            \"dino_vs_gt\": util.pytorch_cos_sim(embedding_dinov2, embedding_gt).diag().mean().item(),\n",
    "            \"linear_vs_gt\": util.pytorch_cos_sim(embedding_linear, embedding_gt).diag().mean().item(),\n",
    "            \"linear_vs_dino\": util.pytorch_cos_sim(embedding_linear, embedding_dinov2).diag().mean().item(),\n",
    "            \"widecnn_vs_gt\": util.pytorch_cos_sim(embedding_widecnn, embedding_gt).diag().mean().item(),\n",
    "            \"widecnn_vs_dino\": util.pytorch_cos_sim(embedding_widecnn, embedding_dinov2).diag().mean().item(),\n",
    "            \"shallowcnn_vs_gt\": util.pytorch_cos_sim(embedding_shallowcnn, embedding_gt).diag().mean().item(),\n",
    "            \"shallowcnn_vs_dino\": util.pytorch_cos_sim(embedding_shallowcnn, embedding_dinov2).diag().mean().item(),\n",
    "        })\n",
    "    \n",
    "\n",
    "scores = pd.DataFrame(scores)\n",
    "print(\"======SUB01 ONLY========\")\n",
    "print(scores[scores['sub']==1].mean())\n",
    "print(\"======AVERAGED========\")\n",
    "for col in scores.columns:\n",
    "    print(f'{col}: {scores[col].mean() :.5f} {scores[col].std() :.5f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42fc869d-03d7-4fda-85be-2ed938a45c40",
   "metadata": {},
   "source": [
    "# CLIP-L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a19f8988-abc9-47a9-a5d5-a49dc6712c55",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======SUB01 ONLY========\n",
      "sub                   1.000000\n",
      "dino_vs_gt            0.692639\n",
      "linear_vs_gt          0.562937\n",
      "linear_vs_dino        0.633645\n",
      "widecnn_vs_gt         0.584748\n",
      "widecnn_vs_dino       0.673862\n",
      "shallowcnn_vs_gt      0.571342\n",
      "shallowcnn_vs_dino    0.657476\n",
      "dtype: float64\n",
      "======AVERAGED========\n",
      "sub: 3.75000 2.75379\n",
      "dino_vs_gt: 0.69264 0.00000\n",
      "linear_vs_gt: 0.55721 0.00799\n",
      "linear_vs_dino: 0.62834 0.00908\n",
      "widecnn_vs_gt: 0.57591 0.01311\n",
      "widecnn_vs_dino: 0.66478 0.01555\n",
      "shallowcnn_vs_gt: 0.56650 0.01668\n",
      "shallowcnn_vs_dino: 0.65189 0.02274\n"
     ]
    }
   ],
   "source": [
    "scores = []\n",
    "\n",
    "model_clip = CLIPModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "tokenizer =  AutoTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "\n",
    "for sub in [1, 2, 5, 7]:\n",
    "    # dinov2 captions\n",
    "    # its same for all subjects \n",
    "    dinov2_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_dinov2_captions.pkl')\n",
    "    \n",
    "    # linear model [Brain Diffuser]\n",
    "    linear_captions = load_pickle(f'results/sub0{sub}_linear_w_beam_fmri_captions.pkl')\n",
    "\n",
    "    # CNN wide\n",
    "    widecnn_captions = load_pickle(f'results/sub0{sub}_widecnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    # CNN shallow\n",
    "    shallowcnn_captions = load_pickle(f'results/sub0{sub}_shallowcnn_w_beam_fmri_captions.pkl')\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        # select first caption\n",
    "        embedding_gt = model_clip.get_text_features(**tokenizer([i[0] for i in gt_captions],return_tensors=\"pt\",padding=True))\n",
    "        \n",
    "        embedding_dinov2 = model_clip.get_text_features(**tokenizer(dinov2_captions,return_tensors=\"pt\",padding=True))\n",
    "        embedding_linear = model_clip.get_text_features(**tokenizer(linear_captions,return_tensors=\"pt\",padding=True))\n",
    "        embedding_widecnn = model_clip.get_text_features(**tokenizer(widecnn_captions,return_tensors=\"pt\",padding=True))\n",
    "        embedding_shallowcnn = model_clip.get_text_features(**tokenizer(shallowcnn_captions,return_tensors=\"pt\",padding=True))\n",
    "\n",
    "        scores.append({\n",
    "            \"sub\": sub,\n",
    "            \"dino_vs_gt\": util.pytorch_cos_sim(embedding_dinov2, embedding_gt).diag().mean().item(),\n",
    "            \"linear_vs_gt\": util.pytorch_cos_sim(embedding_linear, embedding_gt).diag().mean().item(),\n",
    "            \"linear_vs_dino\": util.pytorch_cos_sim(embedding_linear, embedding_dinov2).diag().mean().item(),\n",
    "            \"widecnn_vs_gt\": util.pytorch_cos_sim(embedding_widecnn, embedding_gt).diag().mean().item(),\n",
    "            \"widecnn_vs_dino\": util.pytorch_cos_sim(embedding_widecnn, embedding_dinov2).diag().mean().item(),\n",
    "            \"shallowcnn_vs_gt\": util.pytorch_cos_sim(embedding_shallowcnn, embedding_gt).diag().mean().item(),\n",
    "            \"shallowcnn_vs_dino\": util.pytorch_cos_sim(embedding_shallowcnn, embedding_dinov2).diag().mean().item(),\n",
    "        })\n",
    "    \n",
    "scores = pd.DataFrame(scores)\n",
    "print(\"======SUB01 ONLY========\")\n",
    "print(scores[scores['sub']==1].mean())\n",
    "print(\"======AVERAGED========\")\n",
    "for col in scores.columns:\n",
    "    print(f'{col}: {scores[col].mean() :.5f} {scores[col].std() :.5f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d128fddb-18b4-4e43-9c58-1dedc46557de",
   "metadata": {},
   "source": [
    "# Some examples sorted by CLIP-L cosine distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1ceb8672-46ea-40ab-8286-43a04df12c50",
   "metadata": {},
   "outputs": [],
   "source": [
    "widecnn_captions = load_pickle(f'results/sub01_widecnn_w_beam_fmri_captions.pkl')\n",
    "\n",
    "with torch.no_grad():\n",
    "    embedding_widecnn = model_clip.get_text_features(**tokenizer(widecnn_captions,return_tensors=\"pt\",padding=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7bd9c0ba-8a7b-44de-8d31-5d85e723167f",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ids = util.pytorch_cos_sim(embedding_widecnn, embedding_gt).diag().sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "48d05df6-1214-4f7c-baa2-4e7ef45697f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True caption:  ['Two giraffes that are standing in the grass.', 'Two giraffes are standing together in the field.']\n",
      "fMRI pred caption:  A couple of giraffes standing in a field.\n",
      "Dinov2 pred caption:  A couple of giraffes standing in a field.\n",
      "==================================\n",
      "True caption:  ['Group of people standing on the side of a busy city street.']\n",
      "fMRI pred caption:  A group of people walking on a city street.\n",
      "Dinov2 pred caption:  A group of people that are standing in the street.\n",
      "==================================\n",
      "True caption:  ['A man on a surfboard in the ocean.']\n",
      "fMRI pred caption:  A person on a surfboard on a body of water.\n",
      "Dinov2 pred caption:  A person on a surfboard in the ocean.\n",
      "==================================\n",
      "True caption:  ['a bath room with a toilet a sink and a mirror']\n",
      "fMRI pred caption:  A bathroom with a sink, toilet and sink.\n",
      "Dinov2 pred caption:  A bathroom with a white toilet and a white sink.\n",
      "==================================\n",
      "True caption:  ['A man on a surf board riding a wave in the ocean.', 'A guy surfing on a wave in the ocean.']\n",
      "fMRI pred caption:  A person on a surfboard in the ocean.\n",
      "Dinov2 pred caption:  A person on a surfboard in the ocean.\n",
      "==================================\n",
      "True caption:  ['A person who is riding a surfboard in the ocean.']\n",
      "fMRI pred caption:  A person on a surfboard in the water.\n",
      "Dinov2 pred caption:  A person on a surfboard in the water.\n",
      "==================================\n",
      "True caption:  ['a large red bus driving down the street']\n",
      "fMRI pred caption:  A red double decker bus driving down the street.\n",
      "Dinov2 pred caption:  A red double decker bus driving down a city street.\n",
      "==================================\n",
      "True caption:  ['A group of people that are sitting around a table.']\n",
      "fMRI pred caption:  A group of people that are sitting at a table.\n",
      "Dinov2 pred caption:  A large group of people sitting around a table.\n",
      "==================================\n",
      "True caption:  ['A man riding a skateboard up the side of a ramp.', 'A man riding a skate board doing a trick on a half pipe.']\n",
      "fMRI pred caption:  A man riding a skateboard down a ramp.\n",
      "Dinov2 pred caption:  A man riding a skateboard down a ramp.\n",
      "==================================\n",
      "True caption:  ['A group of people that are sitting around a table.', 'a number of people sitting at a table with many wine glasses', 'Nine people posing for a picture with a lot of wine glasses on the table.', 'A GROUP OF PEOPLE SITTING AT A BAR IN RED BAR STOOLS DRINKING']\n",
      "fMRI pred caption:  A group of people that are sitting around a table.\n",
      "Dinov2 pred caption:  A group of people that are sitting around a table.\n",
      "==================================\n"
     ]
    }
   ],
   "source": [
    "# good examples\n",
    "for i in ids[-10:]:\n",
    "    print('True caption: ', gt_captions[i])\n",
    "    print('fMRI pred caption: ', widecnn_captions[i])\n",
    "    print('Dinov2 pred caption: ', dinov2_captions[i])\n",
    "    print('==================================')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "36762253-ba5c-4a75-934b-45820b7d0c58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True caption:  ['Green beans, purple beans and lavender arranged along with a scissor on a bamboo mat.']\n",
      "fMRI pred caption:  a close up of a plate of food on a table.\n",
      "Dinov2 pred caption:  a close up of a bottle of water on a table top.\n",
      "==================================\n",
      "True caption:  ['A woman in a tank top holds a rainbow colored umbrella.']\n",
      "fMRI pred caption:  A black and white photo of a black and white dog.\n",
      "Dinov2 pred caption:  a close up of a woman holding a cell phone in her left hand.\n",
      "==================================\n",
      "True caption:  ['A little boy in a blue shirt smiles and stands next to a fire hydrant.', 'A little boy is walking by a fire hydrant.']\n",
      "fMRI pred caption:  A black and white cat laying on a couch.\n",
      "Dinov2 pred caption:  A red fire hydrant in front of a building.\n",
      "==================================\n",
      "True caption:  ['A woman walks down the sidewalk in front of a red wall and a yellow fire hydrant.']\n",
      "fMRI pred caption:  A dog that is sitting on the floor of a living room.\n",
      "Dinov2 pred caption:  A black and white photo of a fire hydrant on the side of the street.\n",
      "==================================\n",
      "True caption:  ['Red exterior door with paper decorations and a white dagoba.']\n",
      "fMRI pred caption:  A large room with a clock on the wall.\n",
      "Dinov2 pred caption:  A stop sign in front of a building.\n",
      "==================================\n",
      "True caption:  ['a dog is visible in the reflection of a mirror', 'a reflection of a dog looking out the window in the side view mirror']\n",
      "fMRI pred caption:  A black and white cat laying on top of a pillow.\n",
      "Dinov2 pred caption:  A reflection of a dog in front of a mirror.\n",
      "==================================\n",
      "True caption:  ['A colorful bird stands on a branch and sings']\n",
      "fMRI pred caption:  A large black and white cat is standing in a field.\n",
      "Dinov2 pred caption:  a close up of a leafy green tree in a field of green trees in a field of green trees in a field of green trees in a field of green trees in a field of green trees in a field of green trees in a field of green trees in a field of green trees in a field of green trees in a field of\n",
      "==================================\n",
      "True caption:  ['A person riding a snowboard down a snowy slope.', 'a person is snowboarding down a snowy hill']\n",
      "fMRI pred caption:  A couple of zebras standing in a field.\n",
      "Dinov2 pred caption:  A group of people skiing down a snowy slope.\n",
      "==================================\n",
      "True caption:  ['A Woman feeding food with her hands to a Giraffe.', 'A girl feeding a giraffe some green leaves.', 'A lady with glasses feeds a giraffe with some leaves']\n",
      "fMRI pred caption:  A black and white dog laying on top of a grassy field.\n",
      "Dinov2 pred caption:  A couple of giraffes standing next to each other.\n",
      "==================================\n",
      "True caption:  ['a person in red wearing a woven hat on a boat with fruits']\n",
      "fMRI pred caption:  A group of people riding on the back of a jeep.\n",
      "Dinov2 pred caption:  A person is on a boat in the water.\n",
      "==================================\n"
     ]
    }
   ],
   "source": [
    "# bad examples\n",
    "for i in ids[:10]:\n",
    "    print('True caption: ', gt_captions[i])\n",
    "    print('fMRI pred caption: ', widecnn_captions[i])\n",
    "    print('Dinov2 pred caption: ', dinov2_captions[i])\n",
    "    print('==================================')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad94883-1afc-4735-b362-0a5f5ba71536",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}