"""Classify an image with the INT8 ONNX model in this repo.

Usage: python run_classify.py <image.jpg>
Deps:  pip install onnxruntime numpy pillow
"""

import glob
import os
import sys

import numpy as np
from PIL import Image
import onnxruntime as ort

HERE = os.path.dirname(os.path.abspath(__file__))
MODEL = sorted(glob.glob(os.path.join(HERE, "*.onnx")))[0]
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

sess = ort.InferenceSession(MODEL, providers=["CPUExecutionProvider"])
size = sess.get_inputs()[0].shape[-1]  # static square input baked into the model

img = Image.open(sys.argv[1]).convert("RGB").resize((size, size), Image.BICUBIC)
x = ((np.asarray(img, np.float32) / 255.0 - MEAN) / STD).transpose(2, 0, 1)[None]
logits = sess.run(None, {"input": x.astype(np.float32)})[0][0]
top5 = np.argsort(logits)[-5:][::-1]
print(f"model: {os.path.basename(MODEL)} (input {size}x{size})")
for i in top5:
    print(f"  class {i}: logit {logits[i]:.3f}")
print("Class indices follow the standard ImageNet-1K sorted-synset order.")