alexgk commited on
Commit
dfc0e6f
1 Parent(s): 44d0369

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +26 -0
handler.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from PIL import Image
3
+ from io import BytesIO
4
+ from transformers import pipeline
5
+ import base64
6
+
7
+
8
+ class EndpointHandler():
9
+ def __init__(self, path=""):
10
+ self.pipeline=pipeline("Image-to-Text",model=path)
11
+
12
+ def __call__(self, data: Dict[str, Any]) -> str:
13
+ """
14
+ data args:
15
+ images (:obj:`string`)
16
+ Return:
17
+ A str containing a caption for the text
18
+ """
19
+ inputs = data.pop("inputs", data)
20
+
21
+ # decode base64 image to PIL
22
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
23
+
24
+ # run prediction one image wit provided candiates
25
+ prediction = self.pipeline(images=[image])
26
+ return prediction[0]