File size: 3,322 Bytes
468a6bc | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Custom handler cho Hugging Face Inference Endpoints.
Đặt file này ở GỐC repo model trên Hub (cùng cấp với config.json). Inference
Endpoints tự phát hiện `handler.py` và dùng class `EndpointHandler` thay cho
pipeline mặc định.
Đầu vào chấp nhận:
- ảnh nhị phân thô (Content-Type: image/jpeg, image/png, ...)
- JSON {"inputs": "<chuỗi base64>", "parameters": {...}}
parameters (tuỳ chọn):
score_thresh : float, mặc định 0.7
return_image : bool, mặc định False — True thì trả thêm ảnh overlay base64
language : "vi" | "en", mặc định "vi"
"""
import base64
import io
from typing import Any, Dict, List, Union
import torch
from PIL import Image
from cardamage import AutoModel
MAX_PIXELS = 40_000_000
class EndpointHandler:
def __init__(self, path: str = ""):
torch.set_grad_enabled(False)
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = AutoModel.from_pretrained(path or ".").to(self.device).eval()
# ------------------------------------------------------------------ helpers
@staticmethod
def _decode(data: Union[bytes, str, Dict[str, Any]]) -> Image.Image:
if isinstance(data, (bytes, bytearray)):
raw = bytes(data)
elif isinstance(data, str):
raw = base64.b64decode(data)
else:
raise ValueError("Không nhận diện được định dạng ảnh đầu vào")
image = Image.open(io.BytesIO(raw))
image.load()
if image.width * image.height > MAX_PIXELS:
raise ValueError(f"Ảnh {image.width}x{image.height} vượt quá {MAX_PIXELS} pixel")
return image
# ------------------------------------------------------------------ entrypoint
def __call__(self, data: Dict[str, Any]) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
params = data.get("parameters") or {}
payload = data.get("inputs", data)
image = self._decode(payload)
self.model.score_thresh = float(params.get("score_thresh", 0.7))
result = self.model.predict(image)
language = params.get("language", "vi")
names = result["labels_en"] if language == "en" else result["labels"]
# Định dạng theo quy ước image-segmentation của HF: label / score / mask
detections = [
{
"label": name,
"score": round(float(score), 4),
"class_id": int(cls),
"box": {
"xmin": round(float(box[0]), 2), "ymin": round(float(box[1]), 2),
"xmax": round(float(box[2]), 2), "ymax": round(float(box[3]), 2),
},
"mask_area_px": int(mask.sum()),
}
for name, score, cls, box, mask in zip(
names, result["scores"], result["classes"], result["boxes"], result["masks"]
)
]
if not params.get("return_image", False):
return detections
overlay = self.model.visualize(image, result, language=language)
buf = io.BytesIO()
overlay.save(buf, format="PNG")
return {
"detections": detections,
"image": base64.b64encode(buf.getvalue()).decode("utf-8"),
}
|