Spaces:
Running
Running
import os | |
import json | |
from hashlib import sha1 | |
os.environ["DEEPFACE_HOME"] = "." | |
import pyzipper | |
import numpy as np | |
import gradio as gr | |
from annoy import AnnoyIndex | |
from deepface.commons import functions | |
from deepface.basemodels import Facenet512 | |
# load the face model | |
model = Facenet512.loadModel() | |
input_shape_x, input_shape_y = functions.find_input_shape(model) | |
index = AnnoyIndex(512, "euclidean") | |
index.load(f"face.db") | |
ANNOY_INDEX = json.load(open(f"face.json")) | |
with pyzipper.AESZipFile('persons.zip') as zf: | |
zf.setpassword(b"4321ecafhsats"[::-1]) | |
PERFORMER_DB = json.loads(zf.read('performers.json')) | |
def predict(image, threshold=20.0, results=3): | |
image_array = np.array(image) | |
print(sha1(image_array).hexdigest()) | |
img = functions.preprocess_face( | |
img=image_array, | |
target_size=(input_shape_x, input_shape_y), | |
enforce_detection=True, | |
detector_backend="retinaface", | |
align=True, | |
) | |
img = functions.normalize_input(img, normalization="Facenet2018") | |
face = model.predict(img)[0].tolist() | |
ids, distances = index.get_nns_by_vector( | |
face, 50, search_k=10000, include_distances=True | |
) | |
persons = {} | |
for p, distance in zip(ids, distances): | |
id = ANNOY_INDEX[p] | |
if id in persons: | |
persons[id]["hits"] += 1 | |
persons[id]["distance"] -= 0.5 | |
continue | |
persons[id] = { | |
"id": id, | |
"distance": round(distance, 2), | |
"hits": 1, | |
} | |
if id in PERFORMER_DB: | |
persons[id].update(PERFORMER_DB.get(id)) | |
persons = sorted(persons.values(), key=lambda x: x["distance"]) | |
persons = [p for p in persons if p["distance"] < threshold] | |
return persons[:results] | |
gr.Interface( | |
fn=predict, | |
inputs=[ | |
gr.components.Image(), | |
gr.components.Slider(label="threshold",minimum=0.0, maximum=30.0, value=20.0), | |
gr.components.Slider(label="results", minimum=0, maximum=10, value=3), | |
], | |
outputs=gr.outputs.JSON(label=""), | |
title="Who is in the photo?", | |
description="Upload an image of a person and we'll tell you who it is.", | |
).launch(enable_queue=True, server_name="0.0.0.0") | |