local_search takes care of searching instead of main

local_search.py

@@ -1,3 +1,4 @@
+from charset_normalizer import utils
 from datastore import Datastore
 from vectorizer_deepface import Vectorizer
 from PIL import Image
@@ -6,79 +7,115 @@
 from os import listdir
 from os.path import isfile, join
 import time
-# from deepface import DeepFace
 import tensorflow as tf
-
-# my_devices = tf.config.experimental.list_physical_devices(device_type='CPU')
-# tf.config.experimental.set_visible_devices(devices= my_devices, device_type='CPU')
-
-
-
+from vectorizer_deepface import Vectorizer
+import utils
+import faiss
 
 
 class Searcher:
     def __init__(self, dir_path) -> None:
         self.dir_path = dir_path
         self.query_image = None
-        self.vec = Vectorizer()
         self.data_store = Datastore()
         self.found = []
         self.results = []        
         self.query_image_path = None
+        self.store = None
+        self.vec = None
+        self.mappings = None
+        self.representations = None
+        self.file_name_maps = None
+        self.names = None
+        self.vectors = None
 
     def input_query_image(self, img_file):
+        self.query_image = img_file
 
-        # self.query_image = Image.open(query_image_path)
-        # opencv_query_Image = cv2.cvtColor(np.array(self.query_image), cv2.COLOR_RGB2BGR)
+    def search(self, query_img):
+        opencvImage = cv2.cvtColor(np.array(query_img), cv2.COLOR_RGB2BGR)
+
+        emb_2 = self.vec.vectorize_single(opencvImage)
+        emb_2 = np.array(emb_2, dtype=np.float32)
+        emb_2 = emb_2.reshape(1, 128)
+
+        D, I = self.store.search(emb_2)
+        ret = dict()
+        rank=0
+        file_names = []
+        for id in I[0]:
+            rank+=1
+            ret[str(rank)] = {str(rank) : {self.names[id]:self.file_name_maps[str(id)]}}
+            file_names.append(self.file_name_maps[str(id)])
+        return ret, file_names
 
-        # self.query_image = np.array(self.vec.vectorize_single(opencv_query_Image), dtype = np.float32)
-        # self.query_image = self.query_image.reshape(1, 128)
-        self.query_image = img_file
 
     def initialize_data_store(self):
-        onlyfiles = [f for f in listdir(self.dir_path) if isfile(join(self.dir_path, f))]
-        self.found = []
-        err_count = 0
-        count = 0 
-        for val in onlyfiles:
-            try :
-                image = Image.open(self.dir_path+"/"+val)
+        self.vec = Vectorizer()
+        self.vec.vectorize_lfw("data/lfw-deepfunneled/lfw-deepfunneled")
+
+        self.store = Datastore()
+        self.mappings  = utils.get_mappings()
+        self.representations = utils.get_representations()
+        self.file_name_maps = utils.get_file_name_mappings()
+        self.names, self.vectors = utils.get_vectors()
+
+        idx = 0
+        for name, vector in zip(self.names, self.vectors):
+            name = np.array(name)
+            vector = np.array(vector)
+            # print(idx)
+            # print(vector.shape)
+            # print(name)
+
+            self.store.add(vector, idx, name)
+            idx+=1
+
+        faiss.write_index(self.store.index, "vecdata/vector.index")
+
+        # onlyfiles = [f for f in listdir(self.dir_path) if isfile(join(self.dir_path, f))]
+        # self.found = []
+        # err_count = 0
+        # count = 0 
+        # for val in onlyfiles:
+        #     try :
+        #         image = Image.open(self.dir_path+"/"+val)
 
 
-                opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        #         opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-                embedding = np.array(self.vec.vectorize_single(opencvImage), dtype=np.float32)
-                embedding = embedding.reshape(1, 128)
-                self.data_store.add(embedding, count, "")
+        #         embedding = np.array(self.vec.vectorize_single(opencvImage), dtype=np.float32)
+        #         embedding = embedding.reshape(1, 128)
+        #         self.data_store.add(embedding, count, "")
 
-                self.found.append(self.dir_path+"/"+val)
-                count+=1
+        #         self.found.append(self.dir_path+"/"+val)
+        #         count+=1
 
-            except Exception as e:
-                print(e)
-                err_count+=1
-                # print(err_count)
-                continue
+        #     except Exception as e:
+        #         print(e)
+        #         err_count+=1
+        #         # print(err_count)
+        #         continue
 
-    def search(self):
-        ret = dict()
-        print(type(self.query_image))
-        Distances, Identifiers = self.data_store.search(self.query_image)
-        endtime = time.time()
-        # print(f"time taken {endtime-starttime}")
-        # print(Distances, Identifiers) 
-        # print(link_maps)
+    # def search(self):
+    #     # ret = dict()
+    #     # print(type(self.query_image))
+    #     # Distances, Identifiers = self.data_store.search(self.query_image)
+    #     # endtime = time.time()
+    #     # # print(f"time taken {endtime-starttime}")
+    #     # # print(Distances, Identifiers) 
+    #     # # print(link_maps)
 
-        for val, dis in zip(Identifiers[0], Distances[0]):
-            if val != -1:
-                self.results.append((self.found[val]))
-                print(self.found[val], dis)
-                ret[self.found[val]] = [str(dis)]
+    #     # for val, dis in zip(Identifiers[0], Distances[0]):
+    #     #     if val != -1:
+    #     #         self.results.append((self.found[val]))
+    #     #         print(self.found[val], dis)
+    #     #         ret[self.found[val]] = [str(dis)]
 
-        return ret 
-        # for val in self.results:
-            # metrics = DeepFace.verify(img1_path=self.query_image_path, img2_path=val)
-            # print(metrics)
+    #     # return ret 
+    #     # for val in self.results:
+    #         # metrics = DeepFace.verify(img1_path=self.query_image_path, img2_path=val)
+    #         # print(metrics)
 
 
 

main.py

@@ -14,67 +14,12 @@
 import os
 import io
 from os import getcwd
+from local_search import Searcher
 
 
-store = Datastore()
-vec = Vectorizer()
-# representations, file_name_mappings = vec.vectorize_lfw("data/lfw-deepfunneled/lfw-deepfunneled")
-
-# print(representations)
-# print(file_name_mappings)    
-
-
-mappings = utils.get_mappings()
-representations = utils.get_representations()
-file_name_maps = utils.get_file_name_mappings()
-names, vectors = utils.get_vectors()
-
-print(len(names))
-print(len(vectors))
-
-
-idx = 0
-for name, vector in zip(names, vectors):
-    name = np.array(name)
-    vector = np.array(vector)
-    # print(vector.shape)
-
-    store.add(vector, idx, name)
-    idx+=1
-
-faiss.write_index(store.index, "vecdata/vector.index")
-
-
-
-
-def search(query_img):
-
-    # image = Image.open("test_data/muk.jpg")
-    opencvImage = cv2.cvtColor(np.array(query_img), cv2.COLOR_RGB2BGR)
-
-    emb_2 = vec.vectorize_single(opencvImage) 
-    emb_2 = np.array(emb_2, dtype=np.float32)
-    emb_2 = emb_2.reshape(1, 128)
-    # print(emb_2.shape)
-
-
-
-
-    D, I = store.search(emb_2)
-    # print(D , I)
-    ret = dict()
-    rank=0
-    file_names = []
-    for id in I[0]:
-        # print(id)
-        # print(names[id])
-        # print(file_name_maps[str(id)])
-        rank+=1
-        ret[str(rank)] = {str(rank) : {names[id]:file_name_maps[str(id)]}}
-        file_names.append(file_name_maps[str(id)])
-    return ret, file_names
-
 
+searcher = Searcher("data/lfw-deepfunneled/lfw-deepfunneled")
+searcher.initialize_data_store()
 
 app = FastAPI()
 @app.post("/image/search")
@@ -94,17 +39,11 @@ async def search_api(file: UploadFile = File(...)):
     contents = await file.read()
     image = Image.open(io.BytesIO(contents))
     image.save("query_data/test.jpg")
-    ret, files = search(image) 
+    ret, files = searcher.search(image)
     # downloader = Downloader("query_data/test.jpg")
-
-
-    return ret
 
-@app.get("/")
-def get_file(name_file: str):
-    print(getcwd()+"/"+name_file)
-    return FileResponse(name_file)
 
-if __name__ == "__main__":
-    uvicorn.run(app, debug=True)
+    return ret
 
+if __name__ == "__main__":
+    uvicorn.run(app, debug=True)

utils.py

@@ -2,15 +2,15 @@
 import numpy as np
 from numpy.lib.utils import source
 
-with open("vecdata/representations_deepface.json") as outputFile:
+with open("vecdata/representations_deepface_lfw.json") as outputFile:
         representaions = json.load(outputFile)
         outputFile.close()
 
-with open("vecdata/mappings_deepface.json") as outputFile:
+with open("vecdata/mappings_deepface_lfw.json") as outputFile:
         mappings = json.load(outputFile)
         outputFile.close()
 
-with open("vecdata/file_name_mappings.json") as outputFile:
+with open("vecdata/file_name_mappings_lfw.json") as outputFile:
         file_name_mappings = json.load(outputFile)
         outputFile.close()