photo_search.py

from annoy import AnnoyIndex
from face_recognition import FacenetPytorch
import numpy as np 
from numpy.linalg import norm
import json

class PhotoSearch:
    def __init__ (self, tree_path, data_path, index_size = 512):
        self.tree_path = tree_path
        self.data_path = data_path

        self.t = AnnoyIndex(index_size, 'angular')
        self.t.load(tree_path)
        
        self.grive_id = []
        with open(data_path) as f :
            self.grive_id = json.load(f)        
        self.extractor = FacenetPytorch()

    def search(self, image):
        face_emb = self.extractor.get_embedding(image)
        nns = self.t.get_nns_by_vector(face_emb, 20)
        result = []
        for n in nns:
            face_origin = self.t.get_item_vector(n)
            # print(self.grive_id[n])
            if self.cosine(face_emb, face_origin) < 0.7: continue

            result.append(f'https://drive.google.com/uc?export=view&id={self.grive_id[n]}')
           
        return result
    
    def cosine(self, x, y):
        return np.dot(x,y)/(norm(x)*norm(y))