|
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]]: |
|
""" |
|
Process an incoming request for zero-shot object detection. |
|
|
|
Args: |
|
data (Dict[str, Any]): The input data containing an encoded image and candidate labels. |
|
|
|
Returns: |
|
A list of dictionaries, each containing a label and its corresponding score. |
|
""" |
|
|
|
inputs = data.get("inputs", {}) |
|
|
|
|
|
image = Image.open(BytesIO(base64.b64decode(inputs['image']))) |
|
|
|
|
|
candidate_labels=inputs["candidates"] |
|
|
|
|
|
detection_results = self.pipeline(image=image, candidate_labels=inputs["candidates"], threshold = 0) |
|
|
|
|
|
return detection_results |
|
|
