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()