Spaces:
Running
Running
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() |