Add pdf-damage-classifier ONNX model

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Utilvo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,77 @@
+---
+license: mit
+library_name: onnx
+tags:
+  - onnx
+  - pdf
+  - pdf-repair
+  - tabular-classification
+  - scikit-learn
+datasets:
+  - utilvo/pdf-corruption-lab
+pipeline_tag: tabular-classification
+---
+
+# pdf-damage-classifier
+
+A tiny (75 KB) ONNX model that looks at the **raw bytes** of a PDF and predicts
+what kind of structural damage it has, so a repair tool can choose the right
+recovery strategy. It does not need a PDF parser, so it works on files that are
+too broken to open.
+
+Trained from scratch (scikit-learn `RandomForestClassifier`, exported with
+`skl2onnx`) on synthetic data generated by the scripts in
+[`utilvo/pdf-corruption-lab`](https://huggingface.co/datasets/utilvo/pdf-corruption-lab).
+No third-party models or documents are used.
+
+## Classes
+
+`valid`, `missing_eof`, `invalid_startxref`, `scrambled_xref`,
+`bad_root_reference`, `truncated`
+
+## Input
+
+A float32 vector of 12 structural features computed by
+[`pdf_features.py`](pdf_features.py) (header present, `%%EOF` present,
+`startxref` valid, share of xref entries that point to real objects, trailer
+present, `/Root` object exists, `obj`/`endobj` and `stream`/`endstream` balance,
+log file size). The exact order is in `labels.json`.
+
+## Usage
+
+```bash
+pip install -r requirements.txt
+python predict.py broken.pdf
+# broken.pdf: scrambled_xref (p=0.99)
+```
+
+## Evaluation and limitations (please read)
+
+- Hold-out accuracy on the synthetic test split (300 files) is 100%, and the
+  7 reference files in the dataset repo are all classified correctly.
+- **This number is not evidence of real-world accuracy.** The damage types are
+  generated by simple rules and are mostly separable by the features, so the
+  task is easy by construction.
+- Training files are produced by `pdf-lib` with classic xref tables (no object
+  streams / xref streams) and are small. PDFs produced by other software,
+  encrypted files, incremental updates, and real-world corruption (bit flips,
+  partial overwrites, mixed damage) are **not covered** and will likely be
+  misclassified.
+- Use it as a quick triage hint, not as a validator.
+
+## Reproduce
+
+```bash
+git clone https://huggingface.co/datasets/utilvo/pdf-corruption-lab
+cd pdf-corruption-lab && npm install
+node gen_training.js ../data 200
+cd .. && python train.py data .
+```
+
+## Disclosure
+
+Maintained by the team behind Utilvo (https://utilvo.com/pdf-tools).
+
+## License
+
+MIT.

labels.json

@@ -0,0 +1,24 @@
+{
+  "classes": [
+    "bad_root_reference",
+    "invalid_startxref",
+    "missing_eof",
+    "scrambled_xref",
+    "truncated",
+    "valid"
+  ],
+  "features": [
+    "has_pdf_header",
+    "eof_in_last_1k",
+    "ends_with_eof",
+    "startxref_present",
+    "startxref_numeric",
+    "startxref_points_to_xref",
+    "xref_valid_fraction",
+    "has_trailer",
+    "root_object_exists",
+    "obj_endobj_gap",
+    "stream_endstream_gap",
+    "log10_size"
+  ]
+}

metrics.txt

@@ -0,0 +1,12 @@
+                    precision    recall  f1-score   support
+
+bad_root_reference      1.000     1.000     1.000        50
+ invalid_startxref      1.000     1.000     1.000        50
+       missing_eof      1.000     1.000     1.000        50
+    scrambled_xref      1.000     1.000     1.000        50
+         truncated      1.000     1.000     1.000        50
+             valid      1.000     1.000     1.000        50
+
+          accuracy                          1.000       300
+         macro avg      1.000     1.000     1.000       300
+      weighted avg      1.000     1.000     1.000       300

pdf_damage_classifier.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:038a87bb980f5ce9794c64e9fa4ae2ba670252ad4a43d120605b58fff880ae31
+size 75503

pdf_features.py

@@ -0,0 +1,93 @@
+"""Byte-level structural features used by pdf-damage-classifier.
+
+The same function is used for training and inference, so the ONNX model's
+input vector is defined by FEATURE_NAMES (in this exact order).
+Only the PDF's raw bytes are needed; no PDF parser is involved, so it works
+on files that are too broken to open.
+"""
+import math
+import re
+
+FEATURE_NAMES = [
+    "has_pdf_header",
+    "eof_in_last_1k",
+    "ends_with_eof",
+    "startxref_present",
+    "startxref_numeric",
+    "startxref_points_to_xref",
+    "xref_valid_fraction",
+    "has_trailer",
+    "root_object_exists",
+    "obj_endobj_gap",
+    "stream_endstream_gap",
+    "log10_size",
+]
+
+_XREF_ENTRY = re.compile(rb"(\d{10}) (\d{5}) n")
+_STARTXREF = re.compile(rb"startxref\s+(\S+)")
+_ROOT = re.compile(rb"/Root\s+(\d+)\s+(\d+)\s+R")
+
+
+def extract_features(data: bytes):
+    n = len(data)
+    tail = data[-1024:]
+
+    has_header = float(data[:5] == b"%PDF-")
+    eof_in_tail = float(b"%%EOF" in tail)
+    ends_eof = float(data.rstrip()[-5:] == b"%%EOF")
+
+    m = None
+    for m in _STARTXREF.finditer(tail):
+        pass
+    sx_present = float(m is not None)
+    sx_numeric = 0.0
+    sx_points = 0.0
+    if m is not None:
+        tok = m.group(1)
+        if tok.isdigit():
+            sx_numeric = 1.0
+            off = int(tok)
+            if off < n and data[off : off + 4] == b"xref":
+                sx_points = 1.0
+
+    entries = _XREF_ENTRY.findall(data)
+    if entries:
+        ok = 0
+        for off_s, _gen in entries:
+            off = int(off_s)
+            if off < n and re.match(rb"\d+ \d+ obj", data[off : off + 24]):
+                ok += 1
+        xref_frac = ok / len(entries)
+    else:
+        xref_frac = 0.0
+
+    has_trailer = float(b"trailer" in data)
+
+    rm = _ROOT.search(data)
+    root_exists = 0.0
+    if rm is not None:
+        root_exists = float(re.search(rb"(?:^|\s)" + rm.group(1) + b" " + rm.group(2) + b" obj", data) is not None)
+
+    objs = len(re.findall(rb"\d+ \d+ obj\b", data))
+    endobjs = data.count(b"endobj")
+    # Count stream starts, ignoring the "stream\n" inside "endstream\n".
+    stream_starts = len(re.findall(rb"(?<!end)stream\r?\n", data))
+    endstreams = data.count(b"endstream")
+
+    obj_gap = (objs - endobjs) / max(objs, 1)
+    stream_gap = (stream_starts - endstreams) / max(stream_starts, 1)
+
+    return [
+        has_header,
+        eof_in_tail,
+        ends_eof,
+        sx_present,
+        sx_numeric,
+        sx_points,
+        xref_frac,
+        has_trailer,
+        root_exists,
+        obj_gap,
+        stream_gap,
+        math.log10(max(n, 1)),
+    ]

predict.py

@@ -0,0 +1,23 @@
+"""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})")

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+onnxruntime
+scikit-learn
+skl2onnx
+onnx

train.py

@@ -0,0 +1,42 @@
+"""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")