Upload 21 files

images.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ba21a3d4aa2f70fc5d1834325a60d9ae12f3d9362a01ff17479bcb5017752d9a
-size 651946218
+oid sha256:91d47f3d7ab1beeaf5618e6dc1a8a6ce0047e468fed1ffdaa9995651e29b1a98
+size 22677465

model/db_vit_b_16.index

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e6a435671e9e822062b51f2a1217472f2dac30bdc8443cb5d84001fe2d9c3ff9
-size 13268013
+oid sha256:4a689e3a0812adb264cc25d857a73292fd0d71ded0b73841117a730ee9c93a75
+size 98349

src/build_vector_database.py

@@ -22,21 +22,21 @@ import os
 from modules import FeatureExtractor
 from config import *
 
-
+images_dir = "../webp_images/images"
+data_dir  = "../webp_images"
 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))
+    image_list = sorted(os.listdir(images_dir))
 
     with torch.no_grad():
         # iterate over the images and add their extracted features to the index
         for img_filename in tqdm(image_list):
             # load image
-            img = PIL.Image.open(os.path.join(IMAGES_DIR, img_filename)).convert("RGB")
-            # extract features
+            img = PIL.Image.open(os.path.join(images_dir, img_filename)).convert("RGB")
             output = feature_extractor.extract_features(img)
             # keep only batch dimension
             output = output.view(output.size(0), -1)
@@ -47,7 +47,7 @@ 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(data_dir, f"db_{args.feat_extractor}.index")
     faiss.write_index(index, index_filepath)