File size: 2,379 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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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
|