File size: 2,888 Bytes
f132af8 330da10 f132af8 c29b42f f132af8 c29b42f f132af8 c29b42f f132af8 c29b42f f132af8 40f1b54 f132af8 6cc7e67 f132af8 830a4c2 36c4050 830a4c2 36c4050 830a4c2 f132af8 c29b42f |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import gradio as gr
import numpy as np
import onnxruntime
import cv2
# Declare ONNX session as a global variable
MODEL_PATH = "weights/Glint360K_R200_TopoFR_9784.onnx"
session = onnxruntime.InferenceSession(MODEL_PATH)
def pil_to_cv2(pil_image):
# Convert PIL Image to CV2 format
numpy_image = np.array(pil_image)
# Convert RGB to BGR
cv2_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
return cv2_image
def process_image(pil_img):
if pil_img is None:
img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.uint8)
else:
# Convert PIL image to CV2
img = pil_to_cv2(pil_img)
img = cv2.resize(img, (112, 112))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.transpose(img, (2, 0, 1))
img = img.astype(np.float32)
img = np.expand_dims(img, axis=0)
img = (img / 255.0 - 0.5) / 0.5
return img
def calculate_similarity(img1, img2):
# Image preprocessing
img1_tensor = process_image(img1)
img2_tensor = process_image(img2)
# Extract features using ONNX model
def get_features(img_tensor):
input_name = session.get_inputs()[0].name
features = session.run(None, {input_name: img_tensor})[0]
return features
# Extract features for each image
feat1 = get_features(img1_tensor)
feat2 = get_features(img2_tensor)
# Normalize features (L2 normalization)
feat1 = feat1 / np.linalg.norm(feat1, axis=1, keepdims=True)
feat2 = feat2 / np.linalg.norm(feat2, axis=1, keepdims=True)
# Calculate cosine similarity
cosine_similarity = np.sum(feat1 * feat2, axis=1).item()
return f"Cosine Similarity: {cosine_similarity:.4f}"
# Create Gradio interface with custom layout
with gr.Blocks() as iface:
gr.Markdown("# Face Recognition with [TopoFR](https://github.com/DanJun6737/TopoFR)")
gr.Markdown("Compare two faces to calculate their cosine similarity.")
with gr.Row():
img1_input = gr.Image(label="Reference Face", type="pil")
img2_input = gr.Image(label="Other Face", type="pil")
with gr.Row():
similarity_output = gr.Text(label="Results")
btn = gr.Button("Compare Faces")
btn.click(
fn=calculate_similarity,
inputs=[img1_input, img2_input],
outputs=similarity_output
)
# Add examples
gr.Examples(
examples=[
["examples/yong1.png", "examples/yong2.png"],
["examples/yong1.png", "examples/yong3.png"],
["examples/yong2.png", "examples/yong3.png"],
["examples/yong1.png", "examples/barboon.jpeg"],
["examples/yong2.png", "examples/barboon.jpeg"],
["examples/yong3.png", "examples/barboon.jpeg"],
],
inputs=[img1_input, img2_input],
label="Example Image Pairs"
)
# Launch the interface
iface.launch() |