{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Interacting with CLIP.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1369964d45004b5e95a058910b2a33e6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_12e23e2819094ee0a079d4eb77cfc4f9",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_7a5f52e56ede4ac3abe37a3ece007dc9",
              "IPY_MODEL_ce8b0faa1a1340b5a504d7b3546b3ccb"
            ]
          }
        },
        "12e23e2819094ee0a079d4eb77cfc4f9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "7a5f52e56ede4ac3abe37a3ece007dc9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_5e6adc4592124a4581b85f4c1f3bab4d",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 169001437,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 169001437,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_4a61c10fc00c4f04bb00b82e942da210"
          }
        },
        "ce8b0faa1a1340b5a504d7b3546b3ccb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_b597cd6f6cd443aba4bf4491ac7f957e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 169001984/? [00:06&lt;00:00, 25734958.25it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_161969cae25a49f38aacd1568d3cac6c"
          }
        },
        "5e6adc4592124a4581b85f4c1f3bab4d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "4a61c10fc00c4f04bb00b82e942da210": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "b597cd6f6cd443aba4bf4491ac7f957e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "161969cae25a49f38aacd1568d3cac6c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPHN7PJgKOzb"
      },
      "source": [
        "# Interacting with CLIP\n",
        "\n",
        "This is a self-contained notebook that shows how to download and run CLIP models, calculate the similarity between arbitrary image and text inputs, and perform zero-shot image classifications."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53N4k0pj_9qL"
      },
      "source": [
        "# Preparation for Colab\n",
        "\n",
        "Make sure you're running a GPU runtime; if not, select \"GPU\" as the hardware accelerator in Runtime > Change Runtime Type in the menu. The next cells will install the `clip` package and its dependencies, and check if PyTorch 1.7.1 or later is installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0BpdJkdBssk9",
        "outputId": "4d9b51f8-d255-4868-97f6-be0a67dadfae"
      },
      "source": [
        "! pip install ftfy regex tqdm\n",
        "! pip install git+https://github.com/openai/CLIP.git"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting ftfy\n",
            "  Downloading ftfy-6.0.3.tar.gz (64 kB)\n",
            "\u001b[?25l\r\u001b[K     |█████                           | 10 kB 34.6 MB/s eta 0:00:01\r\u001b[K     |██████████▏                     | 20 kB 20.7 MB/s eta 0:00:01\r\u001b[K     |███████████████▎                | 30 kB 16.4 MB/s eta 0:00:01\r\u001b[K     |████████████████████▍           | 40 kB 15.2 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▌      | 51 kB 7.0 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▋ | 61 kB 8.2 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 64 kB 2.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: regex in /usr/local/lib/python3.7/dist-packages (2019.12.20)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (4.41.1)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from ftfy) (0.2.5)\n",
            "Building wheels for collected packages: ftfy\n",
            "  Building wheel for ftfy (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for ftfy: filename=ftfy-6.0.3-py3-none-any.whl size=41934 sha256=ce00f21233e5e1a5c3e84204d651ae0c17403484418c330c69ebae7297ca7003\n",
            "  Stored in directory: /root/.cache/pip/wheels/19/f5/38/273eb3b5e76dfd850619312f693716ac4518b498f5ffb6f56d\n",
            "Successfully built ftfy\n",
            "Installing collected packages: ftfy\n",
            "Successfully installed ftfy-6.0.3\n",
            "Collecting git+https://github.com/openai/CLIP.git\n",
            "  Cloning https://github.com/openai/CLIP.git to /tmp/pip-req-build-tbmjxrgj\n",
            "  Running command git clone -q https://github.com/openai/CLIP.git /tmp/pip-req-build-tbmjxrgj\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.7/dist-packages (from clip==1.0) (6.0.3)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.7/dist-packages (from clip==1.0) (2019.12.20)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from clip==1.0) (4.41.1)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from clip==1.0) (1.9.0+cu102)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from clip==1.0) (0.10.0+cu102)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from ftfy->clip==1.0) (0.2.5)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch->clip==1.0) (3.7.4.3)\n",
            "Requirement already satisfied: pillow>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision->clip==1.0) (7.1.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision->clip==1.0) (1.19.5)\n",
            "Building wheels for collected packages: clip\n",
            "  Building wheel for clip (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for clip: filename=clip-1.0-py3-none-any.whl size=1369080 sha256=9059060f3a406b8268cfcbbf03647ab9b59e660d03f86c6120a9dd06f2b82cec\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-7v3mrcvj/wheels/fd/b9/c3/5b4470e35ed76e174bff77c92f91da82098d5e35fd5bc8cdac\n",
            "Successfully built clip\n",
            "Installing collected packages: clip\n",
            "Successfully installed clip-1.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C1hkDT38hSaP",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "70a44964-883d-4fd0-b95a-2c7f2b19aca9"
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "from pkg_resources import packaging\n",
        "\n",
        "print(\"Torch version:\", torch.__version__)\n"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Torch version: 1.9.0+cu102\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eFxgLV5HAEEw"
      },
      "source": [
        "# Loading the model\n",
        "\n",
        "`clip.available_models()` will list the names of available CLIP models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uLFS29hnhlY4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "11779e1e-8bdd-4167-c18e-d26bdd6b67db"
      },
      "source": [
        "import clip\n",
        "\n",
        "clip.available_models()"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['RN50', 'RN101', 'RN50x4', 'RN50x16', 'ViT-B/32', 'ViT-B/16']"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IBRVTY9lbGm8",
        "outputId": "f06fd2fd-6126-475b-87d0-b10aa3b7da49"
      },
      "source": [
        "model, preprocess = clip.load(\"ViT-B/32\")\n",
        "model.cuda().eval()\n",
        "input_resolution = model.visual.input_resolution\n",
        "context_length = model.context_length\n",
        "vocab_size = model.vocab_size\n",
        "\n",
        "print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}\")\n",
        "print(\"Input resolution:\", input_resolution)\n",
        "print(\"Context length:\", context_length)\n",
        "print(\"Vocab size:\", vocab_size)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|███████████████████████████████████████| 338M/338M [00:05<00:00, 63.0MiB/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Model parameters: 151,277,313\n",
            "Input resolution: 224\n",
            "Context length: 77\n",
            "Vocab size: 49408\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "21slhZGCqANb"
      },
      "source": [
        "# Image Preprocessing\n",
        "\n",
        "We resize the input images and center-crop them to conform with the image resolution that the model expects. Before doing so, we will normalize the pixel intensity using the dataset mean and standard deviation.\n",
        "\n",
        "The second return value from `clip.load()` contains a torchvision `Transform` that performs this preprocessing.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d6cpiIFHp9N6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "880cb98e-1e5e-430e-8b59-4bf35fa554f9"
      },
      "source": [
        "preprocess"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Compose(\n",
              "    Resize(size=224, interpolation=bicubic, max_size=None, antialias=None)\n",
              "    CenterCrop(size=(224, 224))\n",
              "    <function _transform.<locals>.<lambda> at 0x7f3a24ffb440>\n",
              "    ToTensor()\n",
              "    Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xwSB5jZki3Cj"
      },
      "source": [
        "# Text Preprocessing\n",
        "\n",
        "We use a case-insensitive tokenizer, which can be invoked using `clip.tokenize()`. By default, the outputs are padded to become 77 tokens long, which is what the CLIP models expects."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qGom156-i2kL",
        "outputId": "050b0ce1-caba-47e1-f4ac-dba994599718"
      },
      "source": [
        "clip.tokenize(\"Hello World!\")"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[49406,  3306,  1002,   256, 49407,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "             0,     0,     0,     0,     0,     0,     0]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4W8ARJVqBJXs"
      },
      "source": [
        "# Setting up input images and texts\n",
        "\n",
        "We are going to feed 8 example images and their textual descriptions to the model, and compare the similarity between the corresponding features.\n",
        "\n",
        "The tokenizer is case-insensitive, and we can freely give any suitable textual descriptions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tMc1AXzBlhzm"
      },
      "source": [
        "import os\n",
        "import skimage\n",
        "import IPython.display\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "\n",
        "from collections import OrderedDict\n",
        "import torch\n",
        "\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "# images in skimage to use and their textual descriptions\n",
        "descriptions = {\n",
        "    \"page\": \"a page of text about segmentation\",\n",
        "    \"chelsea\": \"a facial photo of a tabby cat\",\n",
        "    \"astronaut\": \"a portrait of an astronaut with the American flag\",\n",
        "    \"rocket\": \"a rocket standing on a launchpad\",\n",
        "    \"motorcycle_right\": \"a red motorcycle standing in a garage\",\n",
        "    \"camera\": \"a person looking at a camera on a tripod\",\n",
        "    \"horse\": \"a black-and-white silhouette of a horse\", \n",
        "    \"coffee\": \"a cup of coffee on a saucer\"\n",
        "}"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NSSrLY185jSf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 368
        },
        "outputId": "06451963-5ecb-4ddc-d0a8-24e9b110af7d"
      },
      "source": [
        "original_images = []\n",
        "images = []\n",
        "texts = []\n",
        "plt.figure(figsize=(16, 5))\n",
        "\n",
        "for filename in [filename for filename in os.listdir(skimage.data_dir) if filename.endswith(\".png\") or filename.endswith(\".jpg\")]:\n",
        "    name = os.path.splitext(filename)[0]\n",
        "    if name not in descriptions:\n",
        "        continue\n",
        "\n",
        "    image = Image.open(os.path.join(skimage.data_dir, filename)).convert(\"RGB\")\n",
        "  \n",
        "    plt.subplot(2, 4, len(images) + 1)\n",
        "    plt.imshow(image)\n",
        "    plt.title(f\"{filename}\\n{descriptions[name]}\")\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "\n",
        "    original_images.append(image)\n",
        "    images.append(preprocess(image))\n",
        "    texts.append(descriptions[name])\n",
        "\n",
        "plt.tight_layout()\n"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 1556924 characters]\n",
            "text/plain": [
              "<Figure size 1152x360 with 8 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "image/png": {
              "width": 1116,
              "height": 351
            }
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WEVKsji6WOIX"
      },
      "source": [
        "## Building features\n",
        "\n",
        "We normalize the images, tokenize each text input, and run the forward pass of the model to get the image and text features."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HBgCanxi8JKw"
      },
      "source": [
        "image_input = torch.tensor(np.stack(images)).cuda()\n",
        "text_tokens = clip.tokenize([\"This is \" + desc for desc in texts]).cuda()"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZN9I0nIBZ_vW"
      },
      "source": [
        "with torch.no_grad():\n",
        "    image_features = model.encode_image(image_input).float()\n",
        "    text_features = model.encode_text(text_tokens).float()"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cuxm2Gt4Wvzt"
      },
      "source": [
        "## Calculating cosine similarity\n",
        "\n",
        "We normalize the features and calculate the dot product of each pair."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yKAxkQR7bf3A"
      },
      "source": [
        "image_features /= image_features.norm(dim=-1, keepdim=True)\n",
        "text_features /= text_features.norm(dim=-1, keepdim=True)\n",
        "similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C5zvMxh8cU6m",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 831
        },
        "outputId": "22bca748-ab42-4888-c9da-8f22c21c6185"
      },
      "source": [
        "count = len(descriptions)\n",
        "\n",
        "plt.figure(figsize=(20, 14))\n",
        "plt.imshow(similarity, vmin=0.1, vmax=0.3)\n",
        "# plt.colorbar()\n",
        "plt.yticks(range(count), texts, fontsize=18)\n",
        "plt.xticks([])\n",
        "for i, image in enumerate(original_images):\n",
        "    plt.imshow(image, extent=(i - 0.5, i + 0.5, -1.6, -0.6), origin=\"lower\")\n",
        "for x in range(similarity.shape[1]):\n",
        "    for y in range(similarity.shape[0]):\n",
        "        plt.text(x, y, f\"{similarity[y, x]:.2f}\", ha=\"center\", va=\"center\", size=12)\n",
        "\n",
        "for side in [\"left\", \"top\", \"right\", \"bottom\"]:\n",
        "  plt.gca().spines[side].set_visible(False)\n",
        "\n",
        "plt.xlim([-0.5, count - 0.5])\n",
        "plt.ylim([count + 0.5, -2])\n",
        "\n",
        "plt.title(\"Cosine similarity between text and image features\", size=20)"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Text(0.5, 1.0, 'Cosine similarity between text and image features')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 613800 characters]\n",
            "text/plain": [
              "<Figure size 1440x1008 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "image/png": {
              "width": 1038,
              "height": 796
            },
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alePijoXy6AH"
      },
      "source": [
        "# Zero-Shot Image Classification\n",
        "\n",
        "You can classify images using the cosine similarity (times 100) as the logits to the softmax operation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nqu4GlfPfr-p",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102,
          "referenced_widgets": [
            "1369964d45004b5e95a058910b2a33e6",
            "12e23e2819094ee0a079d4eb77cfc4f9",
            "7a5f52e56ede4ac3abe37a3ece007dc9",
            "ce8b0faa1a1340b5a504d7b3546b3ccb",
            "5e6adc4592124a4581b85f4c1f3bab4d",
            "4a61c10fc00c4f04bb00b82e942da210",
            "b597cd6f6cd443aba4bf4491ac7f957e",
            "161969cae25a49f38aacd1568d3cac6c"
          ]
        },
        "outputId": "ca7a0e3c-e267-4e6e-8a1b-bbab3c0a2462"
      },
      "source": [
        "from torchvision.datasets import CIFAR100\n",
        "\n",
        "cifar100 = CIFAR100(os.path.expanduser(\"~/.cache\"), transform=preprocess, download=True)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz to /root/.cache/cifar-100-python.tar.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1369964d45004b5e95a058910b2a33e6",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=169001437.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Extracting /root/.cache/cifar-100-python.tar.gz to /root/.cache\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C4S__zCGy2MT"
      },
      "source": [
        "text_descriptions = [f\"This is a photo of a {label}\" for label in cifar100.classes]\n",
        "text_tokens = clip.tokenize(text_descriptions).cuda()"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c4z1fm9vCpSR"
      },
      "source": [
        "with torch.no_grad():\n",
        "    text_features = model.encode_text(text_tokens).float()\n",
        "    text_features /= text_features.norm(dim=-1, keepdim=True)\n",
        "\n",
        "text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)\n",
        "top_probs, top_labels = text_probs.cpu().topk(5, dim=-1)"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 931
        },
        "id": "T6Ju_6IBE2Iz",
        "outputId": "e1a155dc-474d-409c-e03d-d41b804648c3"
      },
      "source": [
        "plt.figure(figsize=(16, 16))\n",
        "\n",
        "for i, image in enumerate(original_images):\n",
        "    plt.subplot(4, 4, 2 * i + 1)\n",
        "    plt.imshow(image)\n",
        "    plt.axis(\"off\")\n",
        "\n",
        "    plt.subplot(4, 4, 2 * i + 2)\n",
        "    y = np.arange(top_probs.shape[-1])\n",
        "    plt.grid()\n",
        "    plt.barh(y, top_probs[i])\n",
        "    plt.gca().invert_yaxis()\n",
        "    plt.gca().set_axisbelow(True)\n",
        "    plt.yticks(y, [cifar100.classes[index] for index in top_labels[i].numpy()])\n",
        "    plt.xlabel(\"probability\")\n",
        "\n",
        "plt.subplots_adjust(wspace=0.5)\n",
        "plt.show()"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 1396212 characters]\n",
            "text/plain": [
              "<Figure size 1152x1152 with 16 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "image/png": {
              "width": 917,
              "height": 914
            },
            "needs_background": "light"
          }
        }
      ]
    }
  ]
}
