from facenet_pytorch import InceptionResnetV1, MTCNN
import torchvision.transforms as transforms

class FaceNetModel:
    def __init__(self):
        self.mtcnn = None
        self.resnet = None
        self.transform = None

    def initialize_model(self):
        self.mtcnn = MTCNN()
        self.resnet = InceptionResnetV1(pretrained='vggface2').eval()
        return self.mtcnn, self.resnet

    def get_transform(self):
        self.transform = transforms.Compose([
            transforms.Resize((250, 250)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        return self.transform