File size: 791 Bytes
f954ee9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | """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})")
|