PierreLF38's picture
Update README.md
78d41f4 verified
|
Raw History Blame Contribute Delete
13.9 kB
---
license: apache-2.0
library_name: onnx
pipeline_tag: image-classification
tags:
- onnx
- ocr
- document-ai
- document-layout
- image-quality
- image-classification
datasets:
- allenai/olmOCR-mix-0225
metrics:
- accuracy
- precision
- recall
- f1
---
# Document Gate Keeper v0.5: OCR Crop Quality Classifier (ONNX)
A lightweight binary classifier that tells whether a **document region crop** is legible enough to
be sent to an OCR model. It labels each crop `good` or `bad`.
It is meant to sit between a layout detector and a (vision-language) OCR model. Crops are produced
by a detector such as [`Inceptive/PP-DocLayoutV3-ONNX`](https://huggingface.co/Inceptive/PP-DocLayoutV3-ONNX),
and this model flags the ones that are too blurry, pixelated or compression-damaged. The pipeline can
then skip those regions instead of running expensive inference on them, since unreadable regions
are a common source of hallucinated OCR output.
## What's new in v0.5
- **Apache 2.0 license.** v0.5 is released under [Apache 2.0](LICENSE), so commercial use is now
allowed. v0 was released under CC BY-NC 4.0, for non-commercial use only.
- **FUNSD removed from the training data.** The model is now trained only on crops derived from
[olmOCR-mix-0225](https://huggingface.co/datasets/allenai/olmOCR-mix-0225). FUNSD is available
for non-commercial, research and educational purposes only, and it was what kept v0
non-commercial.
- **CORD dataset added to the training data.**
The previous release, [`Inceptive/document-gate-keeper-v0`](https://huggingface.co/Inceptive/document-gate-keeper-v0),
stays available unchanged. **v0.5 expects different preprocessing**, so do not reuse v0's code (see
[Note for v0 users](#note-for-v0-users)).
| | |
|---|---|
| **Task** | Binary image classification (`bad` / `good`) |
| **Format** | ONNX, opset 17, single file (no external-data sidecar) |
| **Size** | 1.5 MB, 390,370 parameters |
| **Input** | `input_IMAGE`, `float32[N, 3, 640, 640]`, NCHW, RGB, normalized |
| **Output** | `output_0`, `float32[N, 2]`, softmax probabilities, order `[bad, good]` |
| **Runtime** | Any ONNX Runtime execution provider (CPU is fast enough for per-crop use) |
| **Test F1** | 0.9857 on 1,396 held-out crops (see [Evaluation](#evaluation)) |
## Intended use
- **In scope:** crops of printed or scanned document pages (text blocks, titles, tables, form
fields, figures with text), as produced by a document layout detector. Use it to decide whether a
region should be OCR'd or skipped / flagged for review.
- **Out of scope:** natural-scene photos, general-purpose image quality assessment (IQA) or
aesthetic scoring, and full-resolution pages when you need a per-region decision. Also out of
scope: any decision about a document's *content*. The model only judges visual legibility.
## Model inputs and outputs
### Input: `input_IMAGE`
`float32` tensor of shape `[N, 3, 640, 640]`, with a dynamic batch dimension `N`. The image must be
preprocessed **exactly** as during training:
1. **Convert to RGB.** If you load images with OpenCV, convert BGR β†’ RGB first.
2. **Resize to fit.** Resize the image so that its longest side is exactly 640 px, **keeping the
aspect ratio**. Small crops are **upscaled** too.
3. **Center and pad with white.** Paste the image in the center of a 640 Γ— 640 canvas filled with
white, `(255, 255, 255)`. Padding happens in pixel space, *before* normalization.
4. **Scale** pixel values to `[0, 1]` (divide by 255).
5. **Normalize** with the ImageNet statistics: `mean = [0.485, 0.456, 0.406]`,
`std = [0.229, 0.224, 0.225]`, per channel.
6. **Transpose** HWC β†’ CHW and add the batch dimension.
> Do not stretch crops to 640 Γ— 640. A plain resize distorts the aspect ratio, which the model never
> saw during training, and shifts its scores.
The same contract is published in [`config.json`](config.json) (`resize_mode`, `allow_upscale`,
`pad_value`, `pad_position`, `normalization_mean`, `normalization_std`, `classes`, …), and the
snippets below read it from there.
#### Note for v0 users
v0.5's preprocessing differs from v0's in two ways:
| | v0 | v0.5 |
|---|---|---|
| Padding color | black `(0, 0, 0)` | **white `(255, 255, 255)`** |
| Crops smaller than 640 px | never upscaled | **upscaled** so the longest side is 640 px |
Feeding v0.5 with v0's `preprocess()` gives it inputs it was not trained on.
### Output: `output_0`
`float32` tensor of shape `[N, 2]`, holding class probabilities in the order `["bad", "good"]`. The
graph already ends with a `Softmax`, so do **not** apply another one.
The decision rule we use is `P(good) < 0.6 β†’ low quality`. Raise the threshold to skip more regions
(fewer hallucinations, more missed text), or lower it to keep more.
## How to use
Install the dependencies:
```bash
pip install onnxruntime numpy pillow huggingface_hub
```
Classify one crop:
```python
import json
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from PIL import Image
REPO_ID = "Inceptive/document-gate-keeper-v0.5"
# config.json holds the preprocessing contract used during training.
with open(hf_hub_download(repo_id=REPO_ID, filename="config.json")) as f:
config = json.load(f)
SIZE = config["input_size"][0]
PAD_VALUE = tuple(config["pad_value"])
MEAN = np.array(config["normalization_mean"], dtype=np.float32)
STD = np.array(config["normalization_std"], dtype=np.float32)
CLASSES = tuple(config["classes"])
def preprocess(image: Image.Image) -> np.ndarray:
"""PIL image -> float32 [1, 3, 640, 640] tensor, as used during training."""
image = image.convert("RGB")
# Resize so that the longest side is exactly 640 px, keeping the aspect ratio (small crops are upscaled).
w, h = image.size
scale = SIZE / max(w, h)
image = image.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.BILINEAR)
# Center the image on a white 640x640 canvas: padding happens in pixel space, before normalization.
canvas = Image.new("RGB", (SIZE, SIZE), PAD_VALUE)
canvas.paste(image, ((SIZE - image.width) // 2, (SIZE - image.height) // 2))
x = np.asarray(canvas, dtype=np.float32) / 255.0 # HWC, [0, 1]
x = (x - MEAN) / STD # ImageNet normalization
return x.transpose(2, 0, 1)[np.newaxis] # NCHW
model_path = hf_hub_download(repo_id=REPO_ID, filename=config["model_file"])
session = ort.InferenceSession(
model_path,
providers=["CUDAExecutionProvider", "CPUExecutionProvider"], # falls back to CPU
)
image = Image.open("crop.png")
probs = session.run([config["output_name"]], {config["input_name"]: preprocess(image)})[0][0] # already softmaxed
p_bad, p_good = float(probs[0]), float(probs[1])
print(f"P(bad)={p_bad:.3f} P(good)={p_good:.3f} -> {CLASSES[int(probs.argmax())]}")
GOOD_QUALITY_THRESHOLD = config["good_threshold"]
if p_good < GOOD_QUALITY_THRESHOLD:
print("Low-quality crop: skip OCR or flag it for review.")
```
Classify several crops in one call (the batch dimension is dynamic):
```python
paths = ["region_0.png", "region_1.png", "region_2.png"]
batch = np.concatenate([preprocess(Image.open(p)) for p in paths]) # [N, 3, 640, 640]
probs = session.run([config["output_name"]], {config["input_name"]: batch})[0] # [N, 2]
for path, (p_bad, p_good) in zip(paths, probs):
print(path, "good" if p_good >= GOOD_QUALITY_THRESHOLD else "bad", f"(P(good)={p_good:.3f})")
```
Each crop is padded to 640 Γ— 640, so a batch of `N` crops costs `N Γ— 3 Γ— 640 Γ— 640 Γ— 4` bytes of
input (β‰ˆ 4.9 MB per crop). Batch in moderate chunks.
**With a layout detector:** run [`Inceptive/PP-DocLayoutV3-ONNX`](https://huggingface.co/Inceptive/PP-DocLayoutV3-ONNX)
on the page, crop each detected box from the page image at its original resolution, and pass each
crop through `preprocess` above. Do not run the classifier on the detector's own resized input.
## Evaluation
The model was evaluated on a held-out test split (10% of the dataset): 1,396 crops, 688 `bad` and
708 `good`.
| Class | Precision | Recall | F1 | Error rate | Crops |
|---|---|---|---|---|---|
| `bad` | 0.9744 | 0.9971 | 0.9856 | 0.29% | 688 |
| `good` | 0.9971 | 0.9746 | 0.9857 | 2.5% | 708 |
| **Overall** | **0.9857** | **0.9857** | **0.9857** | **1.4%** | **1,396** |
Confusion matrix (rows: true class, columns: predicted class):
| | predicted `bad` | predicted `good` |
|---|---|---|
| **`bad`** | 686 | 2 |
| **`good`** | 18 | 690 |
That is 20 errors out of 1,396 crops, an accuracy of 98.57%. The per-class error rate is
`1 βˆ’ recall`. The overall precision and recall are micro-averaged, so they both equal the accuracy,
and the overall F1 is the macro average of the two class F1 scores. The average test loss is 0.2581.
> These crops come from the same synthetic pipeline as the training crops (see
> [Training data](#training-data)). The numbers measure how well the model detects resampling and
> JPEG damage, not how it performs on real-world scans. Tune the threshold on your own data.
## Model architecture
A small convolutional network:
| Stage | Layers | Output channels |
|---|---|---|
| Block 1 | Conv 3Γ—3 (stride 1, padding 1, no bias) β†’ BatchNorm β†’ ReLU β†’ MaxPool 2Γ—2 | 32 |
| Block 2 | Conv 3Γ—3 β†’ BatchNorm β†’ ReLU β†’ MaxPool 2Γ—2 | 64 |
| Block 3 | Conv 3Γ—3 β†’ BatchNorm β†’ ReLU β†’ MaxPool 2Γ—2 | 128 |
| Block 4 | Conv 3Γ—3 β†’ BatchNorm β†’ ReLU β†’ MaxPool 2Γ—2 | 256 |
| Head | Global average pooling β†’ Flatten β†’ Dropout (p = 0.3) β†’ Linear (256 β†’ 2) β†’ Softmax | 2 |
The 390,370 parameters are 389,410 trainable weights plus the 960 BatchNorm running statistics.
The preprocessing contract is published in [`config.json`](config.json), and its basic fields are
also embedded in the ONNX file's `metadata_props`:
```python
import onnx
print({p.key: p.value for p in onnx.load(model_path).metadata_props})
```
## Training data
The dataset holds **13,947 crops**, split 90% / 10% into about 12,522 training and 1,396 test crops.
It is synthetic and was built from a single public source:
- [**olmOCR-mix-0225**](https://huggingface.co/datasets/allenai/olmOCR-mix-0225) (Allen Institute
for AI), a large and diverse collection of PDF pages.
Construction:
1. **Region cropping.** Every page was run through
[`Inceptive/PP-DocLayoutV3-ONNX`](https://huggingface.co/Inceptive/PP-DocLayoutV3-ONNX), and each
detected layout region was cropped. The classifier therefore sees the same kind of input it gets
at inference time in an OCR pipeline.
2. **Quality degradation.** Each crop was degraded by applying **successive upscales and
downscales** (resampling artifacts and loss of fine detail) and **JPEG compression** (blocking
and ringing artifacts).
3. **Labelling.** Original, untouched crops are labelled `good`. Degraded crops are labelled `bad`.
Unlike v0, no FUNSD data is used.
## Training procedure
| Setting | Value |
|---|---|
| Epochs | 10 |
| Batch size | 16 |
| Optimizer | AdamW, weight decay 1e-4 |
| Learning rate | 5e-4, constant |
| Loss | Cross-entropy with label smoothing 0.1 |
| Input | 640 Γ— 640 Γ— 3: RGB, aspect-preserving resize (upscaling allowed), centered white padding, ImageNet normalization |
The final training loss is 0.2581 (test: 0.2589). With label smoothing 0.1 and two classes, the
lowest reachable loss is about 0.199, so both values are close to that floor despite the ~98.5%
accuracy.
## Limitations and biases
- **Synthetic degradations only.** "Bad" means *resampled and/or JPEG-compressed*. The model was
not trained on motion blur, defocus, noise, low contrast, skew, show-through, stains or
handwriting quality, and it may misjudge real-world scans that fail in those ways.
- **Synthetic test set.** The evaluation above uses held-out crops from the same pipeline, so it
does not measure performance on real degraded scans.
- **Domain.** All source pages come from olmOCR's PDF mix. Other languages and scripts, scanned
forms, unusual layouts or very different rendering resolutions are less well covered.
- **Mostly blank or extreme crops.** A crop that is mostly empty margin, or very thin and elongated
(and so mostly white padding after letterboxing), carries little signal and may be scored `bad`
even when it is clean. Use it on tight layout-detector crops, not on arbitrary rectangles.
- **Small crops are upscaled.** Every crop is resized so that its longest side is 640 px, exactly
as during training, so do not skip this step for small regions. Very small crops carry little
detail, so treat their scores with more caution.
- **Probabilities are not calibrated.** Tune the `P(good)` threshold on your own data if the
cost of skipping a region differs from the cost of OCR'ing a bad one.
## License
This model is released under the **[Apache License 2.0](LICENSE)**, which allows commercial use.
Training data attribution: the training crops are derived from the page images of
[olmOCR-mix-0225](https://huggingface.co/datasets/allenai/olmOCR-mix-0225), which AI2 releases under
[ODC-BY 1.0](https://opendatacommons.org/licenses/by/1-0/). ODC-BY requires attribution (see
[Citation](#citation)). The dataset's GPT-4o-generated transcriptions were not used to train this model.
## Citation
If you use this model, please also cite the dataset it was trained on:
```bibtex
@misc{poznanski2025olmocr,
title = {olmOCR: Unlocking Trillions of Tokens in PDFs with Vision Language Models},
author = {Jake Poznanski and Jon Borchardt and Jason Dunkelberger and Regan Huff and Daniel Lin and Aman Rangapur and Christopher Wilhelm and Kyle Lo and Luca Soldaini},
year = {2025},
eprint = {2502.18443},
archivePrefix = {arXiv},
primaryClass = {cs.CL}
}
```