{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fantastic-grave",
   "metadata": {},
   "source": [
    "# Zero-Shot Image Classification\n",
    "\n",
    "This example shows how [SentenceTransformers](https://www.sbert.net) can be used to map images and texts to the same vector space. \n",
    "\n",
    "We can use this to perform **zero-shot image classification** by providing the names for the labels.\n",
    "\n",
    "As model, we use the [OpenAI CLIP Model](https://github.com/openai/CLIP), which was trained on a large set of images and image alt texts.\n",
    "\n",
    "\n",
    "The images in this example are from [Unsplash](https://unsplash.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "functioning-diesel",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "from PIL import Image\n",
    "import glob\n",
    "import torch\n",
    "import pickle\n",
    "import zipfile\n",
    "from IPython.display import display\n",
    "from IPython.display import Image as IPImage\n",
    "import os\n",
    "from tqdm.autonotebook import tqdm\n",
    "import torch\n",
    "\n",
    "# We use the original CLIP model for computing image embeddings and English text embeddings\n",
    "en_model = SentenceTransformer('clip-ViT-B-32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "revolutionary-recording",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We download some images from our repository which we want to classify\n",
    "img_names = ['eiffel-tower-day.jpg', 'eiffel-tower-night.jpg', 'two_dogs_in_snow.jpg', 'cat.jpg']\n",
    "url = 'https://github.com/UKPLab/sentence-transformers/raw/master/examples/applications/image-search/'\n",
    "for filename in img_names:\n",
    "    if not os.path.exists(filename):\n",
    "        util.http_get(url+filename, filename)\n",
    "\n",
    "# And compute the embeddings for these images\n",
    "img_emb = en_model.encode([Image.open(filepath) for filepath in img_names], convert_to_tensor=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "preceding-hayes",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 129508 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: Paris\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 161084 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: Paris at night\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 104144 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: dog\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 52904 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: cat\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Then, we define our labels as text. Here, we use 4 labels\n",
    "labels = ['dog', 'cat', 'Paris at night', 'Paris']\n",
    "\n",
    "# And compute the text embeddings for these labels\n",
    "en_emb = en_model.encode(labels, convert_to_tensor=True)\n",
    "\n",
    "# Now, we compute the cosine similarity between the images and the labels\n",
    "cos_scores = util.cos_sim(img_emb, en_emb)\n",
    "\n",
    "# Then we look which label has the highest cosine similarity with the given images\n",
    "pred_labels = torch.argmax(cos_scores, dim=1)\n",
    "\n",
    "# Finally we output the images + labels\n",
    "for img_name, pred_label in zip(img_names, pred_labels):\n",
    "    display(IPImage(img_name, width=200))\n",
    "    print(\"Predicted label:\", labels[pred_label])\n",
    "    print(\"\\n\\n\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "generic-cooperative",
   "metadata": {},
   "source": [
    "# Zero-Shot Image Classification\n",
    "The original CLIP Model only works for English, hence, we used [Multilingual Knowlegde Distillation](https://arxiv.org/abs/2004.09813) to make this model work with 50+ languages.\n",
    "\n",
    "For this, we msut load the *clip-ViT-B-32-multilingual-v1* model to encode our labels.\n",
    "We can define our labels in 50+ languages and can also mix the languages we have"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "cellular-network",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 129508 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: Париж\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 161084 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: 巴黎晚上\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 104144 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: Hund\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 52904 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted label: gato\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "multi_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')\n",
    "\n",
    "# Then, we define our labels as text. Here, we use 4 labels\n",
    "labels = ['Hund',     # German: dog\n",
    "          'gato',     # Spanish: cat \n",
    "          '巴黎晚上',  # Chinese: Paris at night\n",
    "          'Париж'     # Russian: Paris\n",
    "         ]\n",
    "\n",
    "# And compute the text embeddings for these labels\n",
    "txt_emb = multi_model.encode(labels, convert_to_tensor=True)\n",
    "\n",
    "# Now, we compute the cosine similarity between the images and the labels\n",
    "cos_scores = util.cos_sim(img_emb, txt_emb)\n",
    "\n",
    "# Then we look which label has the highest cosine similarity with the given images\n",
    "pred_labels = torch.argmax(cos_scores, dim=1)\n",
    "\n",
    "# Finally we output the images + labels\n",
    "for img_name, pred_label in zip(img_names, pred_labels):\n",
    "    display(IPImage(img_name, width=200))\n",
    "    print(\"Predicted label:\", labels[pred_label])\n",
    "    print(\"\\n\\n\")"
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}