File size: 1,583 Bytes
bbb90ac 03dec79 bbb90ac 03dec79 bbb90ac 03dec79 bbb90ac 03dec79 bbb90ac 03dec79 bbb90ac 03dec79 bbb90ac 03dec79 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from transformers import pipeline
from PIL import Image
from io import BytesIO
import base64
from typing import Dict, List, Any
class EndpointHandler():
def __init__(self, model_path=""):
# Initialize the zero-shot object detection pipeline with the specified model
# and set the device to GPU for faster computation.
self.pipeline = pipeline(task="zero-shot-object-detection", model=model_path, device=0)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Handles incoming requests for zero-shot object detection, decoding the image
and predicting labels based on provided candidates.
Args:
data (Dict[str, Any]): The input data containing an encoded image and candidate labels.
Returns:
List[Dict[str, Any]]: Predictions with labels and scores for the detected objects.
"""
# Decode the base64-encoded image to a PIL Image object for processing.
image_data = data.get("inputs", {}).get('image', '')
image = Image.open(BytesIO(base64.b64decode(image_data)))
# Extract candidate labels from the input data.
candidate_labels = data.get("inputs", {}).get("candidates", [])
# Perform zero-shot object detection using the provided image and candidate labels.
detection_results = self.pipeline(image=image, candidate_labels=candidate_labels)
# Return the detection results directly, which should match the expected output structure.
return detection_results
|