File size: 1,209 Bytes
499d14f
c0f3d61
0d90fff
499d14f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2cb0016
 
499d14f
 
 
 
 
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
import faiss
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:
  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()