import torch
import numpy as np
from facenet_pytorch import MTCNN, InceptionResnetV1

class FacenetEmbedder:
    def __init__(self):
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.mtcnn = MTCNN(device=self.device)
        self.resnet = InceptionResnetV1(pretrained='vggface2', device=self.device).eval()

    def detect_face(self, batch):
        faces = self.mtcnn.detect(batch)
        return faces

    def encode(self, batch):
        face_batch = self.mtcnn(batch)
        face_batch = [i for i in face_batch if i is not None]
        if face_batch:
            aligned = torch.stack(face_batch)
            if self.device.type == "cuda": 
                aligned = aligned.to(self.device)

            embeddings = self.resnet(aligned).detach().cpu()
            return embeddings.tolist()
        else: return None