{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import os\n",
    "import torch\n",
    "import clip\n",
    "from scripts.clip_wrapper import ClipWrapper\n",
    "from scripts.plot import visualize_vandt_heatmap\n",
    "from scripts.methods import vision_heatmap_iba, text_heatmap_iba\n",
    "from transformers import CLIPProcessor, CLIPModel, CLIPTokenizerFast\n",
    "from PIL import Image, ImageOps\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(model, image, text, vbeta=0.1, vvar=1, vlayer=9, tbeta=0.1, tvar=1, tlayer=9, bb=None):\n",
    "    # Preprocess image\n",
    "    image = Image.open(image).convert('RGB')\n",
    "    image_feat = processor(images=image, return_tensors=\"pt\")['pixel_values'].to(device) # 3*224*224\n",
    "    # Tokenize text\n",
    "    text_ids = torch.tensor([tokenizer.encode(text, add_special_tokens=True)]).to(device)\n",
    "    text_words = tokenizer.convert_ids_to_tokens(text_ids[0].tolist())\n",
    "    # Train information bottleneck on image\n",
    "    print(\"Training M2IB on the image...\")\n",
    "    vmap = vision_heatmap_iba(text_ids, image_feat, model, vlayer, vbeta, vvar)\n",
    "    # Train information bottleneck on text\n",
    "    print(\"Training M2IB on the text...\")\n",
    "    tmap = text_heatmap_iba(text_ids, image_feat, model, tlayer, tbeta, tvar)\n",
    "    image_under = processor(images=image, return_tensors=\"pt\", do_normalize=False)['pixel_values'][0].permute(1,2,0) # no normalization\n",
    "    if bb:\n",
    "        scaled_bb = []\n",
    "        for x, y, w, h in bb:\n",
    "            hr, vr = image_under.shape[0]/image.size[0], image_under.shape[1]/image.size[1]\n",
    "            scaled_bb.append((x*hr, y*vr, w*hr, h*vr))\n",
    "        bb = scaled_bb\n",
    "    visualize_vandt_heatmap(tmap, vmap, text_words, image_under, bb=bb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device)\n",
    "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "tokenizer = CLIPTokenizerFast.from_pretrained(\"openai/clip-vit-base-patch32\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conceptual Captions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conceptual Caption provides a large amount of images and their raw descriptions collected from the web. The examples below are from its validation split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the image...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the text...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.26it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 120396 characters]\n",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = \"images/mouse.jpg\"\n",
    "text = \"hungry mouse eating cheese in a mousetrap.\"\n",
    "plot(model, img, text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the image...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the text...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.13it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 143540 characters]\n",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = \"images/kites.jpg\"\n",
    "text = \"the kites flown at festival turned the sky into a kaleidoscope of color .\"\n",
    "plot(model, img, text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the image...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:05<00:00,  1.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the text...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:05<00:00,  2.00it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 164856 characters]\n",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = \"images/castle.jpg\"\n",
    "text = \"where 's the best place to show off your nails ? right in front of the castle , of course !\"\n",
    "plot(model, img, text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MS-CXR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MS-CXR dataset is a radiology dataset that provides 1162 image–sentence pairs of bounding boxes and corresponding phrases. Please download the dataset from https://physionet.org/content/ms-cxr/0.1/. We provide two examples below for illustration purpose.\n",
    "- To be able to run the following code, you need to first download checkpoint finetuned on xray data from https://github.com/rajpurkarlab/CXR-RePaiR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "med_model, _ = clip.load(\"ViT-B/32\", device=device)\n",
    "med_model.load_state_dict(torch.load(\"clip-imp-pretrained_128_6_after_4.pt\", map_location=device))\n",
    "med_model = ClipWrapper(med_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the image...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the text...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.08it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 114776 characters]\n",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = \"images/cxr_upper.jpg\"\n",
    "text = \"ill-defined opacities largely in the upper lobe\"\n",
    "bb = [(347,605,600,1052),(1337,765,899,1127)]\n",
    "plot(med_model, img, text, vbeta = 1, vvar=1, vlayer=9, tbeta=1, tvar=1, tlayer=7, bb=bb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the image...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training M2IB on the text...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Bottleneck: 100%|██████████| 10/10 [00:04<00:00,  2.27it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 111616 characters]\n",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = \"images/cxr_left.jpg\"\n",
    "text = \"local consolidation in the left midlung\"\n",
    "bb = [(1767,932,800,1063)]\n",
    "plot(med_model, img, text, vbeta = 1, vvar=1, vlayer=9, tbeta=1, tvar=1, tlayer=7, bb=bb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
