Image_Search_via_Text.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "dr50u1K6ePKY"
      },
      "source": [
        "# Search Image from Text via CLIP model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bVyxLp9UePKh"
      },
      "source": [
        "In this tutorial, we will create an image search system that retrieves images based on short text descriptions as query.\n",
        "\n",
        "The interest behind this is that in regular search, image description or meta data describing the content of the image needs to be indexed first before retrieving the images via text query. This can be expensive because you need a person to write that description and also information about image content is not always available.\n",
        "\n",
        "We need to look for another solution! What if we can directly compare text with images?\n",
        "\n",
        "To do so, we need to figure out a way to match images and text. One way is finding related images with similar semantics to the query text. This requires us to represent both images and query text in the same embedding space to be able to do the matching. In this case, pre-trained cross-modal models can help us out.\n",
        "\n",
        "For example when we write the word \"dog\" in query we want to be able to retrieve pictures with a dog solely by using the embeddings similarity.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PTrfq5IwePKo"
      },
      "source": [
        "In this tutorial we will guide you to create an Image from Text search application with Jina"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IPRsDqodzMsl",
        "tags": []
      },
      "source": [
        "## ⏰ Installing & Importing Dependencies\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hdk-Jw3HePKx"
      },
      "source": [
        "We will start this tutorial by installing the necessary ***pip*** dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qY1LTZTtX8vW"
      },
      "outputs": [],
      "source": [
        "!pip install Pillow jina transformers>=4.9.1 matplotlib torch>=1.9.0 torchvision>=0.10.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ylw5_HiscFYM"
      },
      "outputs": [],
      "source": [
        "! rm -rf workspace images query data*.zip* SimpleIndexer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jol2xKemzf3r"
      },
      "source": [
        "we will download the data to follow the tutorials"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uxtgtYD9zTB4",
        "tags": []
      },
      "outputs": [],
      "source": [
        "! wget https://open-images.s3.eu-central-1.amazonaws.com/data.zip\n",
        "! unzip data.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWIRxABRePLD"
      },
      "source": [
        "We will import the necessary dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "573TXM_scFa7"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from typing import Dict, Optional, Sequence, Tuple"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eIr_O1-m0UtS"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from jina import Executor, Flow, requests\n",
        "from jina.logging.logger import JinaLogger\n",
        "from jina.types.request import Request"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BY7imDNN1C82"
      },
      "outputs": [],
      "source": [
        "from docarray import Document, DocumentArray"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "guw6sCl30_13"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "WuMuUpSMePLV"
      },
      "source": [
        "## Defining the Indexer Flow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "auskQXHGePLX"
      },
      "source": [
        "Now that we understand the problem and we that have an idea on how to fix it, let's try to imagine what the solution would look like: \n",
        "\n",
        "1. We have a bunch of images with no text description about the content.\n",
        "2. We use a model to create an embedding that represents those images. \n",
        "3. Now we will index and save our embeddings which we will call Documents inside a workspace folder. \n",
        "\n",
        "This is what we will call the index Flow and we will show you how to create it with Jina\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "BTzU9r9XePLZ"
      },
      "source": [
        "### 1. Loading images into a DocumentArray"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ak-Ah31VePLd"
      },
      "outputs": [],
      "source": [
        "images = DocumentArray.from_files(f\"images/*.jpg\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ULUFGRUePLg"
      },
      "source": [
        "Here the `DocumentArray` images only contains filenames of the images, we will see later tjat it is the `Executor` that will load the image from disk"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s4zePJZr0B6e",
        "tags": []
      },
      "source": [
        "### 2. CLIPImageEncoder\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Of6SAHcLePLk"
      },
      "source": [
        "We need to add to our Flow (our pipeline) an encoder for the image, so we are going to create an `Executor`. (see details documentation [here](https://docs.jina.ai/fundamentals/executor/https://docs.jina.ai/fundamentals/executor/))\n",
        "\n",
        "This encoder encodes an image into embeddings using the CLIP model. \n",
        "We want an executor that loads the CLIP model and encodes it during the index flow. \n",
        "\n",
        "Our executor should:\n",
        "* support both **GPU** and **CPU**: That's why we will provision the `device` parameter and use it when encoding.\n",
        "* be able to process documents in batches in order to use our resources effectively: To do so, we will use the \n",
        "parameter `batch_size`This encoder encodes an image into embeddings using the CLIP model. \n",
        "We want an executor that loads the CLIP model and encodes it during the index flow. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CvEo335FePLm"
      },
      "outputs": [],
      "source": [
        "class CLIPImageEncoder(Executor):\n",
        "    \"\"\"Encode image into embeddings using the CLIP model.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        pretrained_model_name: str = \"openai/clip-vit-base-patch32\",\n",
        "        device: str = \"cpu\",\n",
        "        batch_size: int = 32,\n",
        "        *args,\n",
        "        **kwargs,\n",
        "    ):\n",
        "        super().__init__(*args, **kwargs)\n",
        "        self.batch_size = batch_size\n",
        "\n",
        "        self.device = device\n",
        "        self.preprocessor = CLIPFeatureExtractor.from_pretrained(pretrained_model_name)\n",
        "        self.model = CLIPModel.from_pretrained(\n",
        "            pretrained_model_name\n",
        "        )  # load the pretrained clip model from the transformer library\n",
        "\n",
        "        self.model.to(\n",
        "            self.device\n",
        "        ).eval()  # we want to do only inference so we put the model in eval mode\n",
        "\n",
        "    @requests\n",
        "    @torch.inference_mode()  # we don't want to keep track of the gradient during inference\n",
        "    def encode(self, docs: DocumentArray, parameters: dict, **kwargs):\n",
        "\n",
        "        for batch_docs in docs.batch(\n",
        "            batch_size=self.batch_size\n",
        "        ):  # we want to compute the embedding by batch of size batch_size\n",
        "            tensor = self._generate_input_features(\n",
        "                batch_docs\n",
        "            )  # Transformation from raw images to torch tensor\n",
        "            batch_docs.embeddings = (\n",
        "                self.model.get_image_features(**tensor).cpu().numpy()\n",
        "            )  # we compute the embeddings and store it directly in the DocumentArray\n",
        "\n",
        "    def _generate_input_features(self, docs: DocumentArray):\n",
        "        docs.apply(lambda d: d.load_uri_to_image_tensor())\n",
        "        input_features = self.preprocessor(\n",
        "            images=[d.tensor for d in docs],\n",
        "            return_tensors=\"pt\",\n",
        "        )\n",
        "        input_features = {\n",
        "            k: v.to(torch.device(self.device)) for k, v in input_features.items()\n",
        "        }\n",
        "        return input_features"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "1HoiXrgDePLq"
      },
      "source": [
        "### 3. SimpleIndexer\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RbyGpjB-ePLs"
      },
      "source": [
        " Now we want to index and save our embeddings to later perform be able to search within the document.\n",
        "Here we will implement a SimpleIndexer which will store the embedding on disk using the `SQLite` support as a backend in `docarray` ( see further details [here](https://docarray.jina.ai/advanced/document-store/sqlite/?highlight=sqlite)). This indexer is also available on the hub [here](https://hub.jina.ai/executor/zb38xlt4) but we prefer showing you how to create your own one in this tutorial.\n",
        "\n",
        "This indexer expose two endpoints : `/search` and `/index`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7EDrY0CzePMG"
      },
      "outputs": [],
      "source": [
        "class SimpleIndexer(Executor):\n",
        "    def __init__(self, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        print(os.path.join(self.workspace, 'index.db'))\n",
        "        self._index = DocumentArray(\n",
        "            storage='sqlite',\n",
        "            config={\n",
        "                'connection': os.path.join(self.workspace, 'index.db'),\n",
        "                'table_name': 'clip',\n",
        "            },\n",
        "        )\n",
        "\n",
        "    @requests(on='/index')\n",
        "    def index(self, docs: DocumentArray, **kwargs):\n",
        "        self._index.extend(docs)\n",
        "\n",
        "    @requests(on='/search')\n",
        "    def search(self, docs: DocumentArray, **kwargs):\n",
        "        docs.match(self._index)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KxA6exq80x0a",
        "tags": []
      },
      "source": [
        "### 4. Building the index Flow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UpEU1GQDePMJ"
      },
      "source": [
        "Here we compose our first `Flow` which will be in charge of indexing the images in the database"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yAKWrrSOhCDa",
        "tags": []
      },
      "outputs": [],
      "source": [
        "flow_index = (\n",
        "    Flow()\n",
        "    .add(uses=CLIPImageEncoder, name='encoder', uses_with={'device': \"cpu\"})\n",
        "    .add(uses=SimpleIndexer, name='indexer', workspace='workspace')\n",
        ")\n",
        "flow_index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nz-uo6QgePMN"
      },
      "source": [
        "As you can see on these image, each `DocumentArray` will first be encode the indexer in the database"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wk2f3oWVePMO"
      },
      "source": [
        "Now let's actually indexing the data by calling the FLow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KMSss7HPePMP"
      },
      "outputs": [],
      "source": [
        "with flow_index:\n",
        "    flow_index.post(on='/index', inputs=images)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5w4gBy8tePMR"
      },
      "source": [
        "The Flow has indexed the data in the database ! Now Let's define another FLow to query these images with some texts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "qlH3-9QHePMS"
      },
      "source": [
        "## Defining the Search Flow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6vhcEln5ePMU"
      },
      "source": [
        "Now to search for an image using text we do the following \n",
        "\n",
        "1. We embed the query text into the same embedding space as the image.\n",
        "2. We compute similarity between the query embedding and previously saved embeddings. \n",
        "3. We return the best results.\n",
        "\n",
        "This is our query Flow. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "X8GFdBRGePMW"
      },
      "source": [
        "### 1. CLIPTextEncoder\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XG_OMx26ePMY"
      },
      "source": [
        "This part is very similar to the CLIPImageEncoder, however instead of using the clip model to embeds images we are going to embed text. So code changed are very little, mainly using a tokenizer instead of the image preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_dvyy1iKePMZ"
      },
      "outputs": [],
      "source": [
        "class CLIPTextEncoder(Executor):\n",
        "    \"\"\"Encode text into embeddings using the CLIP model.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        encode_text=True,\n",
        "        pretrained_model_name: str = 'openai/clip-vit-base-patch32',\n",
        "        device: str = 'cpu',\n",
        "        batch_size: int = 32,\n",
        "        *args,\n",
        "        **kwargs,\n",
        "    ):\n",
        "        super().__init__(*args, **kwargs)\n",
        "        self.batch_size = batch_size\n",
        "        self.device = device\n",
        "\n",
        "        self.tokenizer = CLIPTokenizer.from_pretrained(\n",
        "            pretrained_model_name\n",
        "        )  # load the tokenizer from the transformer library\n",
        "\n",
        "        self.model = CLIPModel.from_pretrained(\n",
        "            pretrained_model_name\n",
        "        )  # load the pretrained clip model from the transformer library\n",
        "\n",
        "        self.model.eval().to(\n",
        "            device\n",
        "        )  # we want to do only inference so we put the model in eval mode\n",
        "\n",
        "    @requests\n",
        "    @torch.inference_mode()  # we don't want to keep track of the gradient during inference\n",
        "    def encode(self, docs: Optional[DocumentArray], parameters: Dict, **kwargs):\n",
        "\n",
        "        for docs_batch in docs.batch(\n",
        "            batch_size=self.batch_size\n",
        "        ):  # we want to compute the embedding by batch of size batch_size\n",
        "            input_tokens = self._generate_input_tokens(\n",
        "                docs_batch.texts\n",
        "            )  # Transformation from raw texts to torch tensor\n",
        "            docs_batch.embeddings = (\n",
        "                self.model.get_text_features(**input_tokens).cpu().numpy()\n",
        "            )  # we compute the embeddings and store it directly in the DocumentArray\n",
        "\n",
        "    def _generate_input_tokens(self, texts: Sequence[str]):\n",
        "\n",
        "        input_tokens = self.tokenizer(\n",
        "            texts,\n",
        "            max_length=77,\n",
        "            padding='longest',\n",
        "            truncation=True,\n",
        "            return_tensors='pt',\n",
        "        )\n",
        "        input_tokens = {k: v.to(self.device) for k, v in input_tokens.items()}\n",
        "        return input_tokens"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "0DecI7M0ePMb"
      },
      "source": [
        "### 2. Compute similarity between the query embedding and previously saved embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Okz33EQ0ePMd"
      },
      "source": [
        "This part will be done by the SimpleInexer that we define above. Let's take a deep dive on how this indexer actually perform the search part"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AbKWVUYOePMe"
      },
      "outputs": [],
      "source": [
        "class DumbSimpleIndexer(\n",
        "    Executor\n",
        "):  # this executor is here for example and is never used\n",
        "    @requests(on='/search')\n",
        "    def search(self, docs: DocumentArray, **kwargs):\n",
        "        docs.match(self._index)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X_T9rrYkePMh"
      },
      "source": [
        "The indexer just call the inner [match method](https://docarray.jina.ai/api/docarray.array.match/?highlight=match#module-docarray.array.match) from `DocumentArray` to perform a cosine similarity search betweem the embedding of the query and the embeding of the indexed images. The idea is that we search for the closests vectors of our query in the semantic space define by clip to find the images which correspond the most to our text query"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "urfNlu1SePMk"
      },
      "source": [
        "### 3. Define how to return the best results\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-IwkBst7ePMl"
      },
      "source": [
        "Here we will define a simple function that will plot the 3 closest image for each query texts. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CINQu_INePMm"
      },
      "outputs": [],
      "source": [
        "def plot_search_results(resp: Request):\n",
        "    for doc in resp.docs:\n",
        "        print(f'Query text: {doc.text}')\n",
        "        print(f'Matches:')\n",
        "        print('-' * 10)\n",
        "        doc.matches[:3].plot_image_sprites()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1GKKn19UePMn"
      },
      "source": [
        "This function will be called as a callback at then end of the Search Flow that we will now define"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "FhtMHPNSePMp"
      },
      "source": [
        "### 4. Building the index Flow\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KMHgJ2jYePMq"
      },
      "source": [
        "let's define the search flow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MUEfi11phPsQ",
        "tags": []
      },
      "outputs": [],
      "source": [
        "flow_search = (\n",
        "    Flow()\n",
        "    .add(uses=CLIPTextEncoder, name='encoder', uses_with={'device': \"cpu\"})\n",
        "    .add(uses=SimpleIndexer, name='indexer', workspace='workspace')\n",
        ")\n",
        "flow_search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BVMwgBV_ePMt"
      },
      "source": [
        "And let's query it"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IE-ZkMn6i-YH"
      },
      "outputs": [],
      "source": [
        "with flow_search:\n",
        "    resp = flow_search.search(\n",
        "        inputs=DocumentArray(\n",
        "            [\n",
        "                Document(text='dog'),\n",
        "                Document(text='cat'),\n",
        "                Document(text='kids on their bikes'),\n",
        "            ]\n",
        "        ),\n",
        "        on_done=plot_search_results,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sSc0B_wEePMz"
      },
      "source": [
        "Here you see that we can recover image of cat,dog or even Kigs on their bike easily !"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BsHuSa-oe9m9"
      },
      "outputs": [],
      "source": [
        "# clean up\n",
        "! rm -rf workspace images query\n",
        "! rm data.zip"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Copy of Image Search via Text.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}