halimb commited on
Commit
4c50bc8
·
verified ·
1 Parent(s): d2d83fc

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +33 -0
handler.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import pipeline
3
+ import transformers
4
+ from PIL import Image
5
+ import base64
6
+ from io import BytesIO
7
+ from PIL import Image
8
+ import requests
9
+ from transformers import pipeline
10
+
11
+
12
+ print('IN HANDLER...')
13
+ print(transformers.__version__)
14
+
15
+ class EndpointHandler():
16
+ def __init__(self, path=""):
17
+ # load pipe
18
+ self.pipe = pipeline("mask-generation", device = 0, points_per_batch = 256)
19
+
20
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
21
+ base64_image = data.pop("inputs",data)
22
+ if base64_image is None:
23
+ raise ValueError("No image provided")
24
+
25
+ if base64_image.startswith('data:image/jpeg;base64,'):
26
+ base64_image = base64_image.replace('data:image/jpeg;base64,', '')
27
+
28
+ image_bytes = base64.b64decode(base64_image)
29
+ image = Image.open(BytesIO(image_bytes))
30
+
31
+ masks = self.pipe(image, points_per_batch = 256)
32
+
33
+ return depth