Skip to content

Instantly share code, notes, and snippets.

@manuthu

manuthu/lancedb-multimodal-myntra-fashion-search-engine.ipynb

Created August 4, 2024 00:31
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save manuthu/6c3845f511a1d160c0afa84344b80fba to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save manuthu/6c3845f511a1d160c0afa84344b80fba to your computer and use it in GitHub Desktop.
Download ZIP
lancedb-multimodal-myntra-fashion-search-engine.ipynb
Display the source blob
Display the rendered blob
Raw
lancedb-multimodal-myntra-fashion-search-engine.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/manuthu/6c3845f511a1d160c0afa84344b80fba/lancedb-multimodal-myntra-fashion-search-engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"source": [
"# LanceDB Multimodal Myntra Fashion Search Engine"
],
"metadata": {
"id": "xgydN0D4TtBF"
}
},
{
"cell_type": "markdown",
"source": [
"## Preliminaries"
],
"metadata": {
"id": "BUElYQzKT41t"
}
},
{
"cell_type": "code",
"source": [
"%%capture\n",
"\n",
"!pip install lancedb\n",
"!pip install open_clip_torch"
],
"metadata": {
"id": "PQj_jlbRUteu"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import os\n",
"import argparse\n",
"import pandas as pd\n",
"from PIL import Image\n",
"from pathlib import Path\n",
"from random import sample\n",
"\n",
"import lancedb\n",
"from lancedb.pydantic import LanceModel, Vector\n",
"from lancedb.embeddings import EmbeddingFunctionRegistry\n",
"\n",
"from typing import Any"
],
"metadata": {
"id": "uddRx_vxT6_t"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"## Download Data\n",
"- For this project you need to download the [Myntra Fashion Product Dataset]( https://www.kaggle.com/datasets/hiteshsuthar101/myntra-fashion-product-dataset) from Kaggle and store it here in the `input` folder.\n",
"- Note that while creating the table you shall pass the path of the folder in which the images are present, example `/content/input/images`"
],
"metadata": {
"id": "GEtZ52fYZNz0"
}
},
{
"cell_type": "markdown",
"source": [
"## Embedding Model"
],
"metadata": {
"id": "JZafdZy-TiJa"
}
},
{
"cell_type": "code",
"source": [
"def register_model(model_name: str) -> Any:\n",
" \"\"\"\n",
" Register a model with the given name using LanceDB's EmbeddingFunctionRegistry.\n",
"\n",
" Args:\n",
" model_name (str): The name of the model to register.\n",
"\n",
" Returns:\n",
" model: The registered model instance.\n",
"\n",
" Usage:\n",
" >>> model = register_model(\"open-clip\")\n",
" \"\"\"\n",
" registry = EmbeddingFunctionRegistry.get_instance()\n",
" model = registry.get(model_name).create()\n",
" return model"
],
"metadata": {
"id": "14TJFhtxTczp"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"## Schema"
],
"metadata": {
"id": "l2bJC-3CTd36"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5upIYuE5TJ13"
},
"outputs": [],
"source": [
"# Register the OpenAI CLIP model\n",
"clip = register_model(\"open-clip\")\n",
"\n",
"\n",
"class Myntra(LanceModel):\n",
" \"\"\"\n",
" Represents a Myntra Schema.\n",
"\n",
" Attributes:\n",
" vector (Vector): The vector representation of the item.\n",
" image_uri (str): The URI of the item's image.\n",
" \"\"\"\n",
"\n",
" vector: Vector(clip.ndims()) = clip.VectorField()\n",
" image_uri: str = clip.SourceField()\n",
"\n",
" @property\n",
" def image(self):\n",
" return Image.open(self.image_uri)\n",
"\n",
"\n",
"# Function to map schema name to schema class\n",
"def get_schema_by_name(schema_name: str) -> Any:\n",
" \"\"\"\n",
" Retrieves the schema object based on the given schema name.\n",
"\n",
" Args:\n",
" schema_name (str): The name of the schema.\n",
"\n",
" Returns:\n",
" object: The schema object corresponding to the given schema name, or None if not found.\n",
"\n",
" Usage:\n",
" >>> schema = get_schema_by_name(\"Myntra\")\n",
" \"\"\"\n",
" schema_map = {\n",
" \"Myntra\": Myntra,\n",
" }\n",
" return schema_map.get(schema_name)"
]
},
{
"cell_type": "markdown",
"source": [
"## Creating a Table"
],
"metadata": {
"id": "dy3LohCmTlW3"
}
},
{
"cell_type": "code",
"source": [
"def create_table(\n",
" database: str,\n",
" table_name: str,\n",
" data_path: str,\n",
" schema: Any = Myntra,\n",
" mode: str = \"overwrite\",\n",
") -> None:\n",
" \"\"\"\n",
" Create a table in the specified vector database and add data to it.\n",
"\n",
" Args:\n",
" database (str): The name of the database to connect to.\n",
" table_name (str): The name of the table to create.\n",
" data_path (str): The path to the data directory.\n",
" schema (Schema, optional): The schema to use for the table. Defaults to Myntra.\n",
" mode (str, optional): The mode for creating the table. Defaults to \"overwrite\".\n",
"\n",
" Returns:\n",
" None\n",
"\n",
" Usage:\n",
" >>> create_table(database=\"~/.lancedb\"\", table_name=\"myntra\", data_path=\"input\")\n",
" \"\"\"\n",
"\n",
" # Connect to the lancedb database\n",
" db = lancedb.connect(database)\n",
"\n",
" # Check if the table already exists in the database\n",
" if table_name in db:\n",
" print(f\"Table {table_name} already exists in the database\")\n",
" table = db[table_name]\n",
"\n",
" # if it does not exist then create a new table\n",
" else:\n",
"\n",
" print(f\"Creating table {table_name} in the database\")\n",
"\n",
" # Create the table with the given schema\n",
" table = db.create_table(table_name, schema=schema, mode=mode)\n",
"\n",
" # Define the Path of the images and obtain the Image uri\n",
" p = Path(data_path).expanduser()\n",
" uris = [str(f) for f in p.glob(\"*.jpg\")]\n",
" print(f\"Found {len(uris)} images in {p}\")\n",
"\n",
" # Sample 1000 images from the data\n",
" # Increase this value for more accurate results but\n",
" # it will take more time to process embeddings\n",
" uris = sample(uris, 100)\n",
"\n",
" # Add the data to the table\n",
" print(f\"Adding {len(uris)} images to the table\")\n",
" table.add(pd.DataFrame({\"image_uri\": uris}))\n",
" print(f\"Added {len(uris)} images to the table\")"
],
"metadata": {
"id": "uDhRy1uTTnY-"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# The data_path should refer to the folder in which the images are located\n",
"\n",
"create_table(database=\"~/.lancedb\", table_name=\"myntra\", data_path=\"/content/input/images\")"
],
"metadata": {
"id": "FP5uK1rcP38K",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "ce98efcf-cc66-4b71-ef58-d6b1512e0583"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Creating table myntra_v4 in the database\n",
"Found 205 images in /content/input/images\n",
"Adding 100 images to the table\n"
]
},
{
"output_type": "stream",
"name": "stderr",
"text": [
"100%|██████████| 64/64 [00:16<00:00, 3.79it/s]\n",
"100%|██████████| 36/36 [00:09<00:00, 3.76it/s]"
]
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"Added 100 images to the table\n"
]
},
{
"output_type": "stream",
"name": "stderr",
"text": [
"\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"## Vector Search"
],
"metadata": {
"id": "3vAWovRJTqhv"
}
},
{
"cell_type": "code",
"source": [
"def run_vector_search(\n",
" database: str,\n",
" table_name: str,\n",
" schema: Any,\n",
" search_query: Any,\n",
" limit: int = 6,\n",
" output_folder: str = \"output\",\n",
") -> None:\n",
" \"\"\"\n",
" This function performs a vector search on the specified database and table using the provided search query.\n",
" The search can be performed on either text or image data. The function retrieves the top 'limit' number of results\n",
" and saves the corresponding images in the 'output_folder' directory. The function assumes if the search query ends\n",
" with '.jpg' or '.png', it is an image search, otherwise it is a text search.\n",
" Args:\n",
" database (str): The path to the database.\n",
" table_name (str): The name of the table.\n",
" schema (Schema): The schema to use for converting search results to Pydantic models.\n",
" search_query (Any): The search query, can be text or image.\n",
" limit (int, optional): The maximum number of results to return. Defaults to 6.\n",
" output_folder (str, optional): The folder to save the output images. Defaults to \"output\".\n",
"\n",
" Returns:\n",
" None\n",
"\n",
" Usage:\n",
" >>> run_vector_search(database=\"~/.lancedb\", table_name=\"myntra\", schema=Myntra, search_query=\"Black Kurta\")\n",
"\n",
" \"\"\"\n",
"\n",
" # Create the output folder if it does not exist\n",
" if os.path.exists(output_folder):\n",
" for file in os.listdir(output_folder):\n",
" os.remove(os.path.join(output_folder, file))\n",
" else:\n",
" os.makedirs(output_folder)\n",
"\n",
" # Connect to the lancedb database\n",
" db = lancedb.connect(database)\n",
"\n",
" # Open the table\n",
" table = db.open_table(table_name)\n",
"\n",
" # Check if the search query is an image or text\n",
" try:\n",
" if search_query.endswith(\".jpg\") or search_query.endswith(\".png\"):\n",
" search_query = Image.open(search_query)\n",
" else:\n",
" search_query = search_query\n",
" except AttributeError as e:\n",
" if str(e) == \"'JpegImageFile' object has no attribute 'endswith'\":\n",
" print(\"Running via Streamlit, search query is already an array so skipping opening image using Pillow\")\n",
" else:\n",
" raise\n",
"\n",
" # Perform the vector search and retrieve the results\n",
" rs = table.search(search_query).limit(limit).to_pydantic(schema)\n",
"\n",
" # Save the images to the output folder\n",
" for i in range(limit):\n",
" image_path = os.path.join(output_folder, f\"image_{i}.jpg\")\n",
" rs[i].image.save(image_path, \"JPEG\")"
],
"metadata": {
"id": "unDrIuzKTrpJ"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"After the search is done, the results will be saved in the `output` folder."
],
"metadata": {
"id": "yPCppZacZ1us"
}
},
{
"cell_type": "markdown",
"source": [
"## Text Search"
],
"metadata": {
"id": "b6vk1e1wdu0x"
}
},
{
"cell_type": "code",
"source": [
"run_vector_search(database=\"~/.lancedb\", table_name=\"myntra\", schema=Myntra, search_query=\"White Kurta\", limit=3, output_folder=\"output\")"
],
"metadata": {
"id": "zBSCG4otP-QZ"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"## Image Search"
],
"metadata": {
"id": "hgVFAc_XdxHo"
}
},
{
"cell_type": "code",
"source": [
"run_vector_search(database=\"~/.lancedb\", table_name=\"myntra\", schema=Myntra, search_query=\"/content/input/images/0.jpg\", limit=3, output_folder=\"output\")"
],
"metadata": {
"id": "MfmUygJEdyZO"
},
"execution_count": null,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment