| import traceback |
| import numpy as np |
| import torch |
| import gradio as gr |
| import timm |
| import torch.nn.functional as F |
|
|
|
|
| MODEL_NAME = "hf_hub:gaunernst/vit_tiny_patch8_112.arcface_ms1mv3" |
|
|
| model = timm.create_model( |
| MODEL_NAME, |
| pretrained=True |
| ) |
|
|
| model.eval() |
|
|
|
|
| def image_to_tensor(image): |
| image = image.convert("RGB") |
| image = image.resize((112, 112)) |
|
|
| img = np.array(image).astype(np.float32) / 255.0 |
|
|
| img = (img - 0.5) / 0.5 |
|
|
| img = np.transpose(img, (2, 0, 1)) |
|
|
| tensor = torch.tensor(img).unsqueeze(0) |
|
|
| return tensor |
|
|
|
|
| def get_embeddings(image): |
| if image is None: |
| raise ValueError("Please upload an image.") |
|
|
| tensor = image_to_tensor(image) |
|
|
| with torch.no_grad(): |
| embedding = model(tensor) |
|
|
| embedding = F.normalize(embedding, p=2, dim=1) |
|
|
| return embedding |
|
|
|
|
| def compare_faces(image1, image2): |
| try: |
| face1 = get_embeddings(image1) |
| face2 = get_embeddings(image2) |
|
|
| similarity = F.cosine_similarity(face1, face2).item() |
| similarity = round(similarity, 4) |
|
|
| if similarity >= 0.6: |
| result = "THIS IS THE SAME PERSON" |
| else: |
| result = "FACE DOES NOT MATCH" |
|
|
| return result, similarity |
|
|
| except Exception: |
| return traceback.format_exc(), 0 |
|
|
|
|
| demo = gr.Interface( |
| fn=compare_faces, |
| inputs=[ |
| gr.Image(type="pil", label="Upload your first face"), |
| gr.Image(type="pil", label="Upload your second face") |
| ], |
| outputs=[ |
| gr.Textbox(label="Result"), |
| gr.Number(label="Similarity") |
| ], |
| title="Face Recognition App", |
| description="Upload two clear cropped face images and compare them." |
| ) |
|
|
| demo.launch() |
