handler.py

"""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"),
        }