Upload 22 files

app.py

@@ -3,13 +3,13 @@ import torch
 import faiss
 import base64
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from io import BytesIO
 from src.modules import FeatureExtractor
-from src.config import DATA_DIR
 from fastapi.middleware.cors import CORSMiddleware
 import zipfile
+from pydantic import BaseModel, Field
 
 app = FastAPI(docs_url="/")
 origins = ["*"]
@@ -22,16 +22,13 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 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.backends.mps.is_built():
     torch.set_default_device("mps")
 
 
-# Helper function to convert base64 string to PIL image
 def base64_to_image(base64_str: str) -> Image.Image:
     try:
         image_data = base64.b64decode(base64_str)
@@ -41,14 +38,12 @@ def base64_to_image(base64_str: str) -> Image.Image:
         raise HTTPException(status_code=400, detail="Invalid Base64 image")
 
 
-# Helper function to convert PIL image to base64 string
 def image_to_base64(image: Image.Image) -> str:
     buffered = BytesIO()
     image.save(buffered, format="JPEG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
-# Helper function to convert PIL image to base64 string
 def image_to_base64(image: Image.Image) -> str:
     buffered = BytesIO()
     image.save(buffered, format="JPEG")
@@ -56,24 +51,26 @@ def image_to_base64(image: Image.Image) -> str:
 
 
 def unzip_folder(zip_file_path, extract_to_path):
-    # Check if the zip file exists
     if not os.path.exists(zip_file_path):
         raise FileNotFoundError(f"Zip file not found: {zip_file_path}")
-
-    # Unzip the folder
     with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
-        zip_ref.extractall(extract_to_path)
+        for member in zip_ref.infolist():
+            filename = member.filename.encode("cp437").decode("utf-8")
+            extracted_path = os.path.join(extract_to_path, filename)
+            os.makedirs(os.path.dirname(extracted_path), exist_ok=True)
+            with zip_ref.open(member) as source, open(extracted_path, "wb") as target:
+                target.write(source.read())
         print(f"Extracted all files to: {extract_to_path}")
 
 
-# Example usage
-zip_file = "./images.zip"  # Replace with your zip file path
-extract_path = "./data"  # Replace with the directory you want to extract to
+zip_file = "./images.zip"
+extract_path = "./data"
 unzip_folder(zip_file, extract_path)
 
 
-image_dir = "./data/images"
-from pydantic import BaseModel, Field
+def is_image_file(filename):
+    valid_extensions = (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff", ".webp")
+    return filename.lower().endswith(valid_extensions)
 
 
 class ImageSearchBody(BaseModel):
@@ -83,29 +80,18 @@ class ImageSearchBody(BaseModel):
 @app.post("/search-image/")
 async def search_image(body: ImageSearchBody):
     try:
-        # Convert the Base64 string to an image
         image = base64_to_image(body.base64_image)
-
-        # Extract features from the image
         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)
         print(D, I)
-
-        image_list = sorted(os.listdir(image_dir))
-
-        # Load the matched image from the directory
-
-        image_list = sorted(os.listdir(image_dir))
+        image_list = sorted([f for f in os.listdir(extract_path) if is_image_file(f)])
+        print(image_list)
         image_name = image_list[int(I[0][0])]
-        matched_image_path = f"{image_dir}/{image_list[int(I[0][0])]}"
+        matched_image_path = f"{extract_path}/{image_list[int(I[0][0])]}"
         matched_image = Image.open(matched_image_path)
-
-        # Convert the matched image to Base64 string
         matched_image_base64 = image_to_base64(matched_image)
 
         return JSONResponse(

images.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0c65c754c9eb4694987f102d2fb9d1b957bfd4dcf44d5e1dbfb4b2e40e590fee
-size 29598354
+oid sha256:cb42184dd971df46852b4c7b7ae6b5a2891abdc4a39006e83923245ae7b5e66b
+size 29594676

model/db_vit_b_16.index

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e26e2231dfeee665f11a89639ef30e94cf47780aa44460353529a15c8f3691b4
+oid sha256:5ca38802b326da27ea0d3fd79c78672da86dd82b61c307d51201092cfaf0c107
 size 276525

src/build_vector_database.py

@@ -1,36 +1,32 @@
-# Description:
-#   This script is used to build the vector database for the images in the dataset.
-#   The script uses the FeatureExtractor class to extract the features from the images and saves them to a Faiss index.
-#
-# Usage:
-#   To use this script, you can run the following commands:
-#       python3 build_vector_database.py
-#       python3 build_vector_database.py --feat_extractor vit_l_32
-#       python3 build_vector_database.py --feat_extractor resnet101
-#
 import faulthandler
+
 faulthandler.enable()
 
 import torch
 from tqdm import tqdm
 import argparse
 import faiss
-import torch
 import PIL
 import os
 
 from modules import FeatureExtractor
 from config import *
 
-images_dir = "../webp_images/images"
-data_dir  = "../webp_images"
+images_dir = "../data"
+model_dir = "../model"
+
+def is_image_file(filename):
+    valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.webp')
+    return filename.lower().endswith(valid_extensions)
+
 def main(args=None):
     # initialize the feature extractor with the base model specified in the arguments
     feature_extractor = FeatureExtractor(base_model=args.feat_extractor)
     # initialize the vector database indexing
     index = faiss.IndexFlatIP(feature_extractor.feat_dims)
-    # get the list of images in sorted order
-    image_list = sorted(os.listdir(images_dir))
+    # get the list of images in sorted order and filter out non-image files
+    image_list = sorted([f for f in os.listdir(images_dir) if is_image_file(f)])
+    # print(image_list)
 
     with torch.no_grad():
         # iterate over the images and add their extracted features to the index
@@ -47,10 +43,9 @@ def main(args=None):
             index.add(output.numpy())
 
     # save the index
-    index_filepath = os.path.join(data_dir, f"db_{args.feat_extractor}.index")
+    index_filepath = os.path.join(model_dir, f"db_{args.feat_extractor}.index")
     faiss.write_index(index, index_filepath)
 
-
 if __name__ == "__main__":
     # parse arguments
     args = argparse.ArgumentParser()
@@ -63,4 +58,4 @@ if __name__ == "__main__":
     args = args.parse_args()
 
     # run the main function
-    main(args)
+    main(args)