PoCInnovation / face_compare.py
Valentin Peron
feat(face_compare): Add face comparison functionality (NO FHE)
7468932
from facenet_pytorch import MTCNN, InceptionResnetV1
import torchvision.transforms as transforms
from PIL import Image
import torch
mtcnn = MTCNN(keep_all=False, device='cpu')
model = InceptionResnetV1(pretrained='vggface2').eval()
def preprocess_face(base64_img):
img = Image.open(base64_img).convert('RGB')
preprocess = transforms.Compose([
transforms.Resize((160, 160)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# Detect face and get the bounding box
box, _ = mtcnn.detect(img)
if box is not None:
# Crop the image using the bounding box
img_cropped = img.crop(box[0]) # box[0] contains (x1, y1, x2, y2)
# Apply the rest of preprocessing (resize, tensor, normalize)
img_preprocessed = preprocess(img_cropped).unsqueeze(0)
return img_preprocessed
else:
print("No face detected")
return None
def compare_faces(base64_img1, base64_img2):
# Load and detect face, then crop the image
# Preprocess both images
img1 = preprocess_face(base64_img1) # Replace with your image path
img2 = preprocess_face(base64_img2) # Replace with your image path
# Check if faces were detected in both images
if img1 is not None and img2 is not None:
# Get the embeddings from the FaceNet model
with torch.no_grad():
emb1 = model(img1)
emb2 = model(img2)
# Calculate cosine similarity
cosine_similarity = torch.nn.functional.cosine_similarity(emb1, emb2)
return cosine_similarity.item()
else:
print("Face detection failed on one or both images.")
return None