"""Classify the damage type of one or more PDF files with the ONNX model.

Usage: python predict.py file1.pdf [file2.pdf ...]
"""
import json
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort

from pdf_features import extract_features

here = Path(__file__).parent
meta = json.loads((here / "labels.json").read_text())
sess = ort.InferenceSession(str(here / "pdf_damage_classifier.onnx"), providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name

for path in sys.argv[1:]:
    x = np.array([extract_features(Path(path).read_bytes())], dtype=np.float32)
    label_idx, probs = sess.run(None, {input_name: x})
    label = meta["classes"][int(label_idx[0])]
    print(f"{Path(path).name}: {label} (p={float(probs[0][int(label_idx[0])]):.2f})")