Warlord-K's picture
Update app.py
2cb0016
raw
history blame
1.19 kB
import faiss
detector = get_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:
source_faces.append(extract_faces(detector, img)[0])
source_embeddings = get_embeddings(source_faces)
index = faiss.IndexFlatL2(4096)
index.add(np.array(source_embeddings))
# set image
image = 'group.jpg'
def find_names(image):
imgs = extract_faces(detector, image)
embeds = get_embeddings(imgs)
D, I = index.search(np.array(embeds), 1)
names = ""
for i in I:
names+=source_imgs[i[0]].split("/")[-1].split(".")[0]
names+= "/n"
return "\n".join(names)
demo = gr.Interface(
find_names,
gr.Image(type="filepath"),
"text",
examples = [
os.path.join(os.path.dirname(__file__), "examples/group1.jpeg"),
os.path.join(os.path.dirname(__file__), "examples/group2.jpeg")
]
)
if __name__ == "__main__":
demo.launch()