batch process, faiss, gpu support, optimise

.gitignore

@@ -2,5 +2,6 @@ env/
 images/
 __pycache__/
 *.tree
+*.index
 secrets.toml
 kaggle.json

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import torch
 import os
 import torchvision
-from annoy import AnnoyIndex
+import faiss
 from PIL import Image
 import traceback
 from tqdm import tqdm
@@ -11,27 +11,41 @@ from slugify import slugify
 import opendatasets as od
 import json
 import argparse
+from streamlit_cropper import st_cropper
+from azure.storage.blob import BlobServiceClient
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms
+import numpy as np
+import faiss.contrib.torch_utils
 
-
+BATCH_SIZE = 200
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 FOLDER = "images/"
 NUM_TREES = 100
 FEATURES = 1000
 FILETYPES = [".png", ".jpg", ".jpeg", ".tiff", ".bmp"]
-
-from azure.storage.blob import BlobServiceClient
+LIBRARIES = [
+    "https://www.kaggle.com/datasets/athota1/caltech101",
+    "https://www.kaggle.com/datasets/gpiosenka/sports-classification",
+    "https://www.kaggle.com/datasets/puneet6060/intel-image-classification",
+    "https://www.kaggle.com/datasets/kkhandekar/image-dataset",
+]
 
 
 @st.cache_resource
 def dl_embeddings():
     """dl pretrained embeddings in production environment instead of creating"""
     # Connect to your Blob Storage account
+    if os.path.isfile(f"{slugify(FOLDER)}.index"):
+        print("Embeddings files already exists, skip download")
+        return
     connect_str = st.secrets["connectionstring"]
     blob_service_client = BlobServiceClient.from_connection_string(connect_str)
 
     # Specify container and blob names
     container_name = "imagessearch"
-    blob_name = f"{slugify(FOLDER)}.tree"
+    blob_name = f"{slugify(FOLDER)}.index"
 
     # Get a reference to the blob
     blob_client = blob_service_client.get_blob_client(
@@ -39,7 +53,7 @@ def dl_embeddings():
     )
 
     # Download the binary data
-    download_file_path = f"{slugify(FOLDER)}.tree"  # Path to save the downloaded file
+    download_file_path = f"{slugify(FOLDER)}.index"  # Path to save the downloaded file
     with open(download_file_path, "wb") as download_file:
         download_file.write(blob_client.download_blob().readall())
 
@@ -56,16 +70,18 @@ def load_dataset():
             },
             f,
         )
-    od.download(
-        "https://www.kaggle.com/datasets/kkhandekar/image-dataset",
-        "images/",
-    )
+    for lib in LIBRARIES:
+        od.download(
+            lib,
+            "images/",
+        )
 
 
 # Load a pre-trained image feature extractor model
 @st.cache_resource
 def load_model():
     """Loads a pre-trained image feature extractor model."""
+    print("Loading pretrained model...")
     model = torch.hub.load(
         "NVIDIA/DeepLearningExamples:torchhub",
         "nvidia_efficientnet_b0",
@@ -104,9 +120,19 @@ def load_images(file_paths):
     return images
 
 
+def load_image(file_path):
+    """Load all the images from file paths."""
+    try:
+        image = Image.open(file_path).resize([224, 224])
+        return image
+    except BaseException as e:
+        print("Error loading ", file_path, e)
+
+
 # Function to preprocess images
 def preprocess_image(image):
     """Preprocesses an image for feature extraction."""
+
     if image.mode == "RGB":  # Already has 3 channels
         pass  # No need to modify
     elif image.mode == "L":  # Grayscale image
@@ -128,57 +154,77 @@ def preprocess_image(image):
     return preprocess(image)
 
 
+class ImageLoader(Dataset):
+    def __init__(self, image_files, transform, load_image):
+        self.transform = transform
+        self.load_image = load_image
+        self.image_files = image_files
+
+    def __len__(self):
+        return len(self.image_files)
+
+    def __getitem__(self, index):
+        return self.transform(self.load_image(self.image_files[index]))
+
+
 # Extract features from a list of images
-def extract_features(images, model):
+def extract_features(file_paths, model):
     """Extracts features from a list of images."""
     print("Extracting features:")
+    loader = DataLoader(
+        ImageLoader(file_paths, transform=preprocess_image, load_image=load_image),
+        batch_size=BATCH_SIZE,
+    )
     features = []
-    for image in images:
-        with torch.no_grad():
-            feature = model(preprocess_image(image).unsqueeze(0)).squeeze(0)
-            features.append(feature.numpy())
-    return features
+    model = model.to(DEVICE)
+    with torch.no_grad():
+        for batch_idx, images in enumerate(tqdm(loader)):
+            images = images.to(DEVICE)
+            features.append(model(images))
+    return torch.cat(features)
 
 
 # Build an Annoy index for efficient similarity search
 def build_annoy_index(features):
     """Builds an Annoy index for efficient similarity search."""
-    print("Building annoy index:")
+    print("Building faiss index:")
     f = features[0].shape[0]  # Feature dimensionality
-    t = AnnoyIndex(f, "angular")  # Use angular distance for image features
-    for i, feature in tqdm(enumerate(features)):
-        t.add_item(i, feature)
-    t.build(NUM_TREES)  # Adjust num_trees for accuracy vs. speed trade-off
-    return t
+    index = faiss.IndexIDMap(faiss.IndexFlatIP(f))
+    index.add_with_ids(
+        features.cpu().detach().numpy(), np.array(range(len(features)))
+    )  # Adjust num_trees for accuracy vs. speed trade-off
+    print("built faiss index:")
+    return index
 
 
 # Perform reverse image search
-def search_similar_images(uploaded_file, f=FEATURES, num_results=5):
+def search_similar_images(query_image, num_results, f=FEATURES):
     """Finds similar images based on a query image feature."""
-    index = AnnoyIndex(f, "angular")
-    index.load(f"{slugify(FOLDER)}.tree")
-    query_image = Image.open(uploaded_file)
-    model = load_model()
+    index = faiss.read_index(f"{slugify(FOLDER)}.index")
+    model = load_model().to(DEVICE)
     # Extract features and search
-    query_feature = (
-        model(preprocess_image(query_image).unsqueeze(0)).squeeze(0).detach().numpy()
-    )
-    nearest_neighbors, distances = index.get_nns_by_vector(
-        query_feature, num_results, include_distances=True
+    proc_image = preprocess_image(query_image).unsqueeze(0).to(DEVICE)
+    query_feature = model(proc_image)
+    query_feature = query_feature.cpu().detach().numpy()
+    distances, nearest_neighbors = index.search(
+        query_feature,
+        num_results,
     )
-    return query_image, nearest_neighbors, distances
+    return query_image, nearest_neighbors[0], distances[0]
 
 
 @st.cache_data
 def save_embedding(folder=FOLDER):
-    if os.path.isfile(f"{slugify(FOLDER)}.tree"):
+    if os.path.isfile(f"{slugify(FOLDER)}.index"):
+        print("skipping recreating image embeddings")
         return
+    print("Performing image embeddings")
     model = load_model()  # Load the model once
     file_paths = get_all_file_paths(folder_path=folder)
-    images = load_images(file_paths)
-    features = extract_features(images, model)
+    # images = load_images(file_paths)
+    features = extract_features(file_paths, model)
     index = build_annoy_index(features)
-    index.save(f"{slugify(FOLDER)}.tree")
+    faiss.write_index(index, f"{slugify(FOLDER)}.index")
 
 
 def display_image(idx, dist):
@@ -214,11 +260,12 @@ if __name__ == "__main__":
         )
 
         if uploaded_file is not None:
+            query_image = Image.open(uploaded_file)
+            cropped = st_cropper(query_image)
             query_image, nearest_neighbors, distances = search_similar_images(
-                uploaded_file, num_results=n_matches
+                cropped.resize([256, 256]), n_matches
             )
 
-            st.image(query_image.resize([256, 256]), caption="Query Image", width=200)
             st.subheader("Similar Images:")
             cols = st.columns([1] * 5)
             for i, (idx, dist) in enumerate(

requirements.txt

@@ -1,8 +1,10 @@
-annoy
+faiss-cpu
+faiss-gpu
 torch
 torchvision
 streamlit
 tqdm
 python-slugify
 opendatasets
-azure-storage-blob
+azure-storage-blob
+streamlit-cropper