handler.py

from __future__ import annotations

import base64
import io
from pathlib import Path

import numpy as np
from PIL import Image

from models import register_ultralytics_modules


class EndpointHandler:
    def __init__(self, path: str = ""):
        register_ultralytics_modules()

        from ultralytics import YOLO

        weights = Path(path) / "weights" / "symbolic_capsule_network_segmentation.pt"
        self.model = YOLO(str(weights))

    def __call__(self, data: dict) -> list[dict]:
        """
        Args:
            data: {"inputs": <PIL Image | bytes | str path>}

        Returns:
            List of dicts compatible with HF image-segmentation pipeline:
            [{"score": float, "label": str, "mask": "<base64 PNG>"}]
        """
        image = data.get("inputs")
        if isinstance(image, bytes):
            image = Image.open(io.BytesIO(image)).convert("RGB")

        results = self.model.predict(image, imgsz=640, conf=0.25, verbose=False)
        r = results[0]

        if r.boxes is None or r.masks is None:
            return []

        h, w = r.orig_shape
        output = []
        for box, mask_tensor in zip(r.boxes, r.masks.data):
            # Resize binary mask back to original image size
            mask_np = (mask_tensor.cpu().numpy() * 255).astype(np.uint8)
            mask_img = Image.fromarray(mask_np).resize((w, h), Image.NEAREST)

            buf = io.BytesIO()
            mask_img.save(buf, format="PNG")
            mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

            output.append({
                "score": round(float(box.conf), 4),
                "label": self.model.names[int(box.cls)],
                "mask": mask_b64,
            })

        return output