cbir.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0e129fa-b126-416d-9ca0-aecc4028c6a3",
   "metadata": {},
   "source": [
    "# Content Based Image Retrieval\n",
    "\n",
    "This notebook is an example of a simple CBIR system. Using a pretrained **Vision Transformer (ViT)** to extract image features comparisons are made between the feature vectors of the query image and each of the images in the database. The top 5 images ranked by cosine similarity are presented to the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef9bea6-0e49-4f7b-8c25-c1a853d3a8b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "from timm import create_model\n",
    "from torchvision import transforms\n",
    "from torchvision.transforms.functional import to_pil_image, to_tensor\n",
    "from torchvision.utils import make_grid\n",
    "from PIL import Image\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a0c7efc-8d96-43aa-9f22-ecf215533ca9",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "This is the initial image set which makes up the database which we query on. A function to fetch images as a tensor from the provided path is declared. Images are scaled to dimensions of 224 by 224 pixels with 3 channels and normalized. This transformation allows the image to be processed by the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "305d65bb-8b29-469c-a886-d15cba3dd394",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_DIR = r\"D:\\Projs\\datasets\\small\\monarch_150\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d8e571-d3c7-4776-ac00-4f6df926e256",
   "metadata": {},
   "outputs": [],
   "source": [
    "def path_to_tensor(img_path, img_transforms):\n",
    "    image = Image.open(img_path).convert(\"RGB\")\n",
    "    tensor_image = img_transforms(image)\n",
    "    return tensor_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22402ea5-30b3-42cc-895b-465d77e7c32c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImgFolderDataset(Dataset):\n",
    "    def __init__(self, img_dir, transforms):\n",
    "        self.img_dir = img_dir\n",
    "        self.transforms = transforms\n",
    "        self.img_filenames = os.listdir(self.img_dir)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_filenames)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = os.path.join(self.img_dir, self.img_filenames[idx])\n",
    "        tensor_image = path_to_tensor(img_path, self.transforms)\n",
    "        return tensor_image, img_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d252d0f9-4268-4dfe-963e-22b7c1f92209",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transforms = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "dataset = ImgFolderDataset(DATASET_DIR, data_transforms)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6292e83c-1f95-4693-8c3b-773970e96908",
   "metadata": {},
   "source": [
    "## Model\n",
    "Image to feature vector mapping is done with **ViT-Base (ViT-B/32)** described in [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929). The pretrained model is fetched through the `timm` library. A slight modification is done to the original model, the last softmax layer is removed as described in the [Investigating the Vision Transformer Model for Image Retrieval Tasks](https://arxiv.org/abs/2101.03771)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d951531e-3ee5-4344-bc7f-ca1ab2e9af71",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model = create_model(\"vit_base_patch32_224\", pretrained=True)\n",
    "# Remove softmax layer, output is of size (B, 768)\n",
    "model.head = nn.Identity()\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8455215a-92f4-4232-a30e-61b0489cc41b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Database\n",
    "The image database is a dictionary which maps feature vectors to file paths. The dataset from above is inserted into the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b037f60-b778-4ceb-8921-436ad625b2f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_dict = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60ad3bf3-0dc9-4123-8ac5-3f09b733f566",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    for img, img_path in dataset:\n",
    "        img = img[None,:].to(device)\n",
    "        img_feats = model(img)\n",
    "        img_dict[img_feats] = img_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c78d3a63-f047-4087-a33d-0f0a12933eb9",
   "metadata": {},
   "source": [
    "## Similarity criterion\n",
    "Image similarity is calculated using cosine similarity between two feature vectors, in this case the query image feature vector and each individual feature vector from the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "317ff517-da81-4ddc-be78-8d3dd2ba08b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_fun = nn.CosineSimilarity(dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a86d32ec-6598-4309-b5d8-0c52f5e3a1bf",
   "metadata": {},
   "source": [
    "## Querying\n",
    "Top 5 images from the database are retrieved based on similarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93dc3ea0-43ac-4a8c-9309-b1dc691726d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_top5(img_path, img_dict, sim_fun, model):\n",
    "    img = path_to_tensor(img_path, data_transforms)\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        img = img[None,:].to(device)\n",
    "        feats = model(img)\n",
    "    comparisons = [(sim_fun(feats, img2_feats), img2_path) for img2_feats, img2_path in img_dict.items()]\n",
    "    top5 = sorted(comparisons, reverse=True)[:5]\n",
    "    top5 = [Image.open(x[1]) for x in top5]\n",
    "    return to_pil_image(make_grid([to_tensor(i) for i in top5], nrow=5, padding=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa184acb-4f31-416b-aeaf-34874305393c",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = r\"D:\\Projs\\datasets\\small\\monarch_651\\0000_000009.png\"\n",
    "query = Image.open(img_path)\n",
    "top5 = find_top5(img_path, img_dict, sim_fun, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc3c3521-1e14-4074-94e0-7c2a6dbad54c",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(query, top5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "262aeec2-596f-4845-acaa-adce0a9a3d81",
   "metadata": {},
   "source": [
    "## Deeplake Database\n",
    "The vector database is created locally or loaded if it already exists. The embedding function passes the argument to the model and returns the embeddings as a NumPy array. The list of paths and images is passed to the VectorStore which calls the embedding function and stores the values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3bb2895-9771-4d61-abbe-8f7fc66679b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from deeplake.core.vectorstore.deeplake_vectorstore import VectorStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53bc6cb3-8203-49a1-ab21-5e7ea3bf86cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store_path = r\"./vector_store\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d6bed58-5735-4d04-b244-5a390406fb8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store = VectorStore(path=vector_store_path, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1e230c6-006e-44c2-913e-836fe48e6142",
   "metadata": {},
   "outputs": [],
   "source": [
    "def embedding_fn(img, model=model):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        if isinstance(img, list):  \n",
    "            embeddings = []\n",
    "            for im in img:\n",
    "                embeddings.append(model(im[None,:].to(device)).to(\"cpu\").numpy()[0])\n",
    "        else:\n",
    "            embeddings = model(img[None,:].to(device)).to(\"cpu\").numpy()[0]\n",
    "        return np.array(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec9716e4-ae07-431b-89c8-9a00044c2dd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    path_list = [img_path for _, img_path in dataset]\n",
    "    img_list = [img for img, _ in dataset]\n",
    "    vector_store.add(\n",
    "        text=path_list,\n",
    "        embedding_function=embedding_fn,\n",
    "        embedding_data=img_list,\n",
    "        metadata=None\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59b37cac-c4c9-425e-b420-b7f4ba077107",
   "metadata": {},
   "source": [
    "## Querying on a DB\n",
    "The VectorStore function is called with the query image and embedding function as arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3af6844-0276-4dc5-8450-773b22a03cb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def top5_db(img_path):\n",
    "    img = path_to_tensor(img_path, data_transforms)\n",
    "    search_results = vector_store.search(embedding_data=img, k=5, embedding_function=embedding_fn)\n",
    "    search_results[\"text\"]\n",
    "    top5 = [Image.open(x) for x in search_results[\"text\"]]\n",
    "    return to_pil_image(make_grid([to_tensor(i) for i in top5], nrow=5, padding=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5de048c-b873-4b3f-b0ab-5f1e2e2bfced",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = r\"D:\\Projs\\datasets\\small\\monarch_651\\0000_000009.png\"\n",
    "query = Image.open(img_path)\n",
    "top5 = top5_db(img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ef19da-7add-4421-ba49-a87c4bf783aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(query, top5)"
   ]
  }
 ],
 "metadata": {
  "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}