{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "33d8b132",
   "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\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": "f77c5476",
   "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": "e1058f87",
   "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": 4,
   "id": "15cacfea",
   "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": "57b3bcfd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/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",
      "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n"
     ]
    },
    {
     "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",
    "        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": "3e88e670",
   "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",
    "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": 7,
   "id": "3af92014",
   "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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "80a160a9",
   "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": 9,
   "id": "dbbaa4b1",
   "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": "590230a8",
   "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": "bf92f374",
   "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",
    "    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_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": 12,
   "id": "c189b26c",
   "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": 13,
   "id": "79078f25",
   "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": 14,
   "id": "5ba707c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.brain_model = clip_fmri.vision_model\n",
    "model.brain_projection = clip_fmri.visual_projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "781a2271",
   "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": "d8fff310",
   "metadata": {},
   "outputs": [],
   "source": [
    "from M2IB.scripts.iba import *\n",
    "class IBAInterpreter:\n",
    "    def __init__(self, model, estim: Estimator, beta, steps=10, lr=1, batch_size=10, progbar=False):\n",
    "        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "        self.model = model.to(self.device)\n",
    "        self.original_layer = estim.get_layer()\n",
    "        self.shape = estim.shape()\n",
    "        self.beta = beta\n",
    "        self.batch_size = batch_size\n",
    "        self.fitting_estimator = torch.nn.CosineSimilarity(eps=1e-6)\n",
    "        self.progbar = progbar\n",
    "        self.lr = lr\n",
    "        self.train_steps = steps\n",
    "        self.bottleneck = InformationBottleneck(estim.mean(), estim.std(), device=self.device)\n",
    "        self.sequential = mySequential(self.original_layer, self.bottleneck)\n",
    "\n",
    "    def text_heatmap(self, text_t, image_t):\n",
    "        saliency, loss_c, loss_f, loss_t = self._run_text_training(text_t, image_t)\n",
    "        saliency = torch.nansum(saliency, -1).cpu().detach().numpy()\n",
    "        saliency = normalize(saliency)\n",
    "        return normalize(saliency)\n",
    "    \n",
    "    def brain_heatmap(self, text_t, image_t):\n",
    "        saliency, loss_c, loss_f, loss_t = self._run_vision_training(text_t, image_t)\n",
    "        saliency = torch.nansum(saliency, -1)[1:] \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",
    "\n",
    "    def vision_heatmap(self, text_t, image_t):\n",
    "        saliency, loss_c, loss_f, loss_t = self._run_vision_training(text_t, image_t)\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",
    "\n",
    "    def get_saliency(self, batch, target_model:str='brain'):\n",
    "        saliency, loss_c, loss_f, loss_t = self._run_training(batch, target_model)\n",
    "        return saliency, loss_c, loss_f, loss_t\n",
    "\n",
    "    def _run_training(self, batch, target_model:str='brain'):\n",
    "        assert batch[target_model] is not None, f\"Please provide {target_model}\"\n",
    "        assert any(not v is None for k, v in batch.items() if k != target_model), f\"must provide values for any other models except {target_model}\"\n",
    "\n",
    "        if target_model=='text':\n",
    "            replace_layer(self.model.text_model, self.original_layer, self.sequential)\n",
    "        elif target_model=='vision':\n",
    "            replace_layer(self.model.vision_model, self.original_layer, self.sequential)\n",
    "        elif target_model=='brain':\n",
    "            replace_layer(self.model.brain_model, self.original_layer, self.sequential)\n",
    "        else:\n",
    "            raise NotImplementedError(f\"target_model {target_model} not implemented!\")\n",
    "        loss_c, loss_f, loss_t = self._train_bottleneck(batch, target_model)\n",
    "        if target_model=='text':\n",
    "            replace_layer(self.model.text_model, self.sequential, self.original_layer)\n",
    "        elif target_model=='vision':\n",
    "            replace_layer(self.model.vision_model, self.sequential, self.original_layer)\n",
    "        elif target_model=='brain':\n",
    "            replace_layer(self.model.brain_model, self.sequential, self.original_layer)\n",
    "        return self.bottleneck.buffer_capacity.mean(axis=0), loss_c, loss_f, loss_t\n",
    "\n",
    "    def _train_bottleneck(self, batch, target_model:str='brain'):\n",
    "        optimizer = torch.optim.Adam(lr=self.lr, params=self.bottleneck.parameters())\n",
    "        # Reset from previous run or modifications\n",
    "        self.bottleneck.reset_alpha()\n",
    "        # Train\n",
    "        self.model.eval()\n",
    "        out = {}\n",
    "        for k,v in batch.items():\n",
    "            if v is not None:\n",
    "                if k == 'text':\n",
    "                    # batch[k] = batch[k].expand(self.batch_size, -1)\n",
    "                    batch[k]['input_ids'] = batch[k]['input_ids'].expand(self.batch_size, -1)\n",
    "                    batch[k]['attention_mask'] = batch[k]['attention_mask'].expand(self.batch_size, -1)\n",
    "                elif k == 'vision':\n",
    "                    batch[k] = batch[k].expand(self.batch_size, -1, -1, -1)\n",
    "                elif k == 'brain':\n",
    "                    batch[k] = batch[k].expand(self.batch_size, -1, -1, -1, -1)\n",
    "                else:\n",
    "                    raise NotImplementedError(f\"batch key {k} not implemented!\")\n",
    "        for _ in tqdm(range(self.train_steps), desc=\"Training Bottleneck\",\n",
    "                      disable=not self.progbar):\n",
    "            optimizer.zero_grad()\n",
    "            for k,v in batch.items():\n",
    "                if v is not None:\n",
    "                    if k == 'text':\n",
    "                        out['text'] = self.model.get_text_features(**batch[k])\n",
    "                    elif k == 'vision':\n",
    "                        out['vision'] = self.model.get_image_features(batch[k])\n",
    "                    elif k == 'brain':\n",
    "                        out['brain'] = self.model.get_brain_features(batch[k])\n",
    "\n",
    "            loss_c, loss_f, loss_t = 0, 0, 0\n",
    "            for k in out:\n",
    "                if k != target_model:\n",
    "                    _loss_c, _loss_f, _loss_t = self.calc_loss(outputs=out[target_model], labels=out[k])\n",
    "                    loss_c += _loss_c\n",
    "                    loss_f += _loss_f\n",
    "                    loss_t += _loss_t\n",
    "            loss_t.backward()\n",
    "            optimizer.step(closure=None)\n",
    "        return loss_c, loss_f, loss_t \n",
    "\n",
    "    def calc_loss(self, outputs, labels):\n",
    "        \"\"\" Calculate the combined loss expression for optimization of lambda \"\"\"\n",
    "        compression_term = self.bottleneck.buffer_capacity.mean()\n",
    "        fitting_term = self.fitting_estimator(outputs, labels).mean()\n",
    "        total =  self.beta * compression_term - fitting_term\n",
    "        return compression_term, fitting_term, total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "04a8a65e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from M2IB.scripts.methods import extract_feature_map, extract_bert_layer, get_compression_estimator\n",
    "def get_saliency_iba(batch, target_model, model, layer_idx, beta, var, lr=1, train_steps=10, progbar=True):\n",
    "    features = extract_feature_map(model.get_submodule(f'{target_model}_model'), layer_idx, batch[target_model])\n",
    "    layer = extract_bert_layer(model.get_submodule(f'{target_model}_model'), layer_idx)\n",
    "    compression_estimator = get_compression_estimator(var, layer, features)\n",
    "    reader = IBAInterpreter(model, compression_estimator, beta=beta, lr=lr, steps=train_steps, progbar=progbar)\n",
    "    return reader.get_saliency(batch, target_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8236b236",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==== 基础参数 ====\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)\n",
    "        # saliency = torch.nn.functional.relu(saliency).cpu().detach().numpy()\n",
    "        saliency = normalize(saliency.cpu().detach().numpy())\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",
    "        # saliency = torch.nn.functional.relu(saliency)\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 = torch.nn.functional.relu(saliency).cpu().detach()\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": "code",
   "execution_count": 19,
   "id": "d7012b6a",
   "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": 20,
   "id": "78232dc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(idx):\n",
    "    batch = {}\n",
    "    data = test_dict[f'subj{args.subj:02d}'][idx]\n",
    "\n",
    "    n = len(data['text'])\n",
    "    random_idx = random.randint(0, n-1)\n",
    "    input_ids = tokenized_captions[str(data['cocoid'])]['input_ids'][random_idx]\n",
    "    attention_mask = tokenized_captions[str(data['cocoid'])]['attention_mask'][random_idx]\n",
    "    # text_ids = torch.tensor([processor.tokenizer.encode(data['text'][random_idx], add_special_tokens=True)])\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",
    "    batch.update({'text': {'input_ids': input_ids.to(device).unsqueeze(0), 'attention_mask': attention_mask.to(device).unsqueeze(0)}})\n",
    "    # batch.update({'text': text_ids.to(device)})\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": 88,
   "id": "e533f974",
   "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": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imgidx = 13 # 8 # 49 # 121\n",
    "\n",
    "utils.np_to_Image(stimulus[test_dict[f'subj{args.subj:02d}'][imgidx]['stim_idx']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "e378b4f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/data/private/src/fMRI/autodl-a100-40/kaixinkeng/src/M2IB/scripts/iba.py:95: 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",
      "  self.mean = torch.tensor(mean, dtype=torch.float, device=self.device, requires_grad=False)\n",
      "Training Bottleneck: 100%|██████████| 20/20 [00:01<00:00, 10.35it/s]\n"
     ]
    }
   ],
   "source": [
    "batch = get_batch(imgidx)\n",
    "target_model = 'text'\n",
    "# batch['text']=None\n",
    "batch['vision']=None\n",
    "# batch['brain']=None\n",
    "saliency, _, _, _ = get_saliency_iba(batch, target_model, model, layer_idx=6, beta=0.1, var=1, train_steps=20, progbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "e365d38d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([11, 1024])"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if target_model=='text':\n",
    "    saliency = saliency[batch['text']['attention_mask'][0].bool()]\n",
    "saliency.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ef23f3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 379892 characters]",
      "text/plain": [
       "<Figure size 400x800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "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, target_model)\n",
    "rgba_colors = generate_shades_with_alpha(tmap[1:-1])\n",
    "\n",
    "fig, axs = plt.subplots(2, 1)\n",
    "fig.set_size_inches(4,8) \n",
    "\n",
    "axs[0].imshow(utils.np_to_Image(stimulus[test_dict[f'subj{args.subj:02d}'][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=150, max_height=2.4, fontsize=12)\n",
    "axs[1].set_xlim(0, 150)\n",
    "axs[1].set_ylim(-5, 2)\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": 92,
   "id": "10884268",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 20/20 [00:05<00:00,  3.53it/s]\n"
     ]
    }
   ],
   "source": [
    "batch = get_batch(imgidx)\n",
    "target_model = 'vision'\n",
    "batch['text']=None\n",
    "# batch['vision']=None\n",
    "# batch['brain']=None\n",
    "saliency, _, _, _ = get_saliency_iba(batch, target_model, model, layer_idx=6, beta=0.1, var=1, train_steps=20, progbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc1242af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 333992 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# img = utils.np_to_Image(stimulus[test_dict[f'subj{args.subj:02d}'][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()\n",
    "\n",
    "img = utils.np_to_Image(stimulus[test_dict[f'subj{args.subj:02d}'][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=(4, 4))\n",
    "plt.imshow(overlay); plt.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "06749f9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/data/private/src/fMRI/autodl-a100-40/kaixinkeng/src/M2IB/scripts/iba.py:95: 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",
      "  self.mean = torch.tensor(mean, dtype=torch.float, device=self.device, requires_grad=False)\n",
      "Training Bottleneck: 100%|██████████| 10/10 [00:03<00:00,  3.10it/s]\n"
     ]
    }
   ],
   "source": [
    "imgidx=13\n",
    "batch = get_batch(imgidx)\n",
    "target_model = 'brain'\n",
    "# batch['text']=None\n",
    "# batch['vision']=None\n",
    "# batch['brain']=None\n",
    "saliency, _, _, _ = get_saliency_iba(batch, target_model, model, layer_idx=6, beta=0.1, var=0.1, lr=1, train_steps=10, progbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52414168",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_156437/1261442525.py:36: 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, 19336 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": null,
   "id": "4a158491",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed to get connection\n",
      "** (inkscape:193508): CRITICAL **: 21:09:52.747: dbus_g_proxy_new_for_name: assertion 'connection != NULL' failed\n",
      "** (inkscape:193508): CRITICAL **: 21:09:52.747: dbus_g_proxy_call: assertion 'DBUS_IS_G_PROXY (proxy)' failed\n",
      "** (inkscape:193508): CRITICAL **: 21:09:52.747: dbus_g_connection_register_g_object: assertion 'connection != NULL' failed\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 831064 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",
    "voxel_vol = cortex.Volume(bmap.numpy(), subject=subject, xfmname=xfm, cmap='RdBu_r', vmin=0, vmax=1)#,cmap=magma 'RGrB_tsi'CyanBlueGrayRedPink,BuBkRd RdBu_r\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
}
