File size: 1,040 Bytes
b506301 7a2a4a4 b506301 7a2a4a4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from typing import Dict, List, Any
import torch
import numpy as np
from PIL import Image
from io import BytesIO
import base64
from facenet_pytorch import MTCNN, InceptionResnetV1
class EndpointHandler():
def __init__(self, path=""):
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 __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
imageData = data.get("inputs").get("image")
image = Image.open(BytesIO(base64.b64decode(imageData)))
face_batch = self.mtcnn([image])
face_batch = [i for i in face_batch if i is not None]
if face_batch:
aligned = torch.stack(face_batch)
if self.device.type == "cuda":
aligned = aligned.to(self.device)
embeddings = self.resnet(aligned).detach().cpu()
return embeddings.tolist()
else: return None |