notebooks/02_build_images_vector_db.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Multimodal Image Search App with Titan Embeddings and LangChain\n",
    "\n",
    "This post demonstrates how to combine [Titan Multimodal Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-multiemb-models.html), [LangChain](https://python.langchain.com/docs/get_started/introduction) and [FAISS](https://python.langchain.com/docs/integrations/vectorstores/faiss/) to build a capable image search application. Titan's embeddings allow representing images and text in a common dense vector space, enabling natural language querying of images. FAISS provides a fast, scalable way to index and search those vectors. And LangChain offers abstractions to hook everything together and surface relevant image results based on a user's query.\n",
    "\n",
    "By following the steps outlined, you'll be able to preprocess images, generate embeddings, load them into FAISS, and write a simple application that takes in a natural language query, searches the FAISS index, and returns the most semantically relevant images. It's a great example of the power of combining modern AI technologies to build applications.\n",
    "\n",
    "![Diagram](data/build_images_vector_db.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "Before we begin, make sure you have the following prerequisites set up:\n",
    "\n",
    "1. **Install boto3:** This is the [AWS SDK for Python ](https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingTheBotoAPI.html)that allows interacting with AWS services. Install with `pip install boto3`.\n",
    "2. [Configure AWS credentials](https://docs.aws.amazon.com/braket/latest/developerguide/braket-using-boto3.html): Boto3 needs credentials to make API calls to AWS.\n",
    "3. **Amazon Titan Multimodal Embeddings:** Follow the instructions in the [Amazon Bedrock User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html) to access the Titan Multimodal Embeddings model.\n",
    "4. **LangChain:** Install the LangChain library by running `pip install langchain`.\n",
    "5. **FAISS:** Install the FAISS vector database by running `pip install faiss-cpu` faiss-cpu (or faiss-gpu if you have a GPU available). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://catalog.workshops.aws/building-with-amazon-bedrock/en-US/image-labs/bedrock-image-search\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!brew install postgresql\n",
    "#!pip install boto3\n",
    "#!pip install langchain\n",
    "#!pip install faiss-cpu\n",
    "#!pip install Pillow\n",
    "# or install requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import boto3\n",
    "import json\n",
    "import base64\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_aws import BedrockEmbeddings\n",
    "from io import BytesIO\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock_client              = boto3.client(\"bedrock-runtime\", region_name='us-east-1') \n",
    "boto3_bedrock               = boto3.client('bedrock', region_name='us-east-1')\n",
    "default_embedding_dimmesion = os.environ.get(\"DEFAULT_EMBEDDING_DIMENSION\", \"1024\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validate the connection\n",
    "We can check the client works by trying out the list_foundation_models() method, which will tell us all the models available for us to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[models['modelId'] for models in boto3_bedrock.list_foundation_models()['modelSummaries']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Embeddings for Images\n",
    "The first step is to determine whether we will be processing text or images. We identify this using the `get_multimodal_vector` function which takes the input and utilizes the Amazon Titan model through the [InvokeModel](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html) API from [Amazon Bedrock](https://aws.amazon.com/bedrock/) to generate a joint embedding vector for the image or text, as applicable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calls Bedrock to get a vector from either an image, text, or both\n",
    "def get_multimodal_vector(input_image_base64=None, input_text=None):\n",
    "    \n",
    "    session = boto3.Session()\n",
    "\n",
    "    bedrock = session.client(service_name='bedrock-runtime', region_name='us-east-1') #creates a Bedrock client\n",
    "    \n",
    "    request_body = {}\n",
    "    \n",
    "    if input_text:\n",
    "        request_body[\"inputText\"] = input_text\n",
    "        \n",
    "    if input_image_base64:\n",
    "        request_body[\"inputImage\"] = input_image_base64\n",
    "    \n",
    "    body = json.dumps(request_body)\n",
    "    \n",
    "    response = bedrock.invoke_model(\n",
    "    \tbody=body, \n",
    "    \tmodelId=\"amazon.titan-embed-image-v1\",\n",
    "        accept=\"application/json\", \n",
    "    \tcontentType=\"application/json\"\n",
    "\n",
    "    )\n",
    "    \n",
    "    response_body = json.loads(response.get('body').read())\n",
    "    \n",
    "    embedding = response_body.get(\"embedding\")\n",
    "    \n",
    "    return embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`get_vector_from_file` function takes an image file path, encodes the image to base64, generates an embedding vector using [Titan Multimodal Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-multiemb-models.html), and returns the vector - allowing images to be represented as vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#creates a vector from a file\n",
    "def get_vector_from_file(file_path):\n",
    "    with open(file_path, \"rb\") as image_file:\n",
    "        input_image_base64 = base64.b64encode(image_file.read()).decode('utf8')\n",
    "    \n",
    "    vector = get_multimodal_vector(input_image_base64 = input_image_base64)\n",
    "    \n",
    "    return vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creates a list of (path, vector) tuples from a directory.\n",
    "\n",
    "A sample of [Kaggle Animal Image Dataset (90 Different Animals)](https://www.kaggle.com/datasets/iamsouravbanerjee/animal-image-dataset-90-different-animals) is used in this app."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_vectors_from_directory(path_name):\n",
    "    items = []\n",
    "    sub_1 = os.listdir(path_name)\n",
    "    for n in sub_1:\n",
    "        if n.endswith('.jpg'):\n",
    "            file_path = os.path.join(path_name,n)\n",
    "            size_image(file_path)\n",
    "            vector = get_vector_from_file(file_path)\n",
    "            items.append((file_path, vector))\n",
    "        else:\n",
    "            for n_2 in os.listdir(path_name+\"/\"+n):\n",
    "                if n_2.endswith('.jpg'):\n",
    "                    file_path = os.path.join(path_name+\"/\"+n, n_2)\n",
    "                    size_image(file_path)\n",
    "                    vector = get_vector_from_file(file_path)\n",
    "                    items.append((file_path, vector))\n",
    "                else:\n",
    "                    print(\"no a .jpg file: \",n_2)\n",
    "\n",
    "    return items\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def size_image(file_path):\n",
    "    # Maximum image size supported is 2048 x 2048 pixel\n",
    "    image = Image.open(file_path) #open image\n",
    "    width, height = image.size # Get the width and height of the image in pixels\n",
    "    if width > 2048 or height > 2048:\n",
    "        print(f\"Big File:{file_path} , width: {width}, height {height} px\")\n",
    "        dif_width = width - 2048\n",
    "        dif_height = height - 2048\n",
    "        if dif_width > dif_height:\n",
    "            ave = 1-(dif_width/width)\n",
    "            new_width = int(width*ave)\n",
    "            new_height = int(height*ave)\n",
    "        else:\n",
    "            ave = 1-(dif_height/height)\n",
    "            new_width = int(width*ave)\n",
    "            new_height = int(height*ave)\n",
    "        print(f\"New file: {file_path} , width: {new_width}, height {new_height} px\")\n",
    "\n",
    "        new_image = image.resize((new_width, new_height))\n",
    "        # Save New image\n",
    "        new_image.save(file_path)\n",
    " \n",
    "    return\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creates and returns an in-memory vector store to be used in the application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id      = \"amazon.titan-embed-image-v1\"\n",
    "model_kwargs  = {\"embeddingConfig\": { \"outputEmbeddingLength\": 1024}}\n",
    "def get_bedrock_embeddings(model_id, model_kwargs):\n",
    "    return BedrockEmbeddings(model_id=model_id, client=bedrock_client, model_kwargs=model_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_vector_db(path_name):\n",
    "    image_vectors = get_image_vectors_from_directory(path_name)\n",
    "        \n",
    "    text_embeddings = [(\"\", item[1]) for item in image_vectors]\n",
    "    metadatas = [{\"image_path\": item[0]} for item in image_vectors]\n",
    "        \n",
    "    db = FAISS.from_embeddings(\n",
    "        text_embeddings=text_embeddings,\n",
    "        embedding = get_bedrock_embeddings(model_id, model_kwargs),\n",
    "        metadatas = metadatas\n",
    "    )\n",
    "    print(f\"Vector Database:{db.index.ntotal} docs\")\n",
    "    return db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_file = \"animals/animals\"\n",
    "path_name = f\"{path_file}\"\n",
    "db = create_vector_db(path_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Save to a local Vector database.](https://python.langchain.com/docs/integrations/vectorstores/faiss/#as-a-retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_file = \"animals.vdb\"\n",
    "db.save_local(db_file)\n",
    "print(f\"vectordb was saved in {db_file}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://python.langchain.com/docs/how_to/#vector-stores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Query](https://python.langchain.com/docs/integrations/vectorstores/faiss/#querying) by text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"dog\"\n",
    "search_vector = get_multimodal_vector(input_text=query)\n",
    "results = db.similarity_search_by_vector(embedding=search_vector)\n",
    "for res in results:\n",
    "    with open(res.metadata['image_path'], \"rb\") as f:\n",
    "        img = BytesIO(f.read())\n",
    "        image = Image.open(img)\n",
    "        image.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Query](https://python.langchain.com/docs/integrations/vectorstores/faiss/#querying) by Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_path = 'animals/animals/cat/9d21019336.jpg'\n",
    "vector = get_vector_from_file(query_path)\n",
    "results = db.similarity_search_by_vector(embedding=vector)\n",
    "for res in results:\n",
    "    with open(res.metadata['image_path'], \"rb\") as f:\n",
    "        img = BytesIO(f.read())\n",
    "        image = Image.open(img)\n",
    "        image.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load and Query local Vector database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_aws import BedrockEmbeddings  # to create embeddings for the documents.\n",
    "db_file = \"animals.vdb\"\n",
    "bedrock_embeddings          = BedrockEmbeddings(model_id=\"amazon.titan-embed-text-v1\",client=bedrock_client)\n",
    "new_db = FAISS.load_local(db_file, bedrock_embeddings, allow_dangerous_deserialization=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"cat\"\n",
    "search_vector = get_multimodal_vector(input_text=query)\n",
    "results = new_db.similarity_search_by_vector(embedding=search_vector)\n",
    "for res in results:\n",
    "    with open(res.metadata['image_path'], \"rb\") as f:\n",
    "        img = BytesIO(f.read())\n",
    "        image = Image.open(img)\n",
    "        image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_path = 'animals/animals/leopard/9cc45df890.jpg'\n",
    "vector = get_vector_from_file(query_path)\n",
    "results = db.similarity_search_by_vector(embedding=vector)\n",
    "for res in results:\n",
    "    with open(res.metadata['image_path'], \"rb\") as f:\n",
    "        img = BytesIO(f.read())\n",
    "        image = Image.open(img)\n",
    "        image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Delete Vectordb](https://python.langchain.com/docs/integrations/vectorstores/faiss/#delete)\n",
    "\n",
    "You can also delete records from the vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"count before:\", new_db.index.ntotal)\n",
    "new_db.delete([new_db.index_to_docstore_id[0]])\n",
    "print(\"count after:\", new_db.index.ntotal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Delete the entire database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_db.delete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}