insightface-sim / app.py
keyishen's picture
Create app.py
e6321df verified
raw
history blame
1.44 kB
import gradio as gr
import cv2
from insightface.app import FaceAnalysis
import torch
import torch.nn as nn
import numpy as np
np.int = np.int32
np.float = np.float64
np.bool = np.bool_
app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))
def calculate(photo1, photo2):
if photo1 is None or photo2 is None:
return 0
image1 = cv2.imread(photo1)
faces1 = app.get(image1)
if len(faces1) < 1:
return 0
faceid_embeds1 = torch.from_numpy(faces1[0].normed_embedding).unsqueeze(0)
image2 = cv2.imread(photo2)
faces2 = app.get(image2)
if len(faces2) < 1:
return 0
faceid_embeds2 = torch.from_numpy(faces2[0].normed_embedding).unsqueeze(0)
cos = nn.CosineSimilarity(dim=1, eps=1e-10)
cos_similarity = cos(faceid_embeds1, faceid_embeds2)
return cos_similarity
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
face_photo1 = gr.Image(label="Photo of Person 1", type="filepath")
face_photo2 = gr.Image(label="Photo of Person 2", type="filepath")
greet_btn = gr.Button("Calculate")
with gr.Column():
output = gr.JSON()
greet_btn.click(fn=calculate, inputs=[face_photo1, face_photo2], outputs=output, api_name="calculate_face_similarity")
if __name__ == "__main__":
demo.launch()