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 | |