|
|
|
import base64 |
|
from io import BytesIO |
|
import os |
|
from typing import Dict, List, Any |
|
import cv2 |
|
import groundingdino |
|
from groundingdino.util.inference import load_model, load_image, predict, annotate |
|
import tempfile |
|
|
|
|
|
HOME = os.getcwd() |
|
|
|
|
|
PACKAGE_HOME = os.path.dirname(groundingdino.__file__) |
|
CONFIG_PATH = os.path.join(PACKAGE_HOME, "config", "GroundingDINO_SwinT_OGC.py") |
|
|
|
class EndpointHandler(): |
|
def __init__(self, path): |
|
|
|
|
|
self.model = load_model(CONFIG_PATH, os.path.join(path, "weights", "groundingdino_swint_ogc.pth")) |
|
|
|
self.box_threshold = 0.35 |
|
self.text_threshold = 0.25 |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
data args: |
|
inputs (:obj: `str` | `PIL.Image` | `np.array`) |
|
kwargs |
|
Return: |
|
A :obj:`list` | `dict`: will be serialized and returned |
|
""" |
|
inputs = data.pop("inputs") |
|
image_base64 = inputs.pop("image") |
|
prompt = inputs.pop("prompt") |
|
|
|
image_data = base64.b64decode(image_base64) |
|
|
|
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=True) as f: |
|
f.write(image_data) |
|
image_source, image = load_image(f.name) |
|
boxes, logits, phrases = predict( |
|
model=self.model, |
|
image=image, |
|
caption=prompt, |
|
box_threshold=self.box_threshold, |
|
text_threshold=self.text_threshold |
|
) |
|
annotated_frame = annotate(image_source=image_source, boxes=boxes, logits=logits, phrases=phrases) |
|
_, annotated_image = cv2.imencode(".jpg", annotated_frame) |
|
annotated_image_b64 = base64.b64encode(annotated_image).decode("utf-8") |
|
num_found = boxes.size(0) |
|
|
|
return [{ |
|
"image": annotated_image_b64, |
|
"prompt": prompt, |
|
"num_found": num_found, |
|
}] |
|
|
