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