init

handler.py

@@ -0,0 +1,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['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

requirements.txt

@@ -0,0 +1,3 @@
+facenet-pytorch
+torch
+datasets