File size: 1,178 Bytes
6933ba4 | 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 | from __future__ import annotations
import time
from typing import Any
MOCK_FAST = "mock-fast@1.0.0"
MOCK_ACCURATE = "mock-accurate@1.0.0"
def list_models() -> list[str]:
return [MOCK_FAST, MOCK_ACCURATE]
def _image_mean_rgb(image) -> list[float]:
rgb = image.convert("RGB")
pixels = list(rgb.getdata())
n = len(pixels)
if n == 0:
return [0.0, 0.0, 0.0]
r = sum(p[0] for p in pixels) / n
g = sum(p[1] for p in pixels) / n
b = sum(p[2] for p in pixels) / n
return [round(r, 2), round(g, 2), round(b, 2)]
def predict(model_version: str, kind: str, payload: Any) -> dict:
if model_version not in list_models():
raise ValueError(f"Unknown model: {model_version}")
if kind == "image":
output: dict[str, Any] = {"mean_rgb": _image_mean_rgb(payload)}
elif kind == "text":
output = {"length": len(payload)}
else:
raise ValueError(f"Unsupported input kind: {kind}")
metadata: dict[str, Any] = {"model_version": model_version}
if model_version == MOCK_ACCURATE:
time.sleep(0.2)
metadata["confidence"] = 0.95
return {"output": output, "metadata": metadata}
|