fashxp commited on
Commit
ae0228d
1 Parent(s): 4f4e104

added handler

Browse files
Files changed (2) hide show
  1. .gitignore +3 -0
  2. handler.py +29 -0
.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ # PhpStorm / IDEA
2
+ .idea
3
+
handler.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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="openai/clip-vit-large-patch14-336")
11
+
12
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
13
+ """
14
+ data args:
15
+ image (:obj:`string`)
16
+ parameters (:obj:)
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
+ # decode base64 image to PIL
22
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
23
+
24
+ parameters = inputs['parameters']
25
+ candidate_labels = parameters['candidate_labels']
26
+
27
+ # run prediction one image wit provided candiates
28
+ prediction = self.pipeline(images=[image], candidate_labels=candidate_labels)
29
+ return prediction[0]