49 lines
2.1 KiB
Python
49 lines
2.1 KiB
Python
import io
|
|
from typing import Any, Union
|
|
|
|
from PIL import Image
|
|
|
|
from clearml import StorageManager
|
|
|
|
|
|
# Notice Preprocess class Must be named "Preprocess"
|
|
class Preprocess(object):
|
|
def __init__(self):
|
|
# set internal state, this will be called only once. (i.e. not per request)
|
|
pass
|
|
|
|
def preprocess(self, body: Union[bytes, dict], state: dict, collect_custom_statistics_fn=None) -> Any:
|
|
# we expect to get two valid on the dict x0, and x1
|
|
if isinstance(body, bytes):
|
|
# we expect to get a stream of encoded image bytes
|
|
try:
|
|
Image.open(io.BytesIO(body)).convert("RGB")
|
|
except Exception:
|
|
# value error would return 404, we want to return 500 so any other exception
|
|
raise RuntimeError("Image could not be decoded")
|
|
raise ValueError("send JSON with 'url' (or 'image_url') and 'xywh'")
|
|
|
|
if isinstance(body, dict) and ("url" in body.keys() or "image_url" in body.keys()):
|
|
# image is given as url, and is fetched
|
|
url = body.get("image_url") or body.get("url")
|
|
local_file = StorageManager.get_local_copy(remote_url=url)
|
|
xywh = body.get("xywh")
|
|
if not xywh:
|
|
raise ValueError("body must include 'xywh' as 'x,y,w,h'")
|
|
# PersonAttrService.predict 에 전달 (로컬 캐시 경로 + bbox)
|
|
return {"image_url": local_file, "xywh": str(xywh).strip()}
|
|
|
|
raise ValueError("body must include 'url' (or 'image_url') and 'xywh'")
|
|
|
|
def postprocess(self, data: Any, state: dict, collect_custom_statistics_fn=None) -> dict:
|
|
# post process the data returned from the model inference engine
|
|
# data is the return value from model.predict we will put is inside a return value as Y
|
|
if not isinstance(data, dict):
|
|
# this should not happen
|
|
return dict(output={}, status="FAIL")
|
|
|
|
# data is returned as attribute dict from PersonAttrService.predict
|
|
if "output" in data and "status" in data:
|
|
return data
|
|
return dict(output=data, status="PASS")
|