import gradio as gr
from src.model import ImageEmbedding

image_embedding = ImageEmbedding()

def compare_faces(person_image, id_image):
    embeddings1, img1 = image_embedding.get_face_embedding(person_image)
    embeddings2, img2 = image_embedding.get_face_embedding(id_image)
    similarity = image_embedding.cosine_similarity(embeddings1, embeddings2)
    return similarity, img1, img2

with gr.Blocks() as demo:
    gr.Markdown('# Image matcher')
    with gr.Row():
        with gr.Column():
            person_image = gr.Image(label="Person Image", type="numpy", height=256)
            sim_score = gr.Textbox(label="Similarity Score")

        with gr.Column():
            id_image = gr.Image(label="ID Image", type="numpy", height=256)
            submit_button = gr.Button("Submit")

    submit_button.click(fn=compare_faces,
                        inputs=[person_image, id_image],
                        outputs=[sim_score, person_image, id_image])


if __name__ == "__main__":
    demo.launch()