Thomasboosinger commited on
Commit
42e4d7b
1 Parent(s): 3dfe499

Upload 2 files

Browse files
Files changed (2) hide show
  1. handler.py +35 -0
  2. requirements.txt +1 -0
handler.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ from PIL import Image
3
+ from io import BytesIO
4
+ import base64
5
+ from typing import Dict, List, Any
6
+
7
+ class EndpointHandler():
8
+ def __init__(self, model_path=""):
9
+ # Initialize the pipeline with the specified model and set the device to GPU
10
+ self.pipeline = pipeline(task="zero-shot-object-detection", model=model_path, device=0)
11
+
12
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
13
+ """
14
+ Process an incoming request for zero-shot object detection.
15
+
16
+ Args:
17
+ data (Dict[str, Any]): The input data containing an encoded image and candidate labels.
18
+
19
+ Returns:
20
+ A list of dictionaries, each containing a label and its corresponding score.
21
+ """
22
+ # Correctly accessing the 'inputs' key and fixing the typo in 'candidates'
23
+ inputs = data.get("inputs", {})
24
+
25
+ # Decode the base64 image to a PIL image
26
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
27
+
28
+ # Get candidate labels
29
+ candidate_labels=inputs["candidates"]
30
+
31
+ # Correctly passing the image and candidate labels to the pipeline
32
+ detection_results = self.pipeline(image=image, candidate_labels=inputs["candidates"], threshold = 0)
33
+
34
+ # Adjusting the return statement to match the expected output structure
35
+ return detection_results
requirements.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ transformers==4.37.2