Update handler.py

handler.py

@@ -1,14 +1,15 @@
 from typing import Dict, List, Any
 from PIL import Image
 from io import BytesIO
-from transformers import pipeline
+from transformers import AutoProcessor, OmDetTurboForObjectDetection
 import base64
 
 
 class EndpointHandler():
     def __init__(self, path=""):
-       self.pipeline=pipeline("zero-shot-object-detection",model=path)
-        
+        self.processor = AutoProcessor.from_pretrained("Blueway/inference-endpoint-zero-shot-image-classification")
+        self.model = OmDetTurboForObjectDetection.from_pretrained("Blueway/inference-endpoint-zero-shot-image-classification")
+       
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
        data args:
@@ -23,5 +24,16 @@ class EndpointHandler():
         image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
 
         # run prediction one image wit provided candiates
-        prediction = self.pipeline(image=[image], candidate_labels=inputs["candidates"])
-        return prediction[0]
+        inputs = processor(image, text=inputs["candidates"], return_tensors="pt")
+        outputs = model(**inputs)
+        results = processor.post_process_grounded_object_detection(
+            outputs,
+            classes=classes,
+            target_sizes=[image.size[::-1]],
+            score_threshold=0.3,
+            nms_threshold=0.3,
+        )[0]
+        return results
+        #prediction = self.pipeline(image=[image], candidate_labels=inputs["candidates"])
+        #return prediction[0]
+