|
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=""): |
|
|
|
|
|
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. |
|
""" |
|
|
|
image_data = data.get("inputs", {}).get('image', '') |
|
image = Image.open(BytesIO(base64.b64decode(image_data))) |
|
|
|
|
|
candidate_labels = data.get("inputs", {}).get("candidates", []) |
|
|
|
|
|
detection_results = self.pipeline(image=image, candidate_labels=candidate_labels) |
|
|
|
|
|
return detection_results |
|
|
