src/notebooks/zeroshot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zeroshot with CLOOB\n",
    "\n",
    "This notebook shows how to do zeroshot with pretrained CLOOB models.\n",
    "\n",
    "The models and the datasets can be downloaded [here](https://ml.jku.at/research/CLOOB/downloads/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "src_path = os.path.split(os.getcwd())[0]\n",
    "sys.path.insert(0, src_path)\n",
    "\n",
    "import json\n",
    "import random\n",
    "import glob\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.datasets as datasets\n",
    "from tqdm import tqdm\n",
    "from sklearn.metrics import accuracy_score, balanced_accuracy_score\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import clip.clip as clip\n",
    "from clip.clip import _transform\n",
    "from clip.model import CLIPGeneral\n",
    "\n",
    "import training.zeroshot_data as zeroshot_data\n",
    "\n",
    "# Set the GPU\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load checkpoint\n",
    "Add the path of checkpoint here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = '<checkpoint-dir>/checkpoints/cloob_rn50_yfcc_epoch_28.pt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(checkpoint_path)\n",
    "model_config_file = os.path.join(src_path, 'training/model_configs/', checkpoint['model_config_file'])\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Device is \", device)\n",
    "\n",
    "print('Loading model from', model_config_file)\n",
    "assert os.path.exists(model_config_file)\n",
    "with open(model_config_file, 'r') as f:\n",
    "    model_info = json.load(f)\n",
    "model = CLIPGeneral(**model_info)\n",
    "preprocess= _transform(model.visual.input_resolution, is_train=False)\n",
    "\n",
    "if not torch.cuda.is_available():\n",
    "    model.float()\n",
    "else:\n",
    "    model.to(device)\n",
    "\n",
    "sd = checkpoint[\"state_dict\"]\n",
    "sd = {k[len('module.'):]: v for k, v in sd.items()}\n",
    "if 'logit_scale_hopfield' in sd:\n",
    "    sd.pop('logit_scale_hopfield', None)\n",
    "model.load_state_dict(sd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "Add the path of the dataset you want to do evaluate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '<dataset-dir>/zeroshot_datasets/birdsnap/test'\n",
    "dataset = datasets.ImageFolder(data_path, transform=preprocess)\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=256, num_workers=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Zeroshot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zero_shot_classifier(model, classnames, templates, device):\n",
    "    with torch.no_grad():\n",
    "        zeroshot_weights = []\n",
    "        for classname in tqdm(classnames):\n",
    "            texts = [template(classname) for template in templates] #format with class\n",
    "            texts = clip.tokenize(texts).to(device) #tokenize\n",
    "            class_embeddings = model.encode_text(texts)\n",
    "            class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)\n",
    "            class_embedding = class_embeddings.mean(dim=0)\n",
    "            class_embedding /= class_embedding.norm()\n",
    "            zeroshot_weights.append(class_embedding)\n",
    "        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)\n",
    "    return zeroshot_weights\n",
    "\n",
    "def run(model, classifier, dataloader, device, accuracy_metric):\n",
    "    with torch.no_grad():\n",
    "        all_logits = []\n",
    "        all_targets = []\n",
    "        for images, target in tqdm(dataloader):\n",
    "            images = images.to(device)\n",
    "            target = target.to(device)\n",
    "\n",
    "            # predict\n",
    "            image_features = model.encode_image(images)\n",
    "            image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "            logits = image_features @ classifier\n",
    "\n",
    "            all_logits.append(logits.cpu())\n",
    "            all_targets.append(target.cpu())\n",
    "\n",
    "        all_logits = torch.cat(all_logits).numpy()\n",
    "        all_targets = torch.cat(all_targets).numpy()\n",
    "\n",
    "        acc = accuracy_metric(all_targets, all_logits.argmax(axis=1)) * 100.0\n",
    "        return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the classnames and prompt_templates here. \n",
    "\n",
    "For the provided datasets the original CLIP classnames and prompts can be found in [zeroshot_data.py](../training/zeroshot_data.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classnames = zeroshot_data.birdsnap_classnames\n",
    "prompt_templates = zeroshot_data.birdsnap_templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "accuracy_metric = accuracy_score\n",
    "print(\"Calculating the text embeddings for all classes of the dataset\", flush=True)\n",
    "classifier = zero_shot_classifier(model, classnames, prompt_templates, device)\n",
    "print(\"Calculating the image embeddings for all images of the dataset\", flush=True)\n",
    "accuracy = run(model, classifier, dataloader, device, accuracy_score)\n",
    "print('Zeroshot accuracy: ', accuracy.round(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization of the results\n",
    "\n",
    "This visualization shows the top 5 classification results for random images in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_images(data_path, preprocess, count=8, seed=31415):\n",
    "    random.seed(seed)\n",
    "    class_folder_names = sorted(os.listdir(data_path))\n",
    "    original_images = []\n",
    "    processed_images = []\n",
    "    image_labels = []\n",
    "    all_image_paths = glob.glob(os.path.join(data_path, '*', '*'), recursive=True)\n",
    "    select_image_paths = random.sample(all_image_paths, k=count)\n",
    "    for image_path in select_image_paths:\n",
    "        image_labels.append(class_folder_names.index(os.path.split(os.path.split(image_path)[0])[1]))\n",
    "        with Image.open(image_path) as im:\n",
    "            original_images.append(im)\n",
    "            processed_images.append(preprocess(im))\n",
    "\n",
    "    processed_images = torch.stack(processed_images)\n",
    "\n",
    "    return original_images, processed_images, image_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_images, processed_images, image_labels = get_images(data_path, preprocess)\n",
    "model.eval()\n",
    "images = processed_images.to(device)\n",
    "with torch.no_grad():\n",
    "    image_features = model.encode_image(images)\n",
    "    image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "text_probs = (30.0 * image_features @ classifier).softmax(dim=-1)\n",
    "top_probs, top_labels = text_probs.cpu().topk(5, dim=-1)\n",
    "\n",
    "top_probs = top_probs.cpu()\n",
    "top_labels = top_labels.cpu()\n",
    "text_probs = text_probs.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(4,4,figsize=(20,10))\n",
    "for i in range(8):\n",
    "    image = original_images[i]\n",
    "    ax = axs[i%4, 0 + 2*int(i/4)]\n",
    "    ax.imshow(image)\n",
    "    title_tex = classnames[image_labels[i]]\n",
    "    ax.set_title(title_tex)\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "    correct_rank = (text_probs.argsort(dim=1, descending=True)[i] == image_labels[i]).nonzero(as_tuple=False).item() + 1\n",
    "\n",
    "    ax = axs[i%4, 1 + 2*int(i/4)]\n",
    "    title_tex = f'correct_rank: {correct_rank}/{len(classnames)}'\n",
    "    ax.set_title(title_tex)\n",
    "    y = np.arange(top_probs.shape[-1])\n",
    "    ax.barh(y, top_probs[i], zorder=-1, color=[123/255.0,204/255.0,196/255.0,255/255.0])\n",
    "    ax.invert_yaxis()\n",
    "    ax.set_axisbelow(True)\n",
    "    ax.set_yticks(y)\n",
    "    ax.set_xlim([0,1])\n",
    "    ax.yaxis.set_ticks_position('none') \n",
    "    ax.tick_params(axis=\"y\", labelsize=16)\n",
    "    classnames_plot = classnames\n",
    "    ax.set_yticklabels([classnames_plot[index] for index in top_labels[i].numpy()], x=0.05, zorder=1, horizontalalignment='left')\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(hspace=0.4, wspace=0.05)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "41231f278ebceac84b2c2f784f457cfc03b89a7637de3e5e7fcd045d158a53d5"
  },
  "kernelspec": {
   "display_name": "Python 3.6.13 64-bit ('open_clip': conda)",
   "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.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}