Spaces:
Runtime error
Runtime error
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] | |
aligned = torch.stack(face_batch) | |
if self.device.type == "cuda": | |
aligned = aligned.to(self.device) | |
embeddings = self.resnet(aligned).detach().cpu() | |
return embeddings.tolist() |