{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2b949f9f",
   "metadata": {},
   "source": [
    "# BLIP: Inference Demo\n",
    " - [Image Captioning](#Image-Captioning)\n",
    " - [VQA](#VQA)\n",
    " - [Feature Extraction](#Feature-Extraction)\n",
    " - [Image Text Matching](#Image-Text-Matching)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cbcb066b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install requirements\n",
    "import sys\n",
    "if 'google.colab' in sys.modules:\n",
    "    print('Running in Colab.')\n",
    "    !pip3 install transformers==4.15.0 timm==0.4.12 fairscale==0.4.4\n",
    "    !git clone https://github.com/salesforce/BLIP\n",
    "    %cd BLIP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a811a65f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torchvision.transforms.functional import InterpolationMode\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def load_demo_image(image_size,device):\n",
    "    img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' \n",
    "    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')   \n",
    "\n",
    "    w,h = raw_image.size\n",
    "    display(raw_image.resize((w//5,h//5)))\n",
    "    \n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((image_size,image_size),interpolation=InterpolationMode.BICUBIC),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
    "        ]) \n",
    "    image = transform(raw_image).unsqueeze(0).to(device)   \n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f72f4406",
   "metadata": {},
   "source": [
    "# Image Captioning\n",
    "Perform image captioning using finetuned BLIP model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6835daef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 169592 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=409x273 at 0x7F9205179040>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_base_caption.pth\n",
      "caption: a woman sitting on the beach with a dog\n"
     ]
    }
   ],
   "source": [
    "from models.blip import blip_decoder\n",
    "\n",
    "image_size = 384\n",
    "image = load_demo_image(image_size=image_size, device=device)\n",
    "\n",
    "model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'\n",
    "    \n",
    "model = blip_decoder(pretrained=model_url, image_size=image_size, vit='base')\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    # beam search\n",
    "    caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5) \n",
    "    # nucleus sampling\n",
    "    # caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5) \n",
    "    print('caption: '+caption[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fac320a2",
   "metadata": {},
   "source": [
    "# VQA\n",
    "Perform visual question answering using finetuned BLIP model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5e6f3fb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 169592 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=409x273 at 0x7F92049303D0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_vqa.pth\n",
      "answer: on beach\n"
     ]
    }
   ],
   "source": [
    "from models.blip_vqa import blip_vqa\n",
    "\n",
    "image_size = 480\n",
    "image = load_demo_image(image_size=image_size, device=device)     \n",
    "\n",
    "model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth'\n",
    "    \n",
    "model = blip_vqa(pretrained=model_url, image_size=image_size, vit='base')\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "\n",
    "question = 'where is the woman sitting?'\n",
    "\n",
    "with torch.no_grad():\n",
    "    answer = model(image, question, train=False, inference='generate') \n",
    "    print('answer: '+answer[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6100e519",
   "metadata": {},
   "source": [
    "# Feature Extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4f8f21ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 169592 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=409x273 at 0x7FBCE3506550>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base.pth\n"
     ]
    }
   ],
   "source": [
    "from models.blip import blip_feature_extractor\n",
    "\n",
    "image_size = 224\n",
    "image = load_demo_image(image_size=image_size, device=device)     \n",
    "\n",
    "model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base.pth'\n",
    "    \n",
    "model = blip_feature_extractor(pretrained=model_url, image_size=image_size, vit='base')\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "\n",
    "caption = 'a woman sitting on the beach with a dog'\n",
    "\n",
    "multimodal_feature = model(image, caption, mode='multimodal')[0,0]\n",
    "image_feature = model(image, caption, mode='image')[0,0]\n",
    "text_feature = model(image, caption, mode='text')[0,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "201e1146",
   "metadata": {},
   "source": [
    "# Image-Text Matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "49ba5906",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 169592 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=409x273 at 0x7FBCE357CCA0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth\n",
      "text: a woman sitting on the beach with a dog\n",
      "The image and text is matched with a probability of 0.9960\n",
      "The image feature and text feature has a cosine similarity of 0.5262\n"
     ]
    }
   ],
   "source": [
    "from models.blip_itm import blip_itm\n",
    "\n",
    "image_size = 384\n",
    "image = load_demo_image(image_size=image_size,device=device)\n",
    "\n",
    "model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth'\n",
    "    \n",
    "model = blip_itm(pretrained=model_url, image_size=image_size, vit='base')\n",
    "model.eval()\n",
    "model = model.to(device='cpu')\n",
    "\n",
    "caption = 'a woman sitting on the beach with a dog'\n",
    "\n",
    "print('text: %s' %caption)\n",
    "\n",
    "itm_output = model(image,caption,match_head='itm')\n",
    "itm_score = torch.nn.functional.softmax(itm_output,dim=1)[:,1]\n",
    "print('The image and text is matched with a probability of %.4f'%itm_score)\n",
    "\n",
    "itc_score = model(image,caption,match_head='itc')\n",
    "print('The image feature and text feature has a cosine similarity of %.4f'%itc_score)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
