| 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} |
|
|
