| import numpy as np | |
| import os | |
| import torch | |
| from facenet_pytorch import MTCNN, InceptionResnetV1 | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class FacialProcessing: | |
| def __init__(self): | |
| # 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_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 | |