{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/rwightman/12a6526daf2e57f994295b2cce367ac3/interacting-with-clip.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPHN7PJgKOzb"
      },
      "source": [
        "# Interacting with open_clip\n",
        "\n",
        "This is a self-contained notebook that shows how to download and run open_clip models, calculate the similarity between arbitrary image and text inputs, and perform zero-shot image classifications."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z7F87xdw9rBX"
      },
      "source": [
        "## Preparation for colab"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "ck9qjK8z9rBX"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0BpdJkdBssk9",
        "outputId": "514354f7-3d50-4ef9-bbbe-88ecbc382cd9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: open_clip_torch in /usr/local/lib/python3.9/dist-packages (2.16.2)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.9/dist-packages (3.7.1)\n",
            "Requirement already satisfied: torch>=1.9.0 in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (2.0.0+cu118)\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (6.1.1)\n",
            "Requirement already satisfied: timm in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (0.6.13)\n",
            "Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (0.13.4)\n",
            "Requirement already satisfied: protobuf<4 in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (3.20.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (4.65.0)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (2022.10.31)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (0.1.98)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.9/dist-packages (from open_clip_torch) (0.15.1+cu118)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (23.0)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (4.39.3)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (2.8.2)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (1.4.4)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (8.4.0)\n",
            "Requirement already satisfied: importlib-resources>=3.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (5.12.0)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (3.0.9)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (1.0.7)\n",
            "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (1.22.4)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (0.11.0)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.9/dist-packages (from importlib-resources>=3.2.0->matplotlib) (3.15.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.9/dist-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch>=1.9.0->open_clip_torch) (4.5.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.9/dist-packages (from torch>=1.9.0->open_clip_torch) (3.1.2)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.9/dist-packages (from torch>=1.9.0->open_clip_torch) (1.11.1)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.9/dist-packages (from torch>=1.9.0->open_clip_torch) (2.0.0)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.9/dist-packages (from torch>=1.9.0->open_clip_torch) (3.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from torch>=1.9.0->open_clip_torch) (3.11.0)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch>=1.9.0->open_clip_torch) (16.0.1)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch>=1.9.0->open_clip_torch) (3.25.2)\n",
            "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.9/dist-packages (from ftfy->open_clip_torch) (0.2.6)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from huggingface-hub->open_clip_torch) (2.27.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub->open_clip_torch) (6.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.9/dist-packages (from jinja2->torch>=1.9.0->open_clip_torch) (2.1.2)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests->huggingface-hub->open_clip_torch) (2022.12.7)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->huggingface-hub->open_clip_torch) (1.26.15)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.9/dist-packages (from requests->huggingface-hub->open_clip_torch) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests->huggingface-hub->open_clip_torch) (3.4)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.9/dist-packages (from sympy->torch>=1.9.0->open_clip_torch) (1.3.0)\n"
          ]
        }
      ],
      "source": [
        "! pip install open_clip_torch matplotlib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "C1hkDT38hSaP"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import torch"
      ]
    },
    {
      "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",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uLFS29hnhlY4",
        "outputId": "c1cb8995-d525-4f81-fd96-eaba02a2e48e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[('RN50', 'openai'),\n",
              " ('RN50', 'yfcc15m'),\n",
              " ('RN50', 'cc12m'),\n",
              " ('RN50-quickgelu', 'openai'),\n",
              " ('RN50-quickgelu', 'yfcc15m'),\n",
              " ('RN50-quickgelu', 'cc12m'),\n",
              " ('RN101', 'openai'),\n",
              " ('RN101', 'yfcc15m'),\n",
              " ('RN101-quickgelu', 'openai'),\n",
              " ('RN101-quickgelu', 'yfcc15m'),\n",
              " ('RN50x4', 'openai'),\n",
              " ('RN50x16', 'openai'),\n",
              " ('RN50x64', 'openai'),\n",
              " ('ViT-B-32', 'openai'),\n",
              " ('ViT-B-32', 'laion400m_e31'),\n",
              " ('ViT-B-32', 'laion400m_e32'),\n",
              " ('ViT-B-32', 'laion2b_e16'),\n",
              " ('ViT-B-32', 'laion2b_s34b_b79k'),\n",
              " ('ViT-B-32-quickgelu', 'openai'),\n",
              " ('ViT-B-32-quickgelu', 'laion400m_e31'),\n",
              " ('ViT-B-32-quickgelu', 'laion400m_e32'),\n",
              " ('ViT-B-16', 'openai'),\n",
              " ('ViT-B-16', 'laion400m_e31'),\n",
              " ('ViT-B-16', 'laion400m_e32'),\n",
              " ('ViT-B-16', 'laion2b_s34b_b88k'),\n",
              " ('ViT-B-16-plus-240', 'laion400m_e31'),\n",
              " ('ViT-B-16-plus-240', 'laion400m_e32'),\n",
              " ('ViT-L-14', 'openai'),\n",
              " ('ViT-L-14', 'laion400m_e31'),\n",
              " ('ViT-L-14', 'laion400m_e32'),\n",
              " ('ViT-L-14', 'laion2b_s32b_b82k'),\n",
              " ('ViT-L-14-336', 'openai'),\n",
              " ('ViT-H-14', 'laion2b_s32b_b79k'),\n",
              " ('ViT-g-14', 'laion2b_s12b_b42k'),\n",
              " ('ViT-g-14', 'laion2b_s34b_b88k'),\n",
              " ('ViT-bigG-14', 'laion2b_s39b_b160k'),\n",
              " ('roberta-ViT-B-32', 'laion2b_s12b_b32k'),\n",
              " ('xlm-roberta-base-ViT-B-32', 'laion5b_s13b_b90k'),\n",
              " ('xlm-roberta-large-ViT-H-14', 'frozen_laion5b_s13b_b90k'),\n",
              " ('convnext_base', 'laion400m_s13b_b51k'),\n",
              " ('convnext_base_w', 'laion2b_s13b_b82k'),\n",
              " ('convnext_base_w', 'laion2b_s13b_b82k_augreg'),\n",
              " ('convnext_base_w', 'laion_aesthetic_s13b_b82k'),\n",
              " ('convnext_base_w_320', 'laion_aesthetic_s13b_b82k'),\n",
              " ('convnext_base_w_320', 'laion_aesthetic_s13b_b82k_augreg'),\n",
              " ('convnext_large_d', 'laion2b_s26b_b102k_augreg'),\n",
              " ('convnext_large_d_320', 'laion2b_s29b_b131k_ft'),\n",
              " ('convnext_large_d_320', 'laion2b_s29b_b131k_ft_soup'),\n",
              " ('convnext_xxlarge', 'laion2b_s34b_b82k_augreg'),\n",
              " ('convnext_xxlarge', 'laion2b_s34b_b82k_augreg_rewind'),\n",
              " ('convnext_xxlarge', 'laion2b_s34b_b82k_augreg_soup'),\n",
              " ('coca_ViT-B-32', 'laion2b_s13b_b90k'),\n",
              " ('coca_ViT-B-32', 'mscoco_finetuned_laion2b_s13b_b90k'),\n",
              " ('coca_ViT-L-14', 'laion2b_s13b_b90k'),\n",
              " ('coca_ViT-L-14', 'mscoco_finetuned_laion2b_s13b_b90k')]"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "import open_clip\n",
        "open_clip.list_pretrained()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "8pAHe7nD9rBa"
      },
      "outputs": [],
      "source": [
        "model, _, preprocess = open_clip.create_model_and_transforms('convnext_base_w', pretrained='laion2b_s13b_b82k_augreg')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IBRVTY9lbGm8",
        "outputId": "7fc2f6c9-d7c1-4fe4-fda9-dfd372d03834"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model parameters: 179,385,345\n",
            "Context length: 77\n",
            "Vocab size: 49408\n"
          ]
        }
      ],
      "source": [
        "model.eval()\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(\"Context length:\", context_length)\n",
        "print(\"Vocab size:\", vocab_size)"
      ]
    },
    {
      "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",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d6cpiIFHp9N6",
        "outputId": "84943ad9-0b1a-4200-946d-a8fb164ec0b0"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Compose(\n",
              "    Resize(size=256, interpolation=bicubic, max_size=None, antialias=warn)\n",
              "    CenterCrop(size=(256, 256))\n",
              "    <function _convert_to_rgb at 0x7fac58819c10>\n",
              "    ToTensor()\n",
              "    Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ],
      "source": [
        "preprocess"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xwSB5jZki3Cj"
      },
      "source": [
        "# Text Preprocessing\n",
        "\n",
        "We use a case-insensitive tokenizer, which can be invoked using `tokenizer.tokenize()`. By default, the outputs are padded to become 77 tokens long, which is what the CLIP models expects."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "8f9coA_O9rBb"
      },
      "outputs": [],
      "source": [
        "from open_clip import tokenizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qGom156-i2kL",
        "outputId": "f4b4f532-f8ea-424f-d878-8cd014dad5e4"
      },
      "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": {},
          "execution_count": 28
        }
      ],
      "source": [
        "tokenizer.tokenize(\"Hello World!\")"
      ]
    },
    {
      "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",
      "execution_count": 29,
      "metadata": {
        "id": "tMc1AXzBlhzm"
      },
      "outputs": [],
      "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",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 504
        },
        "id": "NSSrLY185jSf",
        "outputId": "d057493e-01e1-4673-b549-0cb8a0a6b2d2"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1600x500 with 8 Axes>"
            ],
            "image/png": "[encoded data, 2938004 characters]\n"
          },
          "metadata": {
            "image/png": {
              "width": 1547,
              "height": 487
            }
          }
        }
      ],
      "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"
      ]
    },
    {
      "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",
      "execution_count": 31,
      "metadata": {
        "id": "HBgCanxi8JKw"
      },
      "outputs": [],
      "source": [
        "image_input = torch.tensor(np.stack(images))\n",
        "text_tokens = tokenizer.tokenize([\"This is \" + desc for desc in texts])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "ZN9I0nIBZ_vW"
      },
      "outputs": [],
      "source": [
        "with torch.no_grad():\n",
        "    image_features = model.encode_image(image_input).float()\n",
        "    text_features = model.encode_text(text_tokens).float()"
      ]
    },
    {
      "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",
      "execution_count": 33,
      "metadata": {
        "id": "yKAxkQR7bf3A"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "C5zvMxh8cU6m",
        "outputId": "dc68e0ed-660b-4a65-82ff-5ae173731497"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Text(0.5, 1.0, 'Cosine similarity between text and image features')"
            ]
          },
          "metadata": {},
          "execution_count": 34
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 2000x1400 with 1 Axes>"
            ],
            "image/png": "[encoded data, 1141928 characters]\n"
          },
          "metadata": {
            "image/png": {
              "width": 1459,
              "height": 1127
            }
          }
        }
      ],
      "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)"
      ]
    },
    {
      "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",
      "execution_count": 35,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Nqu4GlfPfr-p",
        "outputId": "644ddaf0-11c8-48fa-e81a-1488f7728c11"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Files already downloaded and verified\n"
          ]
        }
      ],
      "source": [
        "from torchvision.datasets import CIFAR100\n",
        "\n",
        "cifar100 = CIFAR100(os.path.expanduser(\"~/.cache\"), transform=preprocess, download=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "C4S__zCGy2MT"
      },
      "outputs": [],
      "source": [
        "text_descriptions = [f\"A photo of a {label}\" for label in cifar100.classes]\n",
        "text_tokens = tokenizer.tokenize(text_descriptions)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "c4z1fm9vCpSR"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "T6Ju_6IBE2Iz",
        "outputId": "4de01bf6-9ff5-4b74-d093-fb3bce251ffe"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1600x1600 with 16 Axes>"
            ],
            "image/png": "[encoded data, 2570696 characters]\n"
          },
          "metadata": {
            "image/png": {
              "width": 1264,
              "height": 1295
            }
          }
        }
      ],
      "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()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Interacting with CLIP.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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": 0
}