FaceNet / models.py
mdirshad09's picture
Upload 8 files
2519bba
raw
history blame contribute delete
678 Bytes
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