File size: 678 Bytes
2519bba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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