philschmid HF staff commited on
Commit
592b634
1 Parent(s): ceda65c

Create new file

Browse files
Files changed (1) hide show
  1. handler.py +27 -0
handler.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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("zero-shot-image-classification",model=path)
11
+
12
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
13
+ """
14
+ data args:
15
+ images (:obj:`PIL.Image`)
16
+ candiates (:obj:`list`)
17
+ Return:
18
+ A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}
19
+ """
20
+ inputs = data.pop("inputs", data)
21
+
22
+ # decode base64 image to PIL
23
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
24
+
25
+ # run prediction one image wit provided candiates
26
+ prediction = self.pipeline(images=[image], candidate_labels=inputs["candiates"])
27
+ return prediction[0]