train.py

"""Train the PDF damage classifier and export it to ONNX.

Usage: python train.py [data_dir] [out_dir]
"""
import csv
import json
import sys
from pathlib import Path

import numpy as np
from skl2onnx import to_onnx
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import train_test_split

from pdf_features import FEATURE_NAMES, extract_features

data_dir = Path(sys.argv[1] if len(sys.argv) > 1 else "data")
out_dir = Path(sys.argv[2] if len(sys.argv) > 2 else ".")

X, y = [], []
with open(data_dir / "labels.csv", newline="") as f:
    for row in csv.DictReader(f):
        X.append(extract_features((data_dir / row["file"]).read_bytes()))
        y.append(row["label"])
X = np.array(X, dtype=np.float32)
classes = sorted(set(y))
y_idx = np.array([classes.index(c) for c in y], dtype=np.int64)

X_tr, X_te, y_tr, y_te = train_test_split(X, y_idx, test_size=0.25, stratify=y_idx, random_state=0)
clf = RandomForestClassifier(n_estimators=60, max_depth=8, random_state=0).fit(X_tr, y_tr)
pred = clf.predict(X_te)

report = classification_report(y_te, pred, target_names=classes, digits=3)
print(report)
print(confusion_matrix(y_te, pred))

onx = to_onnx(clf, X_tr[:1], options={id(clf): {"zipmap": False}}, target_opset=15)
(out_dir / "pdf_damage_classifier.onnx").write_bytes(onx.SerializeToString())
(out_dir / "labels.json").write_text(json.dumps({"classes": classes, "features": FEATURE_NAMES}, indent=2))
(out_dir / "metrics.txt").write_text(report)
print("saved pdf_damage_classifier.onnx")