{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 2: Image Search\n",
    "\n",
    "This module uses [OpenAI CLIP](https://huggingface.co/docs/transformers/model_doc/clip) and TiDB Serverless Vector Search.\n",
    "\n",
    "We will use the CLIP model to encode the image to a 512-dimensional vector and store them in TiDB Serverless. Then use the same model to encode the text query and search for the most similar images in TiDB Serverless."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \\\n",
    "    torch==2.6.0 \\\n",
    "    transformers==4.51.1 \\\n",
    "    requests==2.32.3 \\\n",
    "    ipyplot==1.1.2 \\\n",
    "    datasets==3.5.0 \\\n",
    "    sqlalchemy==2.0.36 \\\n",
    "    pymysql==1.1.1 \\\n",
    "    tidb_vector==0.0.14"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial the Database and Table\n",
    "\n",
    "> **Note:**\n",
    ">\n",
    "> - We already set the `TIDB_HOST`, `TIDB_PORT`, `TIDB_USERNAME`, `TIDB_PASSWORD`, and `TIDB_DATABASE` in the environment parameters. If you want to use this code snippet out of TiDB Labs platform, please set them beforehand.\n",
    "> - In this example, we use CLIP to generate text and image embeddings with `512` dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from sqlalchemy import URL, create_engine, Column, Integer\n",
    "from sqlalchemy.orm import declarative_base, sessionmaker\n",
    "from tidb_vector.sqlalchemy import VectorType\n",
    "\n",
    "CLIP_DIMENSION = 512\n",
    "\n",
    "engine = create_engine(URL(\n",
    "    \"mysql+pymysql\",\n",
    "    username=os.getenv(\"TIDB_USERNAME\"),\n",
    "    password=os.getenv(\"TIDB_PASSWORD\"),\n",
    "    host=os.getenv(\"TIDB_HOST\"),\n",
    "    port=int(os.getenv(\"TIDB_PORT\")),\n",
    "    database=os.getenv(\"TIDB_DATABASE\"),\n",
    "    query={\"ssl_verify_cert\": True, \"ssl_verify_identity\": True},\n",
    "))\n",
    "\n",
    "Session = sessionmaker(bind=engine)\n",
    "Base = declarative_base()\n",
    "\n",
    "class ImageSearchTest(Base):\n",
    "    __tablename__ = \"image_search_test\"\n",
    "\n",
    "    id = Column(Integer, primary_key=True)\n",
    "    image_id = Column(Integer)\n",
    "    embedding = Column(VectorType(CLIP_DIMENSION))\n",
    "\n",
    "Base.metadata.drop_all(engine)\n",
    "Base.metadata.create_all(engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial CLIP model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import CLIPProcessor, CLIPModel\n",
    "\n",
    "\n",
    "model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load test images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "\n",
    "imagenet_datasets = datasets.load_dataset('theodor1289/imagenet-1k_tiny', split='train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect the imagenet datasets\n",
    "imagenet_datasets[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "extract the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipyplot\n",
    "\n",
    "imagenet_images = [i['image'] for i in imagenet_datasets]\n",
    "ipyplot.plot_images(imagenet_images, max_images=20, img_width=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the encode function and other helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_images_to_embeddings(images):\n",
    "    # accept a list of images and return the image embeddings\n",
    "    with torch.no_grad():\n",
    "        inputs = processor(images=images, return_tensors=\"pt\")\n",
    "        image_features = model.get_image_features(**inputs)\n",
    "        return image_features.cpu().detach().numpy()\n",
    "\n",
    "def encode_text_to_embedding(text):\n",
    "    # accept a text and return the text embedding\n",
    "    with torch.no_grad():\n",
    "        inputs = processor(text=text, return_tensors=\"pt\")\n",
    "        text_features = model.get_text_features(**inputs)\n",
    "        return text_features.cpu().detach().numpy()[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Store the images and their corresponding image embeddings in TiDB Serverless"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_embedding = encode_images_to_embeddings(imagenet_images)\n",
    "objects = []\n",
    "\n",
    "for i, embedding in enumerate(images_embedding):\n",
    "    img = imagenet_images[i]\n",
    "    objects.append(\n",
    "        ImageSearchTest(\n",
    "            image_id=i,\n",
    "            embedding=embedding\n",
    "        )\n",
    "    )\n",
    "\n",
    "with Session() as session:\n",
    "    session.add_all(objects)\n",
    "    session.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Search for similar images using the text query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import asc\n",
    "\n",
    "query_text = \"dog\"\n",
    "query_text_embedding = encode_text_to_embedding(query_text)\n",
    "\n",
    "with Session() as session:\n",
    "    results = session.query(\n",
    "        ImageSearchTest,\n",
    "        ImageSearchTest.embedding.cosine_distance(query_text_embedding).label(\"distance\"),\n",
    "    ).order_by(\n",
    "        asc(\"distance\")\n",
    "    ).limit(5).all()\n",
    "\n",
    "\n",
    "    similar_images = []\n",
    "    similarities = []\n",
    "    for obj, d in results:\n",
    "        similar_images.append(imagenet_images[obj.image_id])\n",
    "        similarities.append(round(1 - d, 3))\n",
    "\n",
    "# display the similar images\n",
    "ipyplot.plot_images(similar_images, labels=similarities, img_width=100)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}