Upload 161 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model/db_vit_b_16.index filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,16 @@
+# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.11-slim
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Samet Cetin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,10 +1,65 @@
+# image-retrieval
+
+This repository presents an image retrieval pipeline that consists of a vector database storing image features extracted using a pre-trained ResNet/ViT model. Image retrieval pipeline consists of (1) a vector database building script that extracts image features of all images in the collection using a pre-trained ResNet/ViT model and stores them in a vector database by cosine similarity, (2) a search query script that extracts the features for a randomly selected image and returns the most similar images from the vector database.
+
+The construction of the proposed image retrieval pipeline can be divided into three stages:
+
+1 - Place your images (a collection of images that you want to create an image retrieval pipeline on) under the data folders as ```data/images```. In this project, [Landscape Pictures Dataset](https://www.kaggle.com/datasets/arnaud58/landscape-pictures) is used but feel free to use your of own image collection.   
+2 - Extract the features of all images using a pre-trained ResNet model and save all features into a vector database.  
+3 - For a given query image, perform a similarity search by first extracting the features of the query image and then performing a search over the vector database.  
+
+## Installation
+1. Clone the repository to your local first as follows    
+```    
+$ git clone https://github.com/cetinsamet/image-retrieval.git
+```
+
+2. Then recreate the virtual environment on your local using the provided ```.yml``` file and activate the venv as follows    
+```
+$ cd image-retrieval  
+$ conda env create -f environment.yml
+$ conda activate image-retrieval
+```
+
+3. Add the path of your own working directory as WORK_DIR in the configuration file ```src/config/settings.py```as follows
+```
+$ cd src
+$ nano config/settings.py  # or use your favorite text editor and add the path of your own working directory
+
+WORK_DIR = "PATH-TO-YOUR-OWN-WORKING-DIRECTORY"
+```
+
+4. Create a ```data``` folder as follows
+```
+$ mkdir ../data
+```
+
+5. Download data from [here](https://www.kaggle.com/datasets/arnaud58/landscape-pictures) (Landscape Pictures Dataset) and place all image files under ```data/images``` folder.
+
+## Run
+Please perform following steps sequentially to reproduce the image retrieval pipeline that gives the search query examples below.
+
+### Building vector database
+Running ```build_vector_database.py``` extracts the features of all downloaded images by using the selected pre-trained ResNet model (please use ```--feat_extractor``` flag to customize) and save the features into a vector database that is stored under ```data/``` directory.
+
+```    
+$ python3 build_vector_database.py --feat_extractor vit_l_32
+```
+
+### Search query
+Finally, running ```search_query.py``` selects random object images (the amount of the selected images can be customized using ```--n``` flag) and search the most similar object images (the amount of the similar images can be customized using ```--k``` flag) by querying the pre-built vector database. 
+```    
+$ python3 search_query.py --feat_extractor vit_l_32 --n 10 --k 12
+```
+
+Some examples of the resulting search queries are as follows;
+
+![img.png](results/results_vit_l_32/query_001.jpg)
 ---
-title: Image Retrieval
-emoji: 📉
-colorFrom: gray
-colorTo: gray
-sdk: docker
-pinned: false
+![img.png](results/results_vit_l_32/query_002.jpg)
 ---
+![img.png](results/results_vit_l_32/query_005.jpg)
+---
+![img.png](results/results_vit_l_32/query_006.jpg)
+
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,54 @@
+import os
+import torch
+import faiss
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import FileResponse, JSONResponse
+from src.modules import FeatureExtractor
+from src.config import DATA_DIR
+from io import BytesIO
+
+app = FastAPI(docs_url="/")
+
+# Load FAISS index and feature extractor
+index = faiss.read_index("./model/db_vit_b_16.index")
+feature_extractor = FeatureExtractor(base_model="vit_b_16")
+
+# Use MPS (Apple Silicon) if available, otherwise fallback to CPU
+if torch.has_mps:
+    torch.set_default_device("mps")
+
+# Helper function to load image from uploaded file
+def load_image(uploaded_file):
+    image = Image.open(BytesIO(uploaded_file)).convert("RGB")
+    return image
+
+@app.post("/search-image/")
+async def search_image(file: UploadFile = File(...)):
+    try:
+        # Load the uploaded image
+        image = load_image(await file.read())
+
+        # Extract features
+        with torch.no_grad():
+            output = feature_extractor.extract_features(image)
+            output = output.view(output.size(0), -1)
+            output = output / output.norm(p=2, dim=1, keepdim=True)
+
+            # Perform FAISS search for the top 1 similar image
+            D, I = index.search(output.cpu().numpy(), 1)  # Changed from 5 to 1
+
+        # Load the list of image filenames (assuming you have image_list)
+        image_list = sorted(os.listdir(os.path.join(DATA_DIR, "images")))
+
+        # Get the path of the most similar image
+        similar_image_path = os.path.join(DATA_DIR, "images", image_list[int(I[0][0])])
+
+        # Return the image file itself
+        return FileResponse(similar_image_path, media_type="image/jpeg")
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)