import gradio as gr
import cv2
from insightface.app import FaceAnalysis
import torch

app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))

def calculate(photo):
    image = cv2.imread(photo)
    faces = app.get(image)
    image_draw = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
    image_draw = app.draw_on(image_draw, faces)
    return image_draw
    
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            face_photo = gr.Image(label="Photo", type="filepath")
            greet_btn = gr.Button("Calculate")
        with gr.Column():
            output_image = gr.Image(label="Output")
            output = gr.JSON()
    
    greet_btn.click(fn=calculate, inputs=face_photo, outputs=output_image, api_name="calculate_face_embedding")

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