{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "introductory-framework",
   "metadata": {},
   "source": [
    "# Multilingual Joint Image & Text Embeddings \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",
    "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",
    "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",
    "As a source for fotos, we use the [Unsplash Dataset Lite](https://unsplash.com/data), which contains about 25k images. See the [License](https://unsplash.com/license) about the Unsplash images. \n",
    "\n",
    "Note: 25k images is rather small. If you search for really specific terms, the chance are high that no such photo exist in the collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "copyrighted-harmony",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "936d49c346964954ac66ba63bd8858b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0.00/502M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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",
    "\n",
    "# Here we load the multilingual CLIP model. Note, this model can only encode text.\n",
    "# If you need embeddings for images, you must load the 'clip-ViT-B-32' model\n",
    "model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "strategic-mayor",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, we get about 25k images from Unsplash \n",
    "img_folder = 'photos/'\n",
    "if not os.path.exists(img_folder) or len(os.listdir(img_folder)) == 0:\n",
    "    os.makedirs(img_folder, exist_ok=True)\n",
    "    \n",
    "    photo_filename = 'unsplash-25k-photos.zip'\n",
    "    if not os.path.exists(photo_filename):   #Download dataset if does not exist\n",
    "        util.http_get('http://sbert.net/datasets/'+photo_filename, photo_filename)\n",
    "        \n",
    "    #Extract all images\n",
    "    with zipfile.ZipFile(photo_filename, 'r') as zf:\n",
    "        for member in tqdm(zf.infolist(), desc='Extracting'):\n",
    "            zf.extract(member, img_folder)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "adjacent-encyclopedia",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Images: 24996\n"
     ]
    }
   ],
   "source": [
    "# Now, we need to compute the embeddings\n",
    "# To speed things up, we destribute pre-computed embeddings\n",
    "# Otherwise you can also encode the images yourself.\n",
    "# To encode an image, you can use the following code:\n",
    "# from PIL import Image\n",
    "# img_emb = model.encode(Image.open(filepath))\n",
    "\n",
    "use_precomputed_embeddings = True\n",
    "\n",
    "if use_precomputed_embeddings: \n",
    "    emb_filename = 'unsplash-25k-photos-embeddings.pkl'\n",
    "    if not os.path.exists(emb_filename):   #Download dataset if does not exist\n",
    "        util.http_get('http://sbert.net/datasets/'+emb_filename, emb_filename)\n",
    "        \n",
    "    with open(emb_filename, 'rb') as fIn:\n",
    "        img_names, img_emb = pickle.load(fIn)  \n",
    "    print(\"Images:\", len(img_names))\n",
    "else:\n",
    "    #For embedding images, we need the non-multilingual CLIP model\n",
    "    img_model = SentenceTransformer('clip-ViT-B-32')\n",
    "\n",
    "    img_names = list(glob.glob('unsplash/photos/*.jpg'))\n",
    "    print(\"Images:\", len(img_names))\n",
    "    img_emb = img_model.encode([Image.open(filepath) for filepath in img_names], batch_size=128, convert_to_tensor=True, show_progress_bar=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dental-friendly",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, we define a search function.\n",
    "def search(query, k=3):\n",
    "    # First, we encode the query (which can either be an image or a text string)\n",
    "    query_emb = model.encode([query], convert_to_tensor=True, show_progress_bar=False)\n",
    "    \n",
    "    # Then, we use the util.semantic_search function, which computes the cosine-similarity\n",
    "    # between the query embedding and all image embeddings.\n",
    "    # It then returns the top_k highest ranked images, which we output\n",
    "    hits = util.semantic_search(query_emb, img_emb, top_k=k)[0]\n",
    "    \n",
    "    print(\"Query:\")\n",
    "    display(query)\n",
    "    for hit in hits:\n",
    "        print(img_names[hit['corpus_id']])\n",
    "        display(IPImage(os.path.join(img_folder, img_names[hit['corpus_id']]), width=200))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "established-emphasis",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Two dogs playing in the snow'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lyStEjlKNSw.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 107640 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FAcSe7SjDUU.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 59444 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hb6nGDgWztE.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 40916 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "search(\"Two dogs playing in the snow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "varying-feeling",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Eine Katze auf einem Stuhl'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CgGDzMYdYw8.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 163504 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kjERLXaHjXc.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 32620 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I-YJ-gaJNaw.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 176848 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#German: A cat on a chair\n",
    "search(\"Eine Katze auf einem Stuhl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "union-league",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Muchos peces'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "H22jcGTyrS4.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 112668 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CJ_9I6aXSnc.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 44556 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_MJKaRig1Ic.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 248740 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Spanish: Many fish\n",
    "search(\"Muchos peces\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "artistic-reaction",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'棕榈树的沙滩'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crIXKhUDpBI.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 124812 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_6iV1AJZ53s.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 189316 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rv63du1a79E.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 202956 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Chinese: A beach with palm trees\n",
    "search(\"棕榈树的沙滩\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "drawn-feelings",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Закат на пляже'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JC5U3Eyiyr4.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 32552 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5z1QDcisnJ8.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 70572 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rdG4hRoyVR0.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 39116 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Russian: A sunset on the beach\n",
    "search(\"Закат на пляже\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "contrary-drink",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Parkta bir köpek'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ROJLfAbL1Ig.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 122080 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0O9A0F_d1qA.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 195512 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4mdsPUtN0P0.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 200648 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Turkish: A dog in a park\n",
    "search(\"Parkta bir köpek\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "statewide-symphony",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'夜のニューヨーク'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FGjR4IGwP7U.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 137200 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8nCMOFYyXF4.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 95208 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ZAOEjcpdMkc.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 188608 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Japanese: New York at night\n",
    "search(\"夜のニューヨーク\")"
   ]
  }
 ],
 "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
}
