File size: 2,431 Bytes
face243 b10e919 03bedd7 face243 b10e919 355f90d face243 81454e6 face243 03bedd7 face243 03bedd7 face243 03bedd7 |
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 |
import base64
import io
import numpy as np
from PIL import Image
from ultralytics import YOLO
import torch
import os
import time
class EndpointHandler:
def __init__(self, path=""):
model_path = "yolo11n-seg.pt"
# Select device
if torch.backends.mps.is_available():
self.device = "mps"
elif torch.cuda.is_available():
self.device = "cuda"
else:
self.device = "cpu"
# Load YOLO model
self.model = YOLO(model_path)
print(f"✅ Model loaded on {self.device}")
def __call__(self, data):
"""
Expected input:
{
"inputs": <base64 image string>
}
"""
if "inputs" not in data:
return {"error": "No input image provided"}
# Decode image
img_b64 = data["inputs"]
img_bytes = base64.b64decode(img_b64)
image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
image = np.array(image)
t0 = time.perf_counter()
# Run inference
results = self.model(image, device=self.device, verbose=False)[0]
elapsed = time.perf_counter() - t0
detections = []
if results.masks is not None:
for box, mask, conf, cls, poly in zip(
results.boxes.xyxy,
results.masks.data,
results.boxes.conf,
results.boxes.cls,
results.masks.xy
):
x1, y1, x2, y2 = map(int, box.tolist())
poly_np = poly.tolist() # polygon points
detections.append({
"box": [x1, y1, x2, y2],
"confidence": float(conf),
"class_id": int(cls),
"polygon": poly_np
# Optional: "mask": mask.cpu().numpy().astype(np.uint8).tolist()
})
else:
for box, conf, cls in zip(results.boxes.xyxy,
results.boxes.conf,
results.boxes.cls):
x1, y1, x2, y2 = map(int, box.tolist())
detections.append({
"box": [x1, y1, x2, y2],
"confidence": float(conf),
"class_id": int(cls),
"polygon": None
})
return {"detections": detections,"model_latency":elapsed} |