{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "37e940b3-c58b-4aef-b3eb-98a31e0b10c4",
   "metadata": {},
   "source": [
    "# Multimodal Models\n",
    "\n",
    "This notebook is from [AI for Beginners Curriculum](http://aka.ms/ai-beginners).\n",
    "\n",
    "## Experimenting with CLIP\n",
    "\n",
    "[CLIP](https://openai.com/blog/clip/) from OpenAI is released in public, so you can try to use it for different tasks, including zero-shot image classification. Keep in mind that it is quite greedy in terms of resources!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6a187374-f4fc-4318-b89c-486b1e28686a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install git+https://github.com/openai/CLIP.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1574d33a-adaa-49dd-8a12-51a2a3426d18",
   "metadata": {},
   "source": [
    "We will first make sure that we can use GPU if it's available, and then load CLIP model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "907a1bc9-abfa-4290-9c66-729d45ed8cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import clip\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "np.set_printoptions(precision=2,suppress=True)\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model, preprocess = clip.load(\"ViT-B/32\", device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e796a87",
   "metadata": {},
   "source": [
    "Let's also get a subset of cats images from [Oxford-IIIT Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fb3a2099",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://mslearntensorflowlp.blob.core.windows.net/data/oxcats.tar.gz\n",
    "!tar xfz oxcats.tar.gz\n",
    "!rm oxcats.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b9f0991-9987-4a2a-bb51-19cdfb2c35af",
   "metadata": {},
   "source": [
    "### Zero-Shot Image Classification\n",
    "\n",
    "The main thing CLIP can do is to match an image with a text prompt. So, if we take an image, say, of a cat, and then try to match it with text promps \"a cat\", \"a penguin\", \"a bear\" - the first one is likely to have higher probability. Thus we can conlcude that we are dealing with a cat. We don't need to train a model because it has already been pre-trained on a huge dataset - thus it is called **zero-shot**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5b7712b9-7e85-44ac-8a85-75e29f5e86ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Label probs: [[0. 0. 1.]]\n"
     ]
    }
   ],
   "source": [
    "image = preprocess(Image.open(\"oxcats/Maine_Coon_1.jpg\")).unsqueeze(0).to(device)\n",
    "text = clip.tokenize([\"a penguin\", \"a bear\", \"a cat\"]).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    image_features = model.encode_image(image)\n",
    "    text_features = model.encode_text(text)\n",
    "    \n",
    "    logits_per_image, logits_per_text = model(image, text)\n",
    "    probs = logits_per_image.softmax(dim=-1).cpu().numpy()\n",
    "\n",
    "print(\"Label probs:\", probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d66738a-f72c-4334-af71-5021f9942a39",
   "metadata": {},
   "source": [
    "### Intelligent Image Search\n",
    "\n",
    "In the previous example, there was one image the 3 text prompts. We can use CLIP in a different context, eg. we can take many images of a cat, and then select an image the best suits the textual description:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0f8c3b95-3409-4a37-9f40-cca6795e08e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Img Index: 97\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x2f58593b310>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 164728 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "cats_img = [ Image.open(os.path.join(\"oxcats\",x)) for x in os.listdir(\"oxcats\") ] \n",
    "cats = torch.cat([ preprocess(i).unsqueeze(0) for i in cats_img ]).to(device)\n",
    "text = clip.tokenize([\"a very fat gray cat\"]).to(device)\n",
    "with torch.no_grad():\n",
    "    logits_per_image, logits_per_text = model(cats, text)\n",
    "    res = logits_per_text.softmax(dim=-1).argmax().cpu().numpy()\n",
    "\n",
    "print(\"Img Index:\", res)\n",
    "\n",
    "plt.imshow(cats_img[res])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21a91394",
   "metadata": {},
   "source": [
    "## Takeaway\n",
    "\n",
    "Pre-trained CLIP model can be used to perform tasks such as image classification on common objects without domain-specific training. Also, it allows more flexible classification/image search, taking into account spatial configuration of objects on the picture.\n",
    "\n",
    "For another exciting usage of CLIP, look for **VQGAN+CLIP**."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "16af2a8bbb083ea23e5e41c7f5787656b2ce26968575d8763f2c4b17f9cd711f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('py38')",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
