"""Minimal inference example for the Palworld terminal pal classifier.

License: CC BY-NC 4.0. Palworld assets © Pocketpair, Inc.
"""
import json
import numpy as np
import onnxruntime as ort
from PIL import Image

MODEL = "pal_classifier.onnx"
LABELS = "labels.json"
ZH = "zh_names.json"


def classify(cell_image_path, model=MODEL):
    sess = ort.InferenceSession(model, providers=["CPUExecutionProvider"])
    labels = json.load(open(LABELS, encoding="utf-8"))["labels"]
    zh = json.load(open(ZH, encoding="utf-8"))
    img = Image.open(cell_image_path).convert("RGB").resize((96, 96), Image.LANCZOS)
    x = np.asarray(img, np.float32) / 255.0
    x = ((x - 0.5) / 0.5).transpose(2, 0, 1)[None]
    logits = sess.run(None, {sess.get_inputs()[0].name: x})[0][0]
    idx = int(logits.argmax())
    p = np.exp(logits - logits.max())
    p /= p.sum()
    return labels[idx], float(p[idx]), zh.get(labels[idx], "")


if __name__ == "__main__":
    import sys
    path = sys.argv[1] if len(sys.argv) > 1 else "cell.png"
    label, conf, name = classify(path)
    print(f"{label} ({name}) conf={conf:.3f}")