{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "resident-liberal",
   "metadata": {},
   "source": [
    "# Joint Image & Text Embeddings\n",
    "\n",
    "This example shows how SentenceTransformer 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",
    "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": 1,
   "id": "indonesian-yemen",
   "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",
    "torch.set_num_threads(4)\n",
    "\n",
    "\n",
    "\n",
    "#First, we load the respective CLIP model\n",
    "model = SentenceTransformer('clip-ViT-B-32')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "southeast-entity",
   "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": "static-insured",
   "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",
    "    img_names = list(glob.glob('unsplash/photos/*.jpg'))\n",
    "    print(\"Images:\", len(img_names))\n",
    "    img_emb = 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": 11,
   "id": "operating-southeast",
   "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": "southeast-pound",
   "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": [
      "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": [
      "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": [
      "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": 6,
   "id": "tender-blanket",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'A sunset on the beach'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nfQ2QGhdQNc.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 105464 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "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": [
      "vlukOqxOA8o.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 71368 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "search(\"A sunset on the beach\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "genuine-vegetarian",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'London'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UnGcRvTyJOo.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 113312 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S3G8qX4Ft5s.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 161828 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t05kfHeygbE.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 145660 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "search(\"London\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "nuclear-lexington",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'A dog in a park'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IVyZrLp41D0.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 85324 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": [
      "KVeogBZzl4M.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 92616 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "search(\"A dog in a park\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "divine-omaha",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'A beach with palm trees'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7rrgPPljqYU.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 85004 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kmihWgpbDEg.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 162304 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ZyfOq52b0cs.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 90608 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "search(\"A beach with palm trees\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fewer-absorption",
   "metadata": {},
   "source": [
    "# Image-to-Image Search\n",
    "You can use the method also for image-to-image search.\n",
    "\n",
    "To achieve this, you pass `Image.open('path/to/image.jpg')` to the search method.\n",
    "\n",
    "It will then return similar images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "prime-webcam",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 628892 characters]\n",
      "text/plain": [
       "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=640x800 at 0x7F624B577CD0>"
      ]
     },
     "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": [
      "8rmYDezMIE4.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 125640 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mMj5ykKvwAk.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 88324 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nJW1RzBXZcI.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 116964 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Xmo004TUzrI.jpg\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 101340 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "search(Image.open(os.path.join(img_folder, 'lyStEjlKNSw.jpg')), k=5)"
   ]
  }
 ],
 "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
}