photo_search.py

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

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)

        data = np.load(data_path, allow_pickle=True)
        self.face_grive_id = data.f.id
        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.face_grive_id[n])
            if self.cosine(face_emb, face_origin) < 0.7: continue

            result.append('get..')

        return result
    
    def cosine(self, x, y):
        return np.dot(x,y)/(norm(x)*norm(y))