2-process-and-upload.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process each image and upload the vectors to OpenSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import clip\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "# Load the open CLIP model\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using {device}\")\n",
    "model, preprocess = clip.load(\"ViT-B/32\", device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a function that computes the feature vectors for a batch of images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_clip_features(photos_batch):\n",
    "    # Load all the photos from the files\n",
    "    photos = [Image.open(photo_file) for photo_file in photos_batch]\n",
    "    \n",
    "    # Preprocess all photos\n",
    "    photos_preprocessed = torch.stack([preprocess(photo) for photo in photos]).to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        # Encode the photos batch to compute the feature vectors and normalize them\n",
    "        photos_features = model.encode_image(photos_preprocessed)\n",
    "        photos_features /= photos_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "    # Transfer the feature vectors back to the CPU and convert to numpy\n",
    "    return photos_features.cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a function to index embeddings into OpenSearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "SERVICE_URI = os.getenv(\"SERVICE_URI\")\n",
    "\n",
    "from opensearchpy import OpenSearch\n",
    "opensearch = OpenSearch(SERVICE_URI, use_ssl=True)\n",
    "from opensearchpy.helpers import bulk\n",
    "\n",
    "def index_embeddings_to_opensearch(data):\n",
    "    actions = []\n",
    "    for d in data:\n",
    "        action = {\n",
    "            \"_index\": \"photos\",  # Update with your index name\n",
    "            \"_source\": {\n",
    "                \"image_url\": d['image_url'],\n",
    "                \"embedding\": d['embedding'].tolist()\n",
    "            }\n",
    "        }\n",
    "        actions.append(action)\n",
    "    success, _ = bulk(opensearch, actions, index=\"photos\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Iterate over images and process them in batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Path to the directory containing photos\n",
    "image_dir = \"photos\"\n",
    "\n",
    "# Batch size for processing images and indexing embeddings\n",
    "batch_size = 100\n",
    "\n",
    "# List to store embeddings\n",
    "data = []\n",
    "\n",
    "# Process images in batches\n",
    "image_files = os.listdir(image_dir)\n",
    "total_size = len(image_files)\n",
    "for i in range(0, len(image_files), batch_size):\n",
    "    batch_files = image_files[i:i+batch_size]\n",
    "    batch_file_paths = [os.path.join(image_dir, file) for file in batch_files]\n",
    "\n",
    "    # Compute embeddings for the batch of images\n",
    "    batch_embeddings = compute_clip_features(batch_file_paths)\n",
    "\n",
    "    # Create data dictionary for indexing\n",
    "    for file_path, embedding in zip(batch_file_paths, batch_embeddings):\n",
    "        data.append({'image_url': file_path, 'embedding': embedding})\n",
    "\n",
    "    # Check if we have enough data to index\n",
    "    if len(data) >= batch_size:\n",
    "        index_embeddings_to_opensearch(data)\n",
    "        print(f\"Indexed {i+batch_size} embeddings of {total_size} to OpenSearch\")\n",
    "        data = []\n",
    "\n",
    "# Index any remaining data\n",
    "if len(data) > 0:\n",
    "    index_embeddings_to_opensearch(data)\n",
    "\n",
    "print(\"All embeddings indexed successfully.\")\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}