Create app.py

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn.functional as F
+import onnxruntime
+import cv2
+from PIL import Image
+
+# 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 = torch.from_numpy(img).unsqueeze(0).float()
+    img.div_(255).sub_(0.5).div_(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.numpy()})[0]
+        return torch.from_numpy(features)
+    
+    # Extract features for each image
+    feat1 = get_features(img1_tensor)
+    feat2 = get_features(img2_tensor)
+    
+    # Normalize features
+    feat1 = F.normalize(feat1, p=2, dim=1)
+    feat2 = F.normalize(feat2, p=2, dim=1)
+    
+    # Calculate cosine similarity
+    cosine_similarity = torch.sum(feat1 * feat2, dim=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")
+    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="Comparison 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
+    )
+
+# Launch the interface
+iface.launch()