Spaces:
Sleeping
Sleeping
File size: 1,601 Bytes
675d710 |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import os
import faiss
import gradio as gr
from helpers import *
detector = load_detector()
model = load_model()
source_imgs = []
for r, _, f in os.walk(os.getcwd() + "/images"):
for file in f:
if (
(".jpg" in file.lower())
or (".jpeg" in file.lower())
or (".png" in file.lower())
):
exact_path = r + "/" + file
source_imgs.append(exact_path)
source_faces = []
for img in source_imgs:
try:
faces, id = extract_faces(detector, img)
source_faces.append(faces[id])
except Exception as e:
print(f"Skipping {img}, {e}")
source_embeddings = get_embeddings(model, source_faces)
def find_names(image):
imgs, _ = extract_faces(detector, image)
for i, face in enumerate(imgs):
if(face.size[0] * face.size[1] < 1000):
del imgs[i]
embeds = get_embeddings(model, imgs)
d = np.zeros((len(source_embeddings), len(embeds)))
for i, s in enumerate(source_embeddings):
for j, t in enumerate(embeds):
d[i][j] = findCosineDistance(s, t)
ids = np.argmin(d, axis = 0)
names = []
for i in ids:
names.append(source_imgs[i].split("/")[-1].split(".")[0])
recognition(imgs, ids, names, source_faces, d, source_imgs)
return ",".join(names), "Recognition.jpg"
demo = gr.Interface(
find_names,
gr.Image(type="filepath"),
["text", gr.Image(type = "filepath")],
examples = [
os.path.join(os.path.dirname(__file__), "examples/group1.jpg"),
os.path.join(os.path.dirname(__file__), "examples/group2.jpg")
]
)
if __name__ == "__main__":
demo.launch() |