File size: 1,748 Bytes
8795764 | 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 | 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
|