|
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.pop("image", data) |
|
return imageData |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|