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

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
print('running device: ', device)

class FacenetPytorch:

    def __init__(self):
        self.face_detector = FACE_DETECTOR(keep_all=True, device=device, margin=20, thresholds=[0.6, 0.9, 0.9])
        self.resnet = InceptionResnetV1(pretrained='vggface2').eval()
    
    def get_embedding(self, image):
        faces = self.extract_faces(image)
        assert len(faces) == 1, f"total face should be 1, instead of {len(faces)}"
        return np.squeeze(self.resnet(faces[0]).detach().numpy())

    def extract_faces(self, image):
        result = []
        faces = self.face_detector(image)
        if faces is not None:
            print('total face: ', len(faces))
            for face in faces:
                result.append(face.unsqueeze(0))
        return result