Spaces:
Running
Running
| import numpy as np | |
| import os | |
| import torch | |
| from facenet_pytorch import MTCNN, InceptionResnetV1 | |
| import cv2 | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class FacialProcessing: | |
| def __init__(self): | |
| self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') | |
| self.model = cv2.dnn.readNetFromTorch('openface.nn4.small2.v1.t7') | |
| # Set the cache directory to a writable location | |
| os.environ['TORCH_HOME'] = '/tmp/.cache/torch' | |
| self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| self.mtcnn = MTCNN(keep_all=True, device=device) | |
| self.resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device) | |
| def extract_embeddings(self, image_path): | |
| try: | |
| image = cv2.imread(image_path) | |
| if image is None: | |
| logger.error(f"Failed to load image: {image_path}") | |
| return None | |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| faces = self.face_cascade.detectMultiScale(gray, 1.3, 5) | |
| if len(faces) == 0: | |
| logger.warning(f"No face detected in image: {image_path}") | |
| return None | |
| (x, y, w, h) = faces[0] | |
| face = image[y:y+h, x:x+w] | |
| faceBlob = cv2.dnn.blobFromImage(face, 1.0 / 255, (96, 96), (0, 0, 0), swapRB=True, crop=False) | |
| self.model.setInput(faceBlob) | |
| vec = self.model.forward() | |
| return vec.flatten().tolist() | |
| except Exception as e: | |
| logger.error(f"An error occurred while extracting embeddings: {e}") | |
| return None | |
| def extract_embeddings_vgg(self, image): | |
| try: | |
| # Preprocess the image | |
| preprocessed_image = self.mtcnn(image) | |
| if preprocessed_image is None: | |
| logger.warning(f"No face detected in image") | |
| return None | |
| # Extract the face embeddings | |
| embeddings = self.resnet(preprocessed_image.unsqueeze(0)).detach().cpu().numpy().tolist() | |
| if embeddings: | |
| return embeddings[0] | |
| except Exception as e: | |
| logger.error(f"An error occurred while extracting embeddings: {e}") | |
| return None | |