Upload handler.py

handler.py

@@ -1,34 +1,27 @@
-from typing import  Dict, List, Any
-from optimum.onnxruntime import ORTModelForSequenceClassification
-from transformers import pipeline, AutoTokenizer
+from typing import Dict, List, Any
+from PIL import Image
+from io import BytesIO
+from transformers import pipeline
+import base64
 
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # load the optimized model
-        model = ORTModelForSequenceClassification.from_pretrained(path)
-        tokenizer = AutoTokenizer.from_pretrained(path)
-        # create inference pipeline
-        self.pipeline = pipeline("zero-shot-image-classification", model=model, tokenizer=tokenizer)
-
-
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        self.pipeline=pipeline("zero-shot-image-classification",model=path)
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-        Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
-                - "label": A string representing what the label/class is. There can be multiple labels.
-                - "score": A score between 0 and 1 describing how confident the model is for this label/class.
+       data args:
+            images (:obj:`string`)
+            candiates (:obj:`list`)
+      Return:
+            A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}
         """
         inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
 
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
-        else:
-            prediction = self.pipeline(inputs)
-        # postprocess the prediction
-        return prediction
+        # decode base64 image to PIL
+        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+
+        # run prediction one image wit provided candiates
+        prediction = self.pipeline(images=[image], candidate_labels=inputs["candiates"])
+        return prediction[0]