Spaces:
Sleeping
Sleeping
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 | |