PierreLF38's picture
Update README.md
1da3138 verified
|
Raw History Blame Contribute Delete
13.8 kB
metadata
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, 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, 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. FUNSD is available for non-commercial, research and educational purposes only, and it was what kept v0 non-commercial.

The previous release, 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).

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.9951 on 1,016 held-out crops (see 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 (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.8 β†’ 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:

pip install onnxruntime numpy pillow huggingface_hub

Classify one crop:

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

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 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,016 crops, 516 bad and 500 good.

Class Precision Recall F1 Error rate Crops
bad 0.9942 0.9961 0.9952 0.39% 516
good 0.9960 0.9940 0.9950 0.60% 500
Overall 0.9951 0.9951 0.9951 0.49% 1,016

Confusion matrix (rows: true class, columns: predicted class):

predicted bad predicted good
bad 514 2
good 3 497

That is 5 errors out of 1,016 crops, an accuracy of 99.51%. 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.2346.

These crops come from the same synthetic pipeline as the training crops (see 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, and its basic fields are also embedded in the ONNX file's metadata_props:

import onnx
print({p.key: p.value for p in onnx.load(model_path).metadata_props})

Training data

The dataset holds 10,142 crops, split 90% / 10% into about 9,126 training and 1,016 test crops. It is synthetic and was built from a single public source:

  • 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, 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 32
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.2409 (test: 0.2346). 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 ~99.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, which allows commercial use.

Training data attribution: the training crops are derived from the page images of olmOCR-mix-0225, which AI2 releases under ODC-BY 1.0. ODC-BY requires attribution (see 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:

@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}
}