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