{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "d6bbe27d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "import yaml\n",
    "import argparse\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.io as sio\n",
    "import random\n",
    "from copy import deepcopy\n",
    "from tqdm import tqdm, trange\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import cortex\n",
    "\n",
    "import torch\n",
    "import utils\n",
    "import nibabel as nib\n",
    "import h5py\n",
    "from scipy import stats\n",
    "\n",
    "# tf32 data type is faster than standard float32\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "# torch.backends.cudnn.benchmark = True # fixes Conv3D if used\n",
    "# torch.autograd.set_detect_anomaly(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7ca247d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--name ddp2-0.01-qwen25vl-s1-CLIP-ViT-H-14-mask=brain-blocks=12-patch=conv_14-pt_dp=0-at_dp=0-bo_dp=0-txt_scl=1-m_pct=0-bs=128-lr=0.0003-bf=4096-norm_nii --checkpoints_dir ../checkpoints/ --seed 42 --batch_size 32 --subj_list 1 --nsddir /opt/data/private/dataset/nsd --space MNI_2mm --func betas_fithrf_GLMdenoise_RR --clip_model CLIP-ViT-H-14 --norm_nii --num_blocks 12 --patch_size 14 --patch_type conv --patch_drop 0 --attn_drop 0 --block_drop 0 --nii_mask brain --subj 1 --tag last \n"
     ]
    }
   ],
   "source": [
    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
    "if utils.is_interactive():\n",
    "    # global_batch_size and batch_size should already be defined in the 2nd cell block\n",
    "    jupyter_args = \"\"\n",
    "    jupyter_args += \"--name ddp2-0.01-qwen25vl-s1-CLIP-ViT-H-14-mask=brain-blocks=12-patch=conv_14-pt_dp=0-at_dp=0-bo_dp=0-txt_scl=1-m_pct=0-bs=128-lr=0.0003-bf=4096-norm_nii \"\n",
    "    jupyter_args += \"--checkpoints_dir ../checkpoints/ \"\n",
    "    jupyter_args += \"--seed 42 \"\n",
    "    jupyter_args += \"--batch_size 32 \"\n",
    "    jupyter_args += \"--subj_list 1 \"\n",
    "    jupyter_args += \"--nsddir /opt/data/private/dataset/nsd \"\n",
    "    jupyter_args += \"--space MNI_2mm \"\n",
    "    jupyter_args += \"--func betas_fithrf_GLMdenoise_RR \"\n",
    "    jupyter_args += \"--clip_model CLIP-ViT-H-14 \"\n",
    "    jupyter_args += \"--norm_nii \"\n",
    "    jupyter_args += \"--num_blocks 12 \"\n",
    "    jupyter_args += \"--patch_size 14 \"\n",
    "    jupyter_args += \"--patch_type conv \"\n",
    "    jupyter_args += \"--patch_drop 0 \"\n",
    "    jupyter_args += \"--attn_drop 0 \"\n",
    "    jupyter_args += \"--block_drop 0 \"\n",
    "    jupyter_args += \"--nii_mask brain \"\n",
    "    jupyter_args += \"--subj 1 \"\n",
    "    jupyter_args += \"--tag last \"\n",
    "\n",
    "    print(jupyter_args)\n",
    "    jupyter_args = jupyter_args.split()\n",
    "    \n",
    "    # from IPython.display import clear_output # function to clear print outputs in cell\n",
    "    # %load_ext autoreload \n",
    "    # # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n",
    "    # %autoreload 2 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0126effa",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description=\"test model\")\n",
    "parser.add_argument(\"--name\", type=str)\n",
    "parser.add_argument(\"--checkpoints_dir\", type=str, default=\"checkpoints/\")\n",
    "parser.add_argument(\"--seed\", type=int, default=42)\n",
    "parser.add_argument(\"--batch_size\", type=int)\n",
    "parser.add_argument(\"--subj_list\", nargs='+', type=int)\n",
    "parser.add_argument(\"--nsddir\", type=str)\n",
    "parser.add_argument(\"--space\", type=str)\n",
    "parser.add_argument(\"--func\", type=str)\n",
    "parser.add_argument(\"--clip_model\", type=str)\n",
    "parser.add_argument(\"--norm_nii\", action='store_true')\n",
    "parser.add_argument(\"--num_blocks\", type=int)\n",
    "parser.add_argument(\"--patch_size\", type=int)\n",
    "parser.add_argument(\"--patch_type\", type=str, default=\"conv\")\n",
    "parser.add_argument(\"--patch_drop\", type=float, default=0.0)\n",
    "parser.add_argument(\"--attn_drop\", type=float, default=0.0)\n",
    "parser.add_argument(\"--block_drop\", type=float, default=0.0)\n",
    "parser.add_argument(\"--nii_mask\", type=str, default=\"brain\")\n",
    "parser.add_argument(\"--subj\", type=int)\n",
    "parser.add_argument(\"--tag\", type=str)\n",
    "\n",
    "if utils.is_interactive():\n",
    "    args = parser.parse_args(jupyter_args)\n",
    "else:\n",
    "    args = parser.parse_args()\n",
    "\n",
    "utils.seed_everything(args.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "ae141ef5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cuda\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(\"device:\",device)\n",
    "num_devices = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2f378e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading CLIP model from  /opt/data/private/huggingface/models--laion--CLIP-ViT-H-14-laion2B-s32B-b79K\n"
     ]
    }
   ],
   "source": [
    "from typing import Any, Callable, Optional, Union\n",
    "from transformers.utils import auto_docstring\n",
    "from transformers import CLIPConfig, CLIPModel, CLIPProcessor, CLIPImageProcessor\n",
    "from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling\n",
    "class BrainCLIPModel(CLIPModel):\n",
    "    def __init__(self, config: CLIPConfig):\n",
    "        super().__init__(config)\n",
    "        self.brain_model = None\n",
    "        self.brain_projection = None\n",
    "    @auto_docstring\n",
    "    def get_brain_features(\n",
    "        self,\n",
    "        pixel_values: Optional[torch.FloatTensor] = None,\n",
    "        output_attentions: Optional[bool] = None,\n",
    "        output_hidden_states: Optional[bool] = None,\n",
    "        interpolate_pos_encoding: bool = False,\n",
    "    ) -> torch.FloatTensor:\n",
    "        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.\n",
    "        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n",
    "        output_hidden_states = (\n",
    "            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n",
    "        )\n",
    "\n",
    "        brain_outputs: BaseModelOutputWithPooling = self.brain_model(\n",
    "            pixel_values=pixel_values,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            interpolate_pos_encoding=interpolate_pos_encoding,\n",
    "        )\n",
    "\n",
    "        pooled_output = brain_outputs.pooler_output\n",
    "        brain_features = self.brain_projection(pooled_output)\n",
    "\n",
    "        return brain_features\n",
    "\n",
    "clip_model = {'CLIP-ViT-H-14': \"/opt/data/private/huggingface/models--laion--CLIP-ViT-H-14-laion2B-s32B-b79K\", 'clip-vit-l-14': \"/opt/data/private/huggingface/models--openai--clip-vit-large-patch14\", \"CLIP-ViT-bigG-14\": \"/opt/data/private/huggingface/models--laion--CLIP-ViT-bigG-14-laion2B-39B-b160k\"}\n",
    "\n",
    "print(\"Loading CLIP model from \", clip_model.get(args.clip_model, clip_model['CLIP-ViT-H-14']))\n",
    "processor = CLIPProcessor.from_pretrained(clip_model.get(args.clip_model, clip_model['CLIP-ViT-H-14']))\n",
    "model = BrainCLIPModel.from_pretrained(clip_model.get(args.clip_model, clip_model['CLIP-ViT-H-14'])).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46224cec",
   "metadata": {},
   "outputs": [],
   "source": [
    "nii_path = os.path.join(args.nsddir,'nsddata_betas','ppdata_split_pth','{:s}','{:s}','subj{:02d}_betas_session{:02d}', '{:03d}.pth')\n",
    "\n",
    "nsd_info = sio.loadmat(os.path.join(args.nsddir,'nsddata','experiments','nsd','nsd_expdesign.mat'))\n",
    "stim_info = pd.read_csv(os.path.join(args.nsddir,'nsddata','experiments','nsd', 'nsd_stim_info_merged.csv'), index_col=0)\n",
    "cocoId = stim_info['cocoId']\n",
    "stim_sort = (nsd_info['subjectim'][:,nsd_info['masterordering']-1]-1).squeeze()\n",
    "up_sess_bounds = [40, 40, 32, 30, 40, 32, 40, 30] \n",
    "with open(os.path.join(args.nsddir, 'nsddata_stimuli/stimuli/nsd/annotations', 'new_captions.json'), 'r') as f:\n",
    "# with open(os.path.join(args.nsddir, 'nsddata_stimuli/stimuli/nsd/annotations', 'qwen2_5_vl_captions_all.json'), 'r') as f:\n",
    "    captions = json.load(f)\n",
    "# ground truths\n",
    "all_images = torch.load(f\"../../Brain507/tools/all_images.pt\")\n",
    "stimulus = h5py.File(os.path.join(args.nsddir,'nsddata_stimuli','stimuli','nsd','coco_images_224_float16.hdf5'), 'r')['images']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "6c7154db",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_captions = None\n",
    "try:\n",
    "    tokenized_captions = torch.load('../tools/tokenized_newcaptions.pt')\n",
    "    # tokenized_captions = torch.load('../tools/tokenized_captions_qwen25vl.pt')\n",
    "except:\n",
    "    tokenized_captions = {}\n",
    "    for k, v in tqdm(captions.items()):\n",
    "        tokenized_captions[k] =processor.tokenizer(captions[k], padding=\"max_length\", truncation=True, max_length=77, return_tensors=\"pt\")\n",
    "    torch.save(tokenized_captions, '../tools/tokenized_newcaptions.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "aec70818",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subj 1 test set size: 1000\n"
     ]
    }
   ],
   "source": [
    "test_dict = {}\n",
    "if args.norm_nii:\n",
    "    norm_dict = {}\n",
    "for subj in args.subj_list:\n",
    "    test_list = []\n",
    "    stim_unique = pd.Series(stim_sort[subj-1]).unique()\n",
    "    for i in range(len(stim_unique)):\n",
    "        nii_loc = np.where(stim_sort[subj-1]==stim_unique[i])[0]\n",
    "        up_loc = np.where(nii_loc>=up_sess_bounds[subj-1]*750)[0]\n",
    "        if len(up_loc)!=0:\n",
    "            # print(f\"delete nii: {list(up_loc)} in stim: {stim_unique[i]:05d}\")\n",
    "            nii_loc = np.delete(nii_loc, up_loc)\n",
    "        if len(nii_loc)==0:\n",
    "            continue\n",
    "        if stim_unique[i] in nsd_info['sharedix']-1:\n",
    "            # test\n",
    "            test_list.append({\n",
    "                'subj': subj,\n",
    "                # 'stim': stim_path.format(stim_unique[i]),\n",
    "                'stim_idx': stim_unique[i],\n",
    "                'nii': [os.path.join(nii_path.format(args.space,args.func,subj,(nloc//750)+1,nloc%750)) for nloc in nii_loc],\n",
    "                'text': captions[str(cocoId[stim_unique[i]])],\n",
    "                'cocoid': cocoId[stim_unique[i]],\n",
    "            })\n",
    "    print(f\"subj {subj} test set size: {len(test_list)}\")\n",
    "    test_dict['subj{:02d}'.format(subj)] = test_list\n",
    "    if args.norm_nii:\n",
    "        norm_dict.update({'subj{:02d}_mean'.format(subj): torch.load(f'{args.nsddir}/nsddata_betas/mean_std/subj{subj}_{args.space}_{args.func}_mean.pth')})\n",
    "        norm_dict.update({'subj{:02d}_std'.format(subj): torch.load(f'{args.nsddir}/nsddata_betas/mean_std/subj{subj}_{args.space}_{args.func}_std.pth')})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "06e03ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_padding_list(patch_size):\n",
    "    base_size = [91, 109, 91]\n",
    "    padding_list = []\n",
    "    for dim in reversed(base_size):\n",
    "        pad_size = (patch_size - dim % patch_size) % patch_size\n",
    "        padding_list.extend([pad_size // 2, pad_size // 2 + (1 if pad_size % 2 else 0)])\n",
    "    return padding_list\n",
    "padding_list = cal_padding_list(args.patch_size)\n",
    "\n",
    "if args.nii_mask=='brain':\n",
    "    mask = torch.from_numpy(nib.load(os.path.join(\"{:s}\", \"fsl_tmp\", \"MNI152_T1_2mm_brain_mask.nii.gz\").format(args.nsddir)).get_fdata().astype(np.uint8))\n",
    "    mask = torch.nn.functional.pad(mask, pad=padding_list, mode='constant', value=0)\n",
    "    token_mask = torch.nn.functional.conv3d(mask[None,None].float(), torch.ones((1, 1, args.patch_size, args.patch_size, args.patch_size)).float(), stride=(args.patch_size,args.patch_size,args.patch_size)).byte().flatten(2).squeeze()\n",
    "    token_ids = torch.where(token_mask!=0)[0].tolist()\n",
    "elif args.nii_mask=='roi':\n",
    "    mask = torch.from_numpy(nib.load(\"{:s}/nsddata/ppdata/all/{:s}/roi/nsdgeneral.nii.gz\".format(args.nsddir, args.space)).get_fdata().astype(np.uint8))\n",
    "    mask = torch.nn.functional.pad(mask, pad=padding_list, mode='constant', value=0)\n",
    "    token_mask = torch.nn.functional.conv3d(mask[None,None].float(), torch.ones((1, 1, args.patch_size, args.patch_size, args.patch_size)).float(), stride=(args.patch_size,args.patch_size,args.patch_size)).byte().flatten(2).squeeze()\n",
    "    token_ids = torch.where(token_mask!=0)[0].tolist()\n",
    "else:\n",
    "    token_ids = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee59051",
   "metadata": {},
   "outputs": [],
   "source": [
    "PATCH_SIZE = args.patch_size  # Patch的大小\n",
    "NUM_BLOCKS = args.num_blocks \n",
    "PATCH_TYPE = args.patch_type \n",
    "ATTN_DROP = args.attn_drop\n",
    "BLOCK_DROP = args.block_drop\n",
    "PATCH_DROP = args.patch_drop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b3fb1ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import CLIPVisionConfig, CLIPVisionModelWithProjection\n",
    "from vit3d import ViTEmbeddings3D, ViTConfig, calc_image_size\n",
    "\n",
    "clip_cfg = CLIPVisionConfig(\n",
    "    attention_dropout=ATTN_DROP,\n",
    "    dropout=BLOCK_DROP,\n",
    "    hidden_act=model.config.vision_config.hidden_act,\n",
    "    hidden_size=model.config.vision_config.hidden_size,\n",
    "    # image_size=224,\n",
    "    initializer_factor=1.0,\n",
    "    initializer_range=0.02,\n",
    "    intermediate_size=model.config.vision_config.intermediate_size,\n",
    "    layer_norm_eps=model.config.vision_config.layer_norm_eps,\n",
    "    model_type=\"clip_vision_model\",\n",
    "    num_attention_heads=model.config.vision_config.num_attention_heads,\n",
    "    # num_channels=1,\n",
    "    num_hidden_layers=NUM_BLOCKS,\n",
    "    projection_dim=model.config.vision_config.projection_dim,\n",
    "    patch_size=14,\n",
    ")\n",
    "\n",
    "model_args = deepcopy(clip_cfg)\n",
    "model_args.image_size = calc_image_size(PATCH_SIZE)\n",
    "model_args.patch_size = (PATCH_SIZE, PATCH_SIZE, PATCH_SIZE)\n",
    "model_args.hidden_dropout_prob = PATCH_DROP\n",
    "model_args.num_channels=1\n",
    "model_args.token_ids = token_ids\n",
    "model_args.patch_type = PATCH_TYPE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "c71c523a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "167\n",
      "tensor([ 23,  24,  30,  31,  64,  65,  66,  67,  70,  71,  72,  73,  74,  75,\n",
      "         77,  78,  79,  80,  81,  82,  85,  86,  87,  88,  89,  92,  93,  94,\n",
      "         95, 100, 101, 102, 114, 115, 119, 120, 121, 122, 123, 124, 126, 127,\n",
      "        128, 129, 130, 131, 133, 134, 135, 136, 137, 138, 141, 142, 143, 144,\n",
      "        145, 148, 149, 150, 151, 152, 155, 156, 157, 158, 163, 164, 170, 171,\n",
      "        175, 176, 177, 178, 179, 180, 182, 183, 184, 185, 186, 187, 189, 190,\n",
      "        191, 192, 193, 194, 197, 198, 199, 200, 201, 204, 205, 206, 207, 208,\n",
      "        211, 212, 213, 214, 219, 220, 226, 227, 231, 232, 233, 234, 235, 236,\n",
      "        238, 239, 240, 241, 242, 243, 245, 246, 247, 248, 249, 250, 253, 254,\n",
      "        255, 256, 257, 260, 261, 262, 263, 264, 267, 268, 269, 270, 275, 276,\n",
      "        288, 289, 290, 291, 294, 295, 296, 297, 298, 299, 301, 302, 303, 304,\n",
      "        305, 306, 309, 310, 311, 312, 313, 316, 317, 318, 319, 324, 325])\n",
      "Conv3d(1, 1280, kernel_size=(14, 14, 14), stride=(14, 14, 14))\n",
      "param counts:\n",
      "241,175,040 total\n",
      "241,175,040 trainable\n"
     ]
    }
   ],
   "source": [
    "clip_fmri = CLIPVisionModelWithProjection(clip_cfg)\n",
    "clip_fmri.vision_model.embeddings = ViTEmbeddings3D(model_args)\n",
    "utils.count_params(clip_fmri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "20528e86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "---loading ../checkpoints/ddp2-0.01-qwen25vl-s1-CLIP-ViT-H-14-mask=brain-blocks=12-patch=conv_14-pt_dp=0-at_dp=0-bo_dp=0-txt_scl=1-m_pct=0-bs=128-lr=0.0003-bf=4096-norm_nii/last.pth ckpt---\n",
      "\n",
      "Epoch 149\n"
     ]
    }
   ],
   "source": [
    "# ================= load ckpt ================ #\n",
    "checkpoint_dir = os.path.join(args.checkpoints_dir, args.name)\n",
    "print(f\"\\n---loading {checkpoint_dir}/{args.tag}.pth ckpt---\\n\")\n",
    "\n",
    "state_dict = torch.load(checkpoint_dir+f'/{args.tag}.pth', map_location='cpu')\n",
    "clip_fmri.load_state_dict(state_dict, strict=True)\n",
    "json_save_path = os.path.join(checkpoint_dir, f'{args.tag}.json')\n",
    "with open(json_save_path, 'r') as json_file:\n",
    "    state = json.load(json_file)\n",
    "    epoch_idx = state['epoch_idx']\n",
    "print(\"Epoch\", epoch_idx)\n",
    "clip_fmri = clip_fmri.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a1679826",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.brain_model = clip_fmri.vision_model\n",
    "model.brain_projection = clip_fmri.visual_projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "1fa6064e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_norm_nii(nii, mean, std):\n",
    "    mean = mean.to(nii.device)\n",
    "    std = std.to(nii.device)\n",
    "    nii = (nii - mean) / (std+1e-5)\n",
    "    return nii"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b76217",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(idx):\n",
    "    batch = {}\n",
    "    data = test_dict['subj01'][idx]\n",
    "\n",
    "    text_ids = processor.tokenizer(captions[str(data['cocoid'])], return_tensors=\"pt\")\n",
    "\n",
    "    nii = [torch.load(n) for n in data['nii']]\n",
    "    if len(nii)==1:\n",
    "        nii = torch.stack(nii).repeat(3,1,1,1)\n",
    "    elif len(nii)==2:\n",
    "        nii = torch.stack([nii[0], nii[1], nii[0]])\n",
    "    else:\n",
    "        nii = torch.stack(nii)\n",
    "\n",
    "    stim = torch.from_numpy(stimulus[data['stim_idx']])\n",
    "    if processor is not None:\n",
    "        stim = processor.image_processor(stim, return_tensors=\"pt\", do_rescale=False).pixel_values\n",
    "    \n",
    "    nii = nii.float()/300.\n",
    "    if args.norm_nii:\n",
    "        nii = cal_norm_nii(nii, norm_dict[f'subj{subj:02d}_mean'].unsqueeze(0), norm_dict[f'subj{subj:02d}_std'].unsqueeze(0))\n",
    "    nii = torch.nn.functional.pad(nii, pad=padding_list, mode='constant', value=0)\n",
    "\n",
    "    for k,v in text_ids.items():\n",
    "        text_ids[k] = v.to(device)\n",
    "    batch.update({'text': text_ids})\n",
    "    batch.update({'vision': stim.to(device)})\n",
    "    batch.update({'brain': nii.mean(0, keepdim=True).unsqueeze(0).to(device)})\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "4a2e646f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 19444 characters]",
      "image/png": "[encoded data, 155328 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=224x224>"
      ]
     },
     "execution_count": 142,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imgidx = 13\n",
    "batch = get_batch(imgidx)\n",
    "utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "81760109",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_vector_norm(tensor: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    This method is equivalent to tensor.norm(p=2, dim=-1, keepdim=True) and used to make\n",
    "    model `executorch` exportable. See issue https://github.com/pytorch/executorch/issues/3566\n",
    "    \"\"\"\n",
    "    square_tensor = torch.pow(tensor, 2)\n",
    "    sum_tensor = torch.sum(square_tensor, dim=-1, keepdim=True)\n",
    "    normed_tensor = torch.pow(sum_tensor, 0.5)\n",
    "    return normed_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "e547cfcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.text import Text\n",
    "def generate_shades_with_alpha(scores, min_alpha=0, max_alpha=0.8):\n",
    "    normalized_scores = [(score - min(scores)) / (max(scores) - min(scores)) for score in scores]\n",
    "    red_colors = []\n",
    "    for score in normalized_scores:\n",
    "        alpha = min_alpha + score * (max_alpha - min_alpha)\n",
    "        red_color = (0.2, 0.4, 1, alpha)\n",
    "        red_colors.append(red_color)\n",
    "    return red_colors\n",
    "class TextWithBGColor(Text):\n",
    "    def __init__(self, x, y, text, bgcolor, *args, **kwargs):\n",
    "        super().__init__(x, y, text, *args, **kwargs)\n",
    "        self.bgcolor = bgcolor\n",
    "\n",
    "    def draw(self, renderer, *args, **kwargs):\n",
    "        bbox = dict(facecolor=self.bgcolor, edgecolor=self.bgcolor, boxstyle='round,pad=0.01', alpha=self.bgcolor[3])\n",
    "        self.set_bbox(bbox)\n",
    "        super().draw(renderer, *args, **kwargs)\n",
    "\n",
    "def plot_text_with_colors(ax, tokens, rgba_colors, max_width, max_height, fontsize=12):\n",
    "    x, y = 0, max_height\n",
    "    space_width = fontsize * 0.2  # Width of a single space\n",
    "    for i, (token, color) in enumerate(zip(tokens, rgba_colors)):\n",
    "        text = TextWithBGColor(x, y*0.7, token, color, fontsize=fontsize)\n",
    "        ax.add_artist(text)\n",
    "        \n",
    "        token_width = fontsize * 0.4 * len(token)\n",
    "\n",
    "        if i + 1 < len(tokens):\n",
    "            next_token_width = fontsize * 0.4 * len(tokens[i + 1]) + space_width\n",
    "            if x + token_width + next_token_width >= max_width:\n",
    "                x = 0\n",
    "                y -= 1  # Decrease the space between rows\n",
    "            else:\n",
    "                x += token_width + space_width  # Add the width of a single space after each token\n",
    "        else:\n",
    "            x += token_width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96dbffce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from M2IB.scripts.utils import normalize \n",
    "# ==== 基础参数 ====\n",
    "volume_shape = calc_image_size(PATCH_SIZE)    \n",
    "patch_size = (PATCH_SIZE, PATCH_SIZE, PATCH_SIZE)   \n",
    "grid_shape = tuple(d // p for d, p in zip(volume_shape, patch_size))  # (7, 8, 7)\n",
    "num_tokens = np.prod(grid_shape)  # 392\n",
    "\n",
    "def process_cams(cam):\n",
    "    # cam = torch.nn.functional.relu(cam)\n",
    "    # cam = (cam / cam.norm(dim=-1, keepdim=True))\n",
    "    # cam = (cam-cam.mean(axis=-1, keepdims=True))/cam.std(axis=-1, keepdims=True)\n",
    "    cam=cam/(cam.abs().max(dim=-1, keepdims=True).values)\n",
    "    # cam = (cam-cam.mean())/cam.std()\n",
    "    # cam=cam/(cam.abs().max())\n",
    "    # (cam - cam.min(dim=-1, keepdims=True).values) / (cam.max(dim=-1, keepdims=True).values - cam.min(dim=-1, keepdims=True).values)\n",
    "    return cam\n",
    "def heatmap(saliency, target_model):\n",
    "    if target_model=='text':\n",
    "        # saliency = torch.nansum(saliency, -1).cpu().detach().numpy()\n",
    "        saliency = normalize(saliency)\n",
    "        return saliency\n",
    "    elif target_model=='vision':\n",
    "        # saliency = torch.nansum(saliency, -1)[1:] # Discard the first because it's the CLS token\n",
    "        dim = int(saliency.numel() ** 0.5)\n",
    "        saliency = saliency.reshape(1, 1, dim, dim)\n",
    "        saliency = torch.nn.functional.interpolate(saliency, size=224, mode='bilinear')\n",
    "        saliency = saliency.squeeze().cpu().detach().numpy()\n",
    "        return normalize(saliency)\n",
    "    elif target_model=='brain':\n",
    "        global token_ids\n",
    "        # saliency = torch.nansum(saliency, -1)[1:].cpu().detach() # Discard the first because it's the CLS token\n",
    "        saliency = process_cams(saliency)\n",
    "        # saliency = (saliency / saliency.norm(dim=-1, keepdim=True))\n",
    "        token_ids = torch.tensor(token_ids)\n",
    "        saliency_grid = torch.full((num_tokens,), float(0)) \n",
    "        saliency_grid[token_ids] = saliency\n",
    "        # cam_grid=(cam_grid+1)/2\n",
    "        saliency_volume = saliency_grid.reshape(grid_shape)    # shape: [7, 8, 7]\n",
    "        saliency_volume = saliency_volume.unsqueeze(0).unsqueeze(0)  # [1,1,7,8,7]\n",
    "        saliency_volume = torch.nn.functional.interpolate(\n",
    "            saliency_volume, size=volume_shape, mode=\"trilinear\", align_corners=False\n",
    "        )[0, 0]  # shape: [98, 112, 98]\n",
    "        saliency_volume *= mask\n",
    "        saliency_volume = normalize(saliency_volume)\n",
    "        saliency = saliency_volume[padding_list[0]:-padding_list[1], padding_list[2]:-padding_list[3], padding_list[4]:-padding_list[5]]\n",
    "        saliency=saliency.permute(2,1,0)\n",
    "        return saliency, saliency_volume"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f9418b7",
   "metadata": {},
   "source": [
    "## rise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "003b8b3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.transform import resize\n",
    "import copy\n",
    "from torch import nn\n",
    "class RISEV(nn.Module):\n",
    "    def __init__(self, model, input_size, gpu_batch=100):\n",
    "        super(RISEV, self).__init__()\n",
    "        self.model = model\n",
    "        self.input_size = input_size\n",
    "        self.gpu_batch = gpu_batch\n",
    "        self.loss_fn = nn.CosineSimilarity(eps=1e-6)\n",
    "\n",
    "    def generate_masks(self, N, s, p1, savepath='masks.npy'):\n",
    "        cell_size = np.ceil(np.array(self.input_size) / s)\n",
    "        up_size = (s + 1) * cell_size\n",
    "\n",
    "        grid = np.random.rand(N, s, s) < p1\n",
    "        grid = grid.astype('float32')\n",
    "\n",
    "        self.masks = np.empty((N, *self.input_size))\n",
    "\n",
    "        for i in tqdm(range(N), desc='Generating filters'):\n",
    "            # Random shifts\n",
    "            x = np.random.randint(0, cell_size[0])\n",
    "            y = np.random.randint(0, cell_size[1])\n",
    "            # Linear upsampling and cropping\n",
    "            self.masks[i, :, :] = resize(grid[i], up_size, order=1, mode='reflect',\n",
    "                                         anti_aliasing=False)[x:x + self.input_size[0], y:y + self.input_size[1]]\n",
    "        self.masks = self.masks.reshape(-1, 1, *self.input_size)\n",
    "        np.save(savepath, self.masks)\n",
    "        self.masks = torch.from_numpy(self.masks).half()\n",
    "        self.masks = self.masks.cuda()\n",
    "        self.N = N\n",
    "        self.p1 = p1\n",
    "\n",
    "    def load_masks(self, filepath, p1=0.1):\n",
    "        self.masks = np.load(filepath)\n",
    "        self.masks = torch.from_numpy(self.masks).half().cuda()\n",
    "        self.N = self.masks.shape[0]\n",
    "        self.p1 = p1\n",
    "        \n",
    "\n",
    "    def forward(self, image, text_features):\n",
    "        N = self.N\n",
    "        _, _, H, W = image.size()\n",
    "        stack = torch.mul(self.masks, image.data)\n",
    "        p = []\n",
    "        for i in range(0, N, self.gpu_batch):\n",
    "            image_features = self.model.get_image_features(stack[i:min(i + self.gpu_batch, N)])\n",
    "            p.append(self.loss_fn(image_features, text_features).detach().cpu().unsqueeze(-1))\n",
    "        p = torch.cat(p)\n",
    "        sal = torch.matmul(p.data.transpose(0, 1).half(), self.masks.cpu().view(N, H * W))\n",
    "        sal = sal.view((1, H, W))\n",
    "        sal = sal / N / self.p1\n",
    "        sal = sal.mean(0, keepdim=True)\n",
    "        sal = (sal - sal.min()) / (sal.max() - sal.min())\n",
    "        return sal\n",
    "    \n",
    "def rise_v(model,image, brain_features, N=6000, s=14, p1=0.1, savepath='masks/masks_2d_{:d}.npy'):\n",
    "    exp = RISEV(model, (224, 224), gpu_batch=20)\n",
    "    if os.path.exists(savepath.format(s)):\n",
    "        exp.load_masks(savepath.format(s))\n",
    "    else:\n",
    "        exp.generate_masks(N, s, p1, savepath=savepath.format(s))\n",
    "    sal = exp(image, brain_features)\n",
    "    return sal\n",
    "\n",
    "\n",
    "\n",
    "class RISET(nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super(RISET, self).__init__()\n",
    "        self.ori_model = copy.deepcopy(model)\n",
    "        self.model = copy.deepcopy(model)\n",
    "        self.loss_fn = nn.CosineSimilarity(eps=1e-6)\n",
    "\n",
    "    def generate_masks(self,input_size, N, s, p1):\n",
    "        cell_size = np.ceil(np.array(input_size) / s)\n",
    "        up_size = (s + 1) * cell_size\n",
    "        \n",
    "        np.random.seed(0)\n",
    "\n",
    "        grid = np.random.rand(N, s) < p1\n",
    "        grid = grid.astype('float32')\n",
    "\n",
    "        masks = np.empty((N, *input_size))\n",
    "\n",
    "        for i in range(N):\n",
    "            x = np.random.randint(0, cell_size[0])\n",
    "            y = np.random.randint(0, cell_size[1])\n",
    "            masks[i, :, :] = resize(grid[i], up_size, order=1, mode='reflect',\n",
    "                                         anti_aliasing=False)[x:x + input_size[0], y:y + input_size[1]]\n",
    "        masks = masks.reshape(-1, 1, *input_size).squeeze()\n",
    "        masks = torch.from_numpy(masks).float()\n",
    "        masks = masks.cuda()\n",
    "        N = N\n",
    "        p1 = p1\n",
    "        return masks, N, p1\n",
    "\n",
    "\n",
    "    def forward(self, tid, brain_features, N=200, s=8, p1=0.1):\n",
    "        masks, N, p1 = self.generate_masks((tid.shape[-1],brain_features.shape[-1]), N, s, p1)\n",
    "        sal = list()\n",
    "        for i in range(0, N):\n",
    "            self.model.text_model.embeddings.token_embedding.weight.data[tid,:] = self.ori_model.text_model.embeddings.token_embedding.weight.data[tid,:] * masks[i]\n",
    "            text_features = self.model.get_text_features(tid)\n",
    "            sal.append(self.loss_fn(brain_features, text_features).detach().cpu().unsqueeze(-1) * masks[i].cpu())\n",
    "        sal = torch.stack(sal)\n",
    "        sal = sal.mean(0).sum(-1) / p1\n",
    "        sal_min = sal.min()\n",
    "        sal_max = sal.max()\n",
    "        sal = (sal - sal_min) / (sal_max - sal_min)\n",
    "        return sal\n",
    "\n",
    "\n",
    "\n",
    "def rise_t(model,tid,brain_features):\n",
    "    exp = RISET(model)\n",
    "    sal = exp(tid, brain_features)\n",
    "    return sal\n",
    "\n",
    "\n",
    "class RISEB(nn.Module):\n",
    "    def __init__(self, model, input_size, gpu_batch=100):\n",
    "        super(RISEB, self).__init__()\n",
    "        self.model = model\n",
    "        self.input_size = input_size\n",
    "        self.gpu_batch = gpu_batch\n",
    "        self.loss_fn = nn.CosineSimilarity(eps=1e-6)\n",
    "\n",
    "    def generate_masks(self, N, s, p1, savepath='masks.npy'):\n",
    "        cell_size = np.ceil(np.array(self.input_size) / s)\n",
    "        up_size = (s + 1) * cell_size\n",
    "\n",
    "        grid = np.random.rand(N, s, s, s) < p1\n",
    "        grid = grid.astype('float16')\n",
    "\n",
    "        self.masks = np.empty((N, *self.input_size), dtype='float16')\n",
    "\n",
    "        for i in tqdm(range(N), desc='Generating filters'):\n",
    "            # Random shifts\n",
    "            x = np.random.randint(0, cell_size[0])\n",
    "            y = np.random.randint(0, cell_size[1])\n",
    "            z = np.random.randint(0, cell_size[2])\n",
    "            # Linear upsampling and cropping\n",
    "            self.masks[i, :, :, :] = resize(grid[i], up_size, order=1, mode='reflect',\n",
    "                                         anti_aliasing=False)[x:x + self.input_size[0], y:y + self.input_size[1], z:z + self.input_size[2]]\n",
    "        self.masks = self.masks.reshape(-1, 1, *self.input_size).astype('float16')\n",
    "        np.save(savepath, self.masks)\n",
    "        self.masks = torch.from_numpy(self.masks).half()\n",
    "        self.masks = self.masks.cuda()\n",
    "        self.N = N\n",
    "        self.p1 = p1\n",
    "\n",
    "    def load_masks(self, filepath, p1=0.1):\n",
    "        self.masks = np.load(filepath)\n",
    "        self.masks = torch.from_numpy(self.masks).half().cuda()\n",
    "        self.N = self.masks.shape[0]\n",
    "        self.p1 = p1\n",
    "        \n",
    "\n",
    "    def forward(self, brain, text_features, image_features):\n",
    "        N = self.N\n",
    "        _, _, H, W, D = brain.size()\n",
    "        stack = torch.mul(self.masks, brain.data)\n",
    "        p = []\n",
    "        for i in range(0, N, self.gpu_batch):\n",
    "            brain_features = self.model.get_brain_features(stack[i:min(i + self.gpu_batch, N)])\n",
    "            p.append((self.loss_fn(brain_features, text_features)/2+self.loss_fn(brain_features, image_features)/2).unsqueeze(-1))\n",
    "        p = torch.cat(p)\n",
    "        sal = torch.matmul(p.data.transpose(0, 1).half(), self.masks.view(N, H * W * D)).cpu()\n",
    "        sal = sal.view((1, H, W, D))\n",
    "        sal = sal / N / self.p1\n",
    "        sal = sal.mean(0, keepdim=True)\n",
    "        # sal = (sal - sal.min()) / (sal.max() - sal.min())\n",
    "        return sal\n",
    "    \n",
    "def rise_b(model, brain, text_features, image_features, N=1000, s=14, p1=0.1, savepath='masks/masks_3d_{:d}.npy'):\n",
    "    \n",
    "    exp = RISEB(model, batch['brain'].shape[-3:], gpu_batch=10)\n",
    "    if os.path.exists(savepath.format(s)):\n",
    "        exp.load_masks(savepath.format(s))\n",
    "    else:\n",
    "        exp.generate_masks(N, s, p1, savepath=savepath.format(s))\n",
    "    sal = exp(brain, text_features, image_features)\n",
    "    return sal\n",
    "\n",
    "\n",
    "def rise(model, batch, target_model, s=14):\n",
    "    model.eval()\n",
    "    out = {}\n",
    "    with torch.no_grad(), torch.cuda.amp.autocast():\n",
    "        for k in batch.keys():\n",
    "            if k=='text':\n",
    "                text_embeds = model.get_text_features(**batch[k])\n",
    "                text_embeds = text_embeds / get_vector_norm(text_embeds)\n",
    "                out.update({'text': text_embeds})\n",
    "            elif k=='vision':\n",
    "                vision_embeds = model.get_image_features(batch[k])\n",
    "                vision_embeds = vision_embeds / get_vector_norm(vision_embeds)\n",
    "                out.update({'vision': vision_embeds})\n",
    "            elif k=='brain':\n",
    "                brain_embeds = model.get_brain_features(batch[k])\n",
    "                brain_embeds = brain_embeds / get_vector_norm(brain_embeds)\n",
    "                out.update({'brain': brain_embeds})\n",
    "\n",
    "        if target_model=='text':\n",
    "            sal = rise_t(model, batch['text']['input_ids'][0][batch['text']['attention_mask'][0]==1].unsqueeze(0), out['brain'])\n",
    "        elif target_model=='vision':\n",
    "            sal = rise_v(model, batch['vision'], out['brain'], s=s)\n",
    "        elif target_model=='brain':\n",
    "            sal = rise_b(model, batch['brain'], out['text'], out['vision'], s=s)\n",
    "    return sal.cpu().detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6b6c7c81",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_model = 'text'\n",
    "saliency = rise(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5b3a1c21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84104 characters]",
      "text/plain": [
       "<Figure size 1000x200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target_model = 'text'\n",
    "text_words = processor.tokenizer.convert_ids_to_tokens(batch[target_model]['input_ids'][0][batch[target_model]['attention_mask'][0]==1].tolist())\n",
    "text_words = [x.split('<')[0] for x in text_words[1:-1]]\n",
    "tmap = heatmap(saliency.numpy(), target_model)\n",
    "rgba_colors = generate_shades_with_alpha(tmap[1:-1])\n",
    "\n",
    "fig, axs = plt.subplots(1, 2)\n",
    "fig.set_size_inches(10,2) \n",
    "\n",
    "axs[0].imshow(utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']]))\n",
    "# if bb:\n",
    "#     for x, y, w, h in bb:\n",
    "#         rect = mpl.patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='r', facecolor='none')\n",
    "#         axs[0].add_patch(rect)\n",
    "axs[0].axis('off')\n",
    "\n",
    "plot_text_with_colors(axs[1], text_words, rgba_colors, max_width=500, max_height=2, fontsize=12)\n",
    "axs[1].set_xlim(0, 250)\n",
    "axs[1].set_ylim(-2.5, 2.5)\n",
    "axs[1].axis('off')\n",
    "plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "id": "9807bc8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_model = 'vision'\n",
    "saliency = rise(model, batch, target_model=target_model, s=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "id": "ad1d1bdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 736100 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']]).resize((224, 224))\n",
    "vmap = heatmap(saliency, target_model)\n",
    "vmap = plt.cm.jet(vmap)[..., :3] * 255\n",
    "vmap = Image.fromarray(vmap.astype(np.uint8))\n",
    "overlay = Image.blend(img, vmap, alpha=0.4)\n",
    "\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.subplot(1, 3, 1); plt.imshow(img); plt.title(\"Original Image\"); plt.axis(\"off\")\n",
    "plt.subplot(1, 3, 2); plt.imshow(vmap, cmap='jet'); plt.title(\"Grad-CAM Map\"); plt.axis(\"off\")\n",
    "plt.subplot(1, 3, 3); plt.imshow(overlay); plt.title(\"Overlay\"); plt.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "58392079",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = get_batch(0)\n",
    "target_model = 'brain'\n",
    "saliency = rise(model, batch, target_model=target_model, s=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d84af203",
   "metadata": {},
   "outputs": [],
   "source": [
    "bmap_2d = normalize(bmap_2d)\n",
    "bmap_2d *= mask\n",
    "bmap_2d = normalize(bmap_2d)\n",
    "bmap = bmap_2d[padding_list[0]:-padding_list[1], padding_list[2]:-padding_list[3], padding_list[4]:-padding_list[5]]\n",
    "bmap=bmap.permute(2,1,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "03eed721",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32796 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ==== 可视化某一切片 ====\n",
    "plt.imshow(bmap_2d[:, :, bmap_2d.shape[2] // 2].cpu(), cmap='jet', vmax=1, vmin=0)\n",
    "plt.title(\"Grad-CAM Heatmap (z-mid)\")\n",
    "plt.axis(\"off\")\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "ac661850",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed to get connection\n",
      "** (inkscape:152361): CRITICAL **: 15:57:04.515: dbus_g_proxy_new_for_name: assertion 'connection != NULL' failed\n",
      "** (inkscape:152361): CRITICAL **: 15:57:04.515: dbus_g_proxy_call: assertion 'DBUS_IS_G_PROXY (proxy)' failed\n",
      "** (inkscape:152361): CRITICAL **: 15:57:04.515: dbus_g_connection_register_g_object: assertion 'connection != NULL' failed\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 2086972 characters]",
      "text/plain": [
       "<Figure size 2274x1024 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# import cortex\n",
    "subject = 'fsaverage'\n",
    "xfm = 'atlas_2mm'\n",
    "\n",
    "# 将数据包装成 pycortex volume（带掩膜）\n",
    "voxel_vol = cortex.Volume(bmap.numpy(), subject=subject, xfmname=xfm, vmax=1, vmin=0,cmap='RdBu_r')#,cmap='RGrB_tsi'CyanBlueGrayRedPink,BuBkRd\n",
    "# cortex.webshow(voxel_vol, cmap='hot', colorbar=True, with_rois=True)\n",
    "fig=cortex.quickshow(voxel_vol, with_rois=True, with_colorbar=True, colorbar=True) # , cmap='Blues'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d12aa277",
   "metadata": {},
   "source": [
    "## Chefer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "40df3e05",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chefer(model, batch, target_model='brain'):\n",
    "    model.eval()\n",
    "    out = {}\n",
    "    for k in batch.keys():\n",
    "        if k=='text':\n",
    "            text_outputs = model.text_model(**batch[k], output_attentions=target_model=='text')\n",
    "            text_embeds = model.text_projection(text_outputs.pooler_output)\n",
    "            text_embeds = text_embeds / get_vector_norm(text_embeds)\n",
    "            out.update({'text': {'outputs': text_outputs, 'embeds': text_embeds}})\n",
    "        elif k=='vision':\n",
    "            vision_outputs = model.vision_model(batch[k], output_attentions=target_model=='vision')\n",
    "            vision_embeds = model.visual_projection(vision_outputs.pooler_output)\n",
    "            vision_embeds = vision_embeds / get_vector_norm(vision_embeds)\n",
    "            out.update({'vision': {'outputs': vision_outputs, 'embeds': vision_embeds}})\n",
    "        elif k=='brain':\n",
    "            brain_outputs = model.brain_model(batch[k], output_attentions=target_model=='brain')\n",
    "            brain_embeds = model.brain_projection(brain_outputs.pooler_output)\n",
    "            brain_embeds = brain_embeds / get_vector_norm(brain_embeds)\n",
    "            out.update({'brain': {'outputs': brain_outputs, 'embeds': brain_embeds}})\n",
    "\n",
    "    model.zero_grad()\n",
    "    if target_model=='text':\n",
    "        # logits_per_text_v = torch.matmul(out['text']['embeds'], out['vision']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "        logits_per_text_b = torch.matmul(out['text']['embeds'], out['brain']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "        # logits_per_text = (logits_per_text_v + logits_per_text_b) / 2\n",
    "        logits_per_text = logits_per_text_b\n",
    "        logit = logits_per_text[0,0]\n",
    "        A = out['text']['outputs'].attentions[-1]\n",
    "        grad = torch.autograd.grad(logit, A)[0].detach()\n",
    "        grad = (grad[0, :, -1] * A.detach()[0, :, -1]).clamp(min=0).mean(dim=0)\n",
    "    elif target_model=='vision':\n",
    "        # logits_per_image_t = torch.matmul(out['vision']['embeds'], out['text']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "        logits_per_image_b = torch.matmul(out['vision']['embeds'], out['brain']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "        # logits_per_image = (logits_per_image_t + logits_per_image_b) / 2\n",
    "        logits_per_image = logits_per_image_b\n",
    "        logit = logits_per_image[0,0]\n",
    "        A = out['vision']['outputs'].attentions[-1]\n",
    "        grad = torch.autograd.grad(logit, A)[0].detach().abs()\n",
    "        grad = (grad[0, :, 0, 1:] * A.detach()[0, :, 0, 1:]).clamp(min=0).mean(dim=0)\n",
    "    elif target_model=='brain':\n",
    "        logits_per_brain_t = torch.matmul(out['brain']['embeds'], out['text']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "        logits_per_brain_v = torch.matmul(out['brain']['embeds'], out['vision']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "        logits_per_brain = (logits_per_brain_t + logits_per_brain_v) / 2\n",
    "        logit = logits_per_brain[0,0]\n",
    "        A = out['brain']['outputs'].attentions[-1]\n",
    "        grad = torch.autograd.grad(logit, A)[0].detach().abs()\n",
    "        grad = (grad[0, :, 0, 1:] * A.detach()[0, :, 0, 1:]).clamp(min=0).mean(dim=0)\n",
    "\n",
    "    return grad.cpu().detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "id": "711065e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1363: UserWarning: Using a non-full backward hook when outputs are generated by different autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_output. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using a non-full backward hook when outputs are generated by different autograd Nodes \"\n",
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1348: UserWarning: Using non-full backward hooks on a Module that does not take as input a single Tensor or a tuple of Tensors is deprecated and will be removed in future versions. This hook will be missing some of the grad_input. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using non-full backward hooks on a Module that does not take as input a \"\n"
     ]
    }
   ],
   "source": [
    "target_model = 'text'\n",
    "saliency = chefer(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "id": "584379dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 83532 characters]",
      "text/plain": [
       "<Figure size 1000x200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target_model = 'text'\n",
    "text_words = processor.tokenizer.convert_ids_to_tokens(batch[target_model]['input_ids'][0][batch[target_model]['attention_mask'][0]==1].tolist())\n",
    "text_words = [x.split('<')[0] for x in text_words[1:-1]]\n",
    "tmap = heatmap(saliency.numpy(), target_model)\n",
    "rgba_colors = generate_shades_with_alpha(tmap[1:-1])\n",
    "\n",
    "fig, axs = plt.subplots(1, 2)\n",
    "fig.set_size_inches(10,2) \n",
    "\n",
    "axs[0].imshow(utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']]))\n",
    "# if bb:\n",
    "#     for x, y, w, h in bb:\n",
    "#         rect = mpl.patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='r', facecolor='none')\n",
    "#         axs[0].add_patch(rect)\n",
    "axs[0].axis('off')\n",
    "\n",
    "plot_text_with_colors(axs[1], text_words, rgba_colors, max_width=500, max_height=2, fontsize=12)\n",
    "axs[1].set_xlim(0, 250)\n",
    "axs[1].set_ylim(-2.5, 2.5)\n",
    "axs[1].axis('off')\n",
    "plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "347a9f74",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1373: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using a non-full backward hook when the forward contains multiple autograd Nodes \"\n"
     ]
    }
   ],
   "source": [
    "target_model = 'vision'\n",
    "saliency = chefer(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "2e9306fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 777372 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']]).resize((224, 224))\n",
    "vmap = heatmap(saliency, target_model)\n",
    "vmap = plt.cm.jet(vmap)[..., :3] * 255\n",
    "vmap = Image.fromarray(vmap.astype(np.uint8))\n",
    "overlay = Image.blend(img, vmap, alpha=0.4)\n",
    "\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.subplot(1, 3, 1); plt.imshow(img); plt.title(\"Original Image\"); plt.axis(\"off\")\n",
    "plt.subplot(1, 3, 2); plt.imshow(vmap, cmap='jet'); plt.title(\"Grad-CAM Map\"); plt.axis(\"off\")\n",
    "plt.subplot(1, 3, 3); plt.imshow(overlay); plt.title(\"Overlay\"); plt.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "84c70747",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1348: UserWarning: Using non-full backward hooks on a Module that does not take as input a single Tensor or a tuple of Tensors is deprecated and will be removed in future versions. This hook will be missing some of the grad_input. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using non-full backward hooks on a Module that does not take as input a \"\n"
     ]
    }
   ],
   "source": [
    "batch = get_batch(13)\n",
    "target_model = 'brain'\n",
    "saliency = chefer(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "bd417842",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_149542/3096259938.py:34: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  token_ids = torch.tensor(token_ids)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22664 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bmap, bmap_2d = heatmap(saliency, target_model)\n",
    "# ==== 可视化某一切片 ====\n",
    "plt.imshow(bmap_2d[:, :, bmap_2d.shape[2] // 2].cpu(), cmap='jet', vmax=1, vmin=0)\n",
    "plt.title(\"Grad-CAM Heatmap (z-mid)\")\n",
    "plt.axis(\"off\")\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "4de366df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed to get connection\n",
      "** (inkscape:154783): CRITICAL **: 16:16:05.409: dbus_g_proxy_new_for_name: assertion 'connection != NULL' failed\n",
      "** (inkscape:154783): CRITICAL **: 16:16:05.409: dbus_g_proxy_call: assertion 'DBUS_IS_G_PROXY (proxy)' failed\n",
      "** (inkscape:154783): CRITICAL **: 16:16:05.409: dbus_g_connection_register_g_object: assertion 'connection != NULL' failed\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1307232 characters]",
      "text/plain": [
       "<Figure size 2274x1024 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# import cortex\n",
    "subject = 'fsaverage'\n",
    "xfm = 'atlas_2mm'\n",
    "\n",
    "# 将数据包装成 pycortex volume（带掩膜）\n",
    "voxel_vol = cortex.Volume(bmap.numpy(), subject=subject, xfmname=xfm, vmax=1, vmin=0,cmap='RdBu_r')#,cmap='RGrB_tsi'CyanBlueGrayRedPink,BuBkRd\n",
    "# cortex.webshow(voxel_vol, cmap='hot', colorbar=True, with_rois=True)\n",
    "fig=cortex.quickshow(voxel_vol, with_rois=True, with_colorbar=True, colorbar=True) # , cmap='Blues'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e316a6b",
   "metadata": {},
   "source": [
    "## GradCAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "93762930",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GradCAMBrainVLM(object):\n",
    "    def __init__(self,model, target_model='brain'):\n",
    "        self.gradients = dict()\n",
    "        self.activations = dict()\n",
    "        def backward_hook(module, grad_input, grad_output):\n",
    "            self.gradients['value'] = grad_output[0].detach().cpu()\n",
    "            # print(grad_output[0])\n",
    "            return None\n",
    "        def forward_hook(module, input, output):\n",
    "            self.activations['value'] = output[0].detach().cpu()\n",
    "            # print(output[0])\n",
    "            return None\n",
    "\n",
    "        if target_model=='text':\n",
    "            target_layer = model.text_model.encoder.layers[21]\n",
    "        elif target_model=='vision':\n",
    "            target_layer = model.vision_model.encoder.layers[29]\n",
    "        elif target_model=='brain':\n",
    "            target_layer = model.brain_model.encoder.layers[9]\n",
    "        # if target_model=='text':\n",
    "        #     target_layer = model.text_model.encoder.layers[0]\n",
    "        # elif target_model=='vision':\n",
    "        #     target_layer = model.vision_model.encoder.layers[0]\n",
    "        # elif target_model=='brain':\n",
    "        #     target_layer = model.brain_model.encoder.layers[0]\n",
    "\n",
    "        self.handles = []\n",
    "        h = target_layer.register_forward_hook(forward_hook)\n",
    "        self.handles.append(h)\n",
    "        target_layer.register_backward_hook(backward_hook)\n",
    "        self.handles.append(h)\n",
    "        self.model = model\n",
    "        self.target_model = target_model\n",
    "\n",
    "\n",
    "    def forward(self, batch):\n",
    "        # self.model.eval()\n",
    "        out = {}\n",
    "        for k in batch.keys():\n",
    "            if k=='text':\n",
    "                text_embeds = self.model.get_text_features(**batch[k])\n",
    "                text_embeds = text_embeds / get_vector_norm(text_embeds)\n",
    "                out.update({'text':text_embeds})\n",
    "            elif k=='vision':\n",
    "                vision_embeds = self.model.get_image_features(batch[k])\n",
    "                vision_embeds = vision_embeds / get_vector_norm(vision_embeds)\n",
    "                out.update({'vision':vision_embeds})\n",
    "            elif k=='brain':\n",
    "                brain_embeds = self.model.get_brain_features(batch[k])\n",
    "                brain_embeds = brain_embeds / get_vector_norm(brain_embeds)\n",
    "                out.update({'brain':brain_embeds})\n",
    "        self.model.zero_grad()\n",
    "        if self.target_model=='text':\n",
    "            # logits_per_text_v = torch.matmul(out['text']['embeds'], out['vision']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "            logits_per_text_b = torch.matmul(out['text'], out['brain'].t()) * self.model.logit_scale.exp().cuda()\n",
    "            # logits_per_text = (logits_per_text_v + logits_per_text_b) / 2\n",
    "            logits_per_text = logits_per_text_b\n",
    "            logit = logits_per_text[0,0]\n",
    "        elif self.target_model=='vision':\n",
    "            # logits_per_image_t = torch.matmul(out['vision']['embeds'], out['text']['embeds'].t()) * model.logit_scale.exp().cuda()\n",
    "            logits_per_image_b = torch.matmul(out['vision'], out['brain'].t()) * self.model.logit_scale.exp().cuda()\n",
    "            # logits_per_image = (logits_per_image_t + logits_per_image_b) / 2\n",
    "            logits_per_image = logits_per_image_b\n",
    "            logit = logits_per_image[0,0]\n",
    "        elif self.target_model=='brain':\n",
    "            logits_per_brain_t = torch.matmul(out['brain'], out['text'].t()) * self.model.logit_scale.exp().cuda()\n",
    "            logits_per_brain_v = torch.matmul(out['brain'], out['vision'].t()) * self.model.logit_scale.exp().cuda()\n",
    "            logits_per_brain = (logits_per_brain_t + logits_per_brain_v) / 2\n",
    "            logit = logits_per_brain[0,0]\n",
    "        logit.backward()\n",
    "        gradients = self.gradients['value']\n",
    "        activations = self.activations['value']\n",
    "        # print(gradients.shape, activations.shape)\n",
    "        saliency_map = (gradients * activations).sum(-1)\n",
    "        # print(saliency_map)\n",
    "        saliency_map = torch.nn.functional.relu(saliency_map)\n",
    "        saliency_map = saliency_map[:,1:] if self.target_model!='text' else saliency_map\n",
    "        return saliency_map\n",
    "\n",
    "    def __call__(self, img,caption):\n",
    "        return self.forward(img,caption)\n",
    "\n",
    "    def remove_hooks(self):\n",
    "        for handle in self.handles:\n",
    "            handle.remove()\n",
    "\n",
    "def grad_cam(model, batch, target_model='brain'):\n",
    "    grad_cam = GradCAMBrainVLM(model, target_model=target_model)\n",
    "    saliency = grad_cam.forward(batch)\n",
    "    grad_cam.remove_hooks()\n",
    "    return saliency.cpu().detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "id": "1072ca35",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1373: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using a non-full backward hook when the forward contains multiple autograd Nodes \"\n",
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1348: UserWarning: Using non-full backward hooks on a Module that does not take as input a single Tensor or a tuple of Tensors is deprecated and will be removed in future versions. This hook will be missing some of the grad_input. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using non-full backward hooks on a Module that does not take as input a \"\n"
     ]
    }
   ],
   "source": [
    "target_model = 'text'\n",
    "saliency = grad_cam(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "e889a092",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 110240 characters]",
      "text/plain": [
       "<Figure size 1000x200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target_model = 'text'\n",
    "text_words = processor.tokenizer.convert_ids_to_tokens(batch[target_model]['input_ids'][0][batch[target_model]['attention_mask'][0]==1].tolist())\n",
    "text_words = [x.split('<')[0] for x in text_words[1:-1]]\n",
    "tmap = heatmap(saliency[0].numpy(), target_model)\n",
    "rgba_colors = generate_shades_with_alpha(tmap[1:-1])\n",
    "\n",
    "fig, axs = plt.subplots(1, 2)\n",
    "fig.set_size_inches(10,2) \n",
    "\n",
    "axs[0].imshow(utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']]))\n",
    "# if bb:\n",
    "#     for x, y, w, h in bb:\n",
    "#         rect = mpl.patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='r', facecolor='none')\n",
    "#         axs[0].add_patch(rect)\n",
    "axs[0].axis('off')\n",
    "\n",
    "plot_text_with_colors(axs[1], text_words, rgba_colors, max_width=250, max_height=2, fontsize=12)\n",
    "axs[1].set_xlim(0, 250)\n",
    "axs[1].set_ylim(-2.5, 2.5)\n",
    "axs[1].axis('off')\n",
    "plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "d36b726d",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_model = 'vision'\n",
    "saliency = grad_cam(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "id": "cd2c1764",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 816960 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = utils.np_to_Image(stimulus[test_dict['subj01'][imgidx]['stim_idx']]).resize((224, 224))\n",
    "vmap = heatmap(saliency, target_model)\n",
    "vmap = plt.cm.jet(vmap)[..., :3] * 255\n",
    "vmap = Image.fromarray(vmap.astype(np.uint8))\n",
    "overlay = Image.blend(img, vmap, alpha=0.4)\n",
    "\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.subplot(1, 3, 1); plt.imshow(img); plt.title(\"Original Image\"); plt.axis(\"off\")\n",
    "plt.subplot(1, 3, 2); plt.imshow(vmap, cmap='jet'); plt.title(\"Grad-CAM Map\"); plt.axis(\"off\")\n",
    "plt.subplot(1, 3, 3); plt.imshow(overlay); plt.title(\"Overlay\"); plt.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "74e6fb5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py:1348: UserWarning: Using non-full backward hooks on a Module that does not take as input a single Tensor or a tuple of Tensors is deprecated and will be removed in future versions. This hook will be missing some of the grad_input. Please use register_full_backward_hook to get the documented behavior.\n",
      "  warnings.warn(\"Using non-full backward hooks on a Module that does not take as input a \"\n"
     ]
    }
   ],
   "source": [
    "target_model = 'brain'\n",
    "saliency = grad_cam(model, batch, target_model=target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "2648cb90",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_149542/3096259938.py:34: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  token_ids = torch.tensor(token_ids)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22276 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bmap, bmap_2d = heatmap(saliency, target_model)\n",
    "# ==== 可视化某一切片 ====\n",
    "plt.imshow(bmap_2d[:, :, bmap_2d.shape[2] // 2].cpu(), cmap='jet', vmax=1, vmin=0)\n",
    "plt.title(\"Grad-CAM Heatmap (z-mid)\")\n",
    "plt.axis(\"off\")\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "cf118d86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed to get connection\n",
      "** (inkscape:154978): CRITICAL **: 16:16:41.747: dbus_g_proxy_new_for_name: assertion 'connection != NULL' failed\n",
      "** (inkscape:154978): CRITICAL **: 16:16:41.747: dbus_g_proxy_call: assertion 'DBUS_IS_G_PROXY (proxy)' failed\n",
      "** (inkscape:154978): CRITICAL **: 16:16:41.747: dbus_g_connection_register_g_object: assertion 'connection != NULL' failed\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1199648 characters]",
      "text/plain": [
       "<Figure size 2274x1024 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# import cortex\n",
    "subject = 'fsaverage'\n",
    "xfm = 'atlas_2mm'\n",
    "\n",
    "# 将数据包装成 pycortex volume（带掩膜）\n",
    "voxel_vol = cortex.Volume(bmap.numpy(), subject=subject, xfmname=xfm, vmax=1, vmin=0,cmap='RdBu_r')#,cmap='RGrB_tsi'CyanBlueGrayRedPink,BuBkRd\n",
    "# cortex.webshow(voxel_vol, cmap='hot', colorbar=True, with_rois=True)\n",
    "fig=cortex.quickshow(voxel_vol, with_rois=True, with_colorbar=True, colorbar=True) # , cmap='Blues'"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
