philschmid commited on
Commit
58d0174
·
1 Parent(s): e048781

Create new file

Browse files
Files changed (1) hide show
  1. pipeline.py +24 -0
pipeline.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from PIL import Image
3
+ from io import BytesIO
4
+ from transformers import pipeline
5
+
6
+
7
+ class PreTrainedPipeline():
8
+ def __init__(self, path=""):
9
+ self.pipeline=pipeline("zero-shot-image-classification",model=path)
10
+
11
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
12
+ """
13
+ data args:
14
+ images (:obj:`PIL.Image`)
15
+ candiates (:obj:`list`)
16
+ Return:
17
+ A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}
18
+ """
19
+ # decode base64 image to PIL
20
+ image = Image.open(BytesIO(base64.b64decode(data['image'])))
21
+
22
+ # run prediction one image wit provided candiates
23
+ prediction = self.pipeline(images=[image], candidate_labels=data["candiates"])
24
+ return prediction[0]