exp_localizability.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
    "\n",
    "from load import *\n",
    "import torchmetrics\n",
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import numpy as np\n",
    "\n",
    "from explainer import gradCAM, interpret\n",
    "from torch.utils.data import DataLoader, RandomSampler\n",
    "from scipy.ndimage import filters\n",
    "\n",
    "seed_everything(hparams['seed'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloader(dataset, seed, bs):\n",
    "    # Create a RandomSampler with a fixed seed for shuffling\n",
    "    sampler = RandomSampler(dataset, replacement=False, num_samples=None, generator=torch.Generator().manual_seed(seed))\n",
    "    dataloader = DataLoader(\n",
    "        dataset, \n",
    "        batch_size=bs,\n",
    "        sampler=sampler,\n",
    "        num_workers=16, \n",
    "        pin_memory=False\n",
    "    )\n",
    "    return dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = ImageNet(hparams['data_dir'], split='val', transform=tfms)\n",
    "# dataset = CUBDataset(hparams['data_dir'], train=False, transform=tfms)\n",
    "# dataset = torchvision.datasets.OxfordIIITPet(root=hparams['data_dir'], transform=tfms, split='test')\n",
    "# dataset = test_set # EuroSAT\n",
    "# dataset = torchvision.datasets.Food101(root=hparams['data_dir'], transform=tfms, split='test')\n",
    "\n",
    "bs = 8\n",
    "seed = 123\n",
    "dataloader = get_dataloader(dataset, seed, bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(hparams['device'])\n",
    "model, preprocess = clip.load(hparams['model_size'], device=device, jit=False) #Best model use ViT-B/32\n",
    "checkpoint = torch.load(\"/path/to/your/model.pt\")\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_insertion_fidelity_batched_blur(model, images, texts, heatmaps, insertion_steps=5):\n",
    "    device = images.device\n",
    "    batch_size, C, H, W = images.shape\n",
    "    cos = nn.CosineSimilarity(dim=1, eps=1e-6)\n",
    "\n",
    "    # Create blurred version of the images\n",
    "    blurred_images = torch.nn.functional.interpolate(images, scale_factor=0.5, mode='bilinear', align_corners=False)\n",
    "    blurred_images = torch.nn.functional.interpolate(blurred_images, size=(H, W), mode='bilinear', align_corners=False)\n",
    "\n",
    "    # 1. Determine the model's output for the completely blurred images\n",
    "    image_embedding = model.encode_image(blurred_images.to(device))\n",
    "    text_embedding = model.encode_text(texts.to(device))\n",
    "    blurred_output = cos(image_embedding, text_embedding)\n",
    "\n",
    "    # Determine the model's output for the original images\n",
    "    image_embedding = model.encode_image(images.to(device))\n",
    "    original_output = cos(image_embedding, text_embedding)\n",
    "\n",
    "    max_differences = torch.abs(original_output - blurred_output)\n",
    "\n",
    "    # 2. Sort pixels by importance for each image in the batch\n",
    "    _, indices = torch.sort(heatmaps.view(batch_size, -1), descending=True, dim=1)\n",
    "    total_pixels = indices.shape[1]\n",
    "\n",
    "    # Placeholder for storing normalized changes\n",
    "    differences = torch.zeros((batch_size, insertion_steps), device=device)\n",
    "\n",
    "    for step in range(1, insertion_steps + 1):\n",
    "        # 3. Gradually insert the image based on pixel importance\n",
    "        fraction = step / insertion_steps\n",
    "        num_insert = int(fraction * total_pixels)\n",
    "\n",
    "        # Start with blurred images\n",
    "        inserted_images = blurred_images.clone()\n",
    "        for idx, image in enumerate(inserted_images):\n",
    "            flat_image = image.view(-1)\n",
    "            flat_original = images[idx].view(-1)\n",
    "            flat_image[indices[idx, :num_insert]] = flat_original[indices[idx, :num_insert]]\n",
    "\n",
    "        # Determine model's output for the inserted images\n",
    "        image_embedding = model.encode_image(inserted_images.to(device))\n",
    "        inserted_output = cos(image_embedding, text_embedding)\n",
    "\n",
    "        # Compute normalized change\n",
    "        difference = torch.abs(inserted_output - blurred_output)\n",
    "        differences[:, step - 1] = difference\n",
    "\n",
    "    # Normalize changes\n",
    "    normalized_changes = torch.clamp(differences / max_differences.unsqueeze(1), 0, 1)\n",
    "\n",
    "    # 4. Calculate normalized fidelity (average over the insertion steps)\n",
    "    fidelity = normalized_changes.mean(dim=1)\n",
    "\n",
    "    return fidelity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_fidelity = 0.0\n",
    "count = 0 \n",
    "for batch_number, batch in enumerate(tqdm(dataloader)):\n",
    "    images, labels = batch\n",
    "\n",
    "    texts = np.array(label_to_classname)[labels].tolist()\n",
    "\n",
    "    concept_list = []\n",
    "    tokenized_concepts_list = []\n",
    "    for i in range(len(texts)):\n",
    "        concepts = gpt_descriptions[texts[i]][:5].copy()\n",
    "        # concepts.insert(0, texts[i])\n",
    "        concept_list.append(concepts)\n",
    "        tokenized_concepts = clip.tokenize(concepts)\n",
    "        tokenized_concepts_list.append(tokenized_concepts)\n",
    "\n",
    "    tokenized_text = clip.tokenize(texts)\n",
    "\n",
    "    images = images.to(device)\n",
    "    texts = tokenized_text.to(device)\n",
    "\n",
    "    attn_map = []\n",
    "    if hparams['model_size'] in ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64']:\n",
    "        for j in range(len(images)):\n",
    "            repeated_input = images[j].unsqueeze(0).repeat(tokenized_concepts_list[0].shape[0], 1, 1, 1)\n",
    "            attn = gradCAM(\n",
    "                model.visual,\n",
    "                repeated_input.to(device),\n",
    "                model.encode_text(tokenized_concepts_list[j].to(device)).float(),\n",
    "                getattr(model.visual, \"layer4\")\n",
    "            )\n",
    "            attn = F.interpolate(\n",
    "                attn.unsqueeze(0),\n",
    "                images.shape[2:],\n",
    "                mode='bicubic',\n",
    "                align_corners=False)\n",
    "            attn = attn.squeeze()\n",
    "            attn_map.append(attn)\n",
    "\n",
    "    elif hparams['model_size'] in ['ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']:\n",
    "        for k in range(len(images)):\n",
    "            R_image = interpret(model=model, image=images[k].unsqueeze(0), texts=tokenized_concepts_list[k].to(device), device=device)\n",
    "            image_relevance = R_image[0]\n",
    "            dim = int(image_relevance.numel() ** 0.5)\n",
    "            R_image = R_image.reshape(-1, dim, dim)\n",
    "            attn = F.interpolate(\n",
    "                R_image.unsqueeze(0),\n",
    "                images.shape[2:],\n",
    "                mode='bicubic',\n",
    "                align_corners=False)\n",
    "            attn = attn.squeeze()\n",
    "            attn_map.append(attn)\n",
    "    \n",
    "    attn_map = torch.stack(attn_map).reshape(-1, 224, 224)\n",
    "    repeated_images = images.repeat(tokenized_concepts_list[0].shape[0], 1, 1,1)\n",
    "    tokenized_concepts_list = torch.stack(tokenized_concepts_list).reshape(-1, 77).to(device)\n",
    "\n",
    "    fidelities = compute_insertion_fidelity_batched_blur(\n",
    "        model=model,\n",
    "        images=repeated_images,\n",
    "        texts=tokenized_concepts_list,\n",
    "        heatmaps=attn_map\n",
    "    )\n",
    "    if np.isnan(torch.mean(fidelities).item()):\n",
    "        continue\n",
    "    else:\n",
    "        total_fidelity += torch.mean(fidelities).item()\n",
    "        count += 1\n",
    "\n",
    "    del images, labels, texts, attn_map, repeated_images, tokenized_concepts_list, fidelities\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "print ('The final explanation fidelity: ' + str(total_fidelity / count))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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.11.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}