squanchd commited on
Commit
7a2a4a4
1 Parent(s): 6539990

should work

Browse files
Files changed (3) hide show
  1. __pycache__/handler.cpython-310.pyc +0 -0
  2. handler.py +11 -12
  3. test.py +3 -1
__pycache__/handler.cpython-310.pyc CHANGED
Binary files a/__pycache__/handler.cpython-310.pyc and b/__pycache__/handler.cpython-310.pyc differ
 
handler.py CHANGED
@@ -14,16 +14,15 @@ class EndpointHandler():
14
  self.resnet = InceptionResnetV1(pretrained='vggface2', device=self.device).eval()
15
 
16
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
17
- imageData = data.pop("image", data)
18
- return imageData
19
- # image = Image.open(BytesIO(base64.b64decode(imageData)))
20
- # face_batch = self.mtcnn([image])
21
- # face_batch = [i for i in face_batch if i is not None]
22
- # if face_batch:
23
- # aligned = torch.stack(face_batch)
24
- # if self.device.type == "cuda":
25
- # aligned = aligned.to(self.device)
26
 
27
- # embeddings = self.resnet(aligned).detach().cpu()
28
- # return embeddings.tolist()
29
- # else: return None
 
14
  self.resnet = InceptionResnetV1(pretrained='vggface2', device=self.device).eval()
15
 
16
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
17
+ imageData = data.get("inputs").get("image")
18
+ image = Image.open(BytesIO(base64.b64decode(imageData)))
19
+ face_batch = self.mtcnn([image])
20
+ face_batch = [i for i in face_batch if i is not None]
21
+ if face_batch:
22
+ aligned = torch.stack(face_batch)
23
+ if self.device.type == "cuda":
24
+ aligned = aligned.to(self.device)
 
25
 
26
+ embeddings = self.resnet(aligned).detach().cpu()
27
+ return embeddings.tolist()
28
+ else: return None
test.py CHANGED
@@ -7,6 +7,8 @@ inputs = {
7
  }
8
 
9
 
10
- result = my_handler(inputs)
 
 
11
 
12
  print(result)
 
7
  }
8
 
9
 
10
+ result = my_handler({
11
+ "inputs": inputs
12
+ })
13
 
14
  print(result)