File size: 1,181 Bytes
18f531e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
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
|