Spaces:
Runtime error
Runtime error
File size: 1,077 Bytes
5d3f355 68a20ec 5d3f355 68a20ec 5d3f355 68a20ec 5d3f355 68a20ec 5d3f355 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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))
|