Update README.md
Browse files
README.md
CHANGED
|
@@ -1,3 +1,126 @@
|
|
| 1 |
---
|
| 2 |
license: mit
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
license: mit
|
| 3 |
+
library_name: onnx
|
| 4 |
+
tags:
|
| 5 |
+
- ocr
|
| 6 |
+
- text-recognition
|
| 7 |
+
- trocr
|
| 8 |
+
- documents
|
| 9 |
+
- transformer
|
| 10 |
+
- onnx
|
| 11 |
+
base_model: microsoft/trocr-base-printed
|
| 12 |
+
pipeline_tag: image-to-text
|
| 13 |
+
language:
|
| 14 |
+
- en
|
| 15 |
---
|
| 16 |
+
|
| 17 |
+
# TrOCR Base Printed (ONNX, fp32 + fp16 bundle)
|
| 18 |
+
|
| 19 |
+
ONNX exports of [microsoft/trocr-base-printed](https://huggingface.co/microsoft/trocr-base-printed) β Microsoft's Transformer-based OCR for printed text. ViT image encoder + GPT-style autoregressive text decoder, trained on SROIE printed-text crops. Recognizes the text in a cropped image of a single line/word.
|
| 20 |
+
|
| 21 |
+
This repo bundles **both fp32 and fp16 precisions** in one download β distribution symmetry, shared tokenizer + config files. Pick a precision via the `.onnx` filename in the `onnx/` subdir.
|
| 22 |
+
|
| 23 |
+
Re-exported from upstream PyTorch weights via a two-step pipeline. Provenance trail: Li et al. β microsoft/trocr-base-printed β `optimum-cli export onnx --task image-to-text-with-past` (fp32 stage) β `onnxconverter_common.float16.convert_float_to_float16(..., keep_io_types=True)` (fp16 cast on the fp32 graph) β these files.
|
| 24 |
+
|
| 25 |
+
Toolchain: `torch 2.4.x` (CUDA 12.4), `transformers 4.45.2`, `optimum[onnxruntime] 1.24.0`, `onnxconverter-common>=1.14`. Full conversion script: [`scripts/export-trocr-base-printed-fp16.ps1`](https://github.com/HeliosophLLC/DatumIngest/blob/main/scripts/export-trocr-base-printed-fp16.ps1) in the DatumIngest repo (despite the `-fp16` suffix on the script name, it produces both precisions in one run).
|
| 26 |
+
|
| 27 |
+
Why the two-step pipeline instead of `optimum-cli ... --dtype fp16` directly: the CUDA path requires a CUDA-enabled torch build in the venv (the other export scripts don't install one), and optimum-cli's fp16 merged-decoder export ships with an If-subgraph wiring bug on this architecture. The onnxconverter-common pass operates on the already-traced fp32 graph in place and sidesteps both issues. `keep_io_types=True` means inputs and outputs stay fp32 at the wire boundary β only internal weights + activations run in half precision β so runtime code feeds the same input tensors regardless of which `.onnx` file it loads.
|
| 28 |
+
|
| 29 |
+
Credit: Minghao Li, Tengchao Lv, Lei Cui, Yijuan Lu, Dinei Florencio, Cha Zhang, Zhoujun Li, Furu Wei (Microsoft Research). Paper: *"TrOCR: Transformer-based Optical Character Recognition with Pre-trained Models"*, 2021.
|
| 30 |
+
|
| 31 |
+
## What this repo contains
|
| 32 |
+
|
| 33 |
+
TrOCR is encoder-decoder, so the export splits into multiple files. **All shared (root) files must be present** along with the precision-specific `.onnx` files you choose to load.
|
| 34 |
+
|
| 35 |
+
### `onnx/` subdir β precision-specific files
|
| 36 |
+
|
| 37 |
+
| File | Variant | Size | Role |
|
| 38 |
+
|---|---|---|---|
|
| 39 |
+
| `encoder_model.onnx` | fp32 | ~700 MB | ViT image encoder |
|
| 40 |
+
| `encoder_model_fp16.onnx` | fp16 | ~350 MB | Half-precision ViT image encoder |
|
| 41 |
+
| `decoder_model_merged.onnx` | fp32 | ~700 MB | Text decoder with KV cache merged into one graph |
|
| 42 |
+
| `decoder_model_merged_fp16.onnx` | fp16 | ~350 MB | Half-precision text decoder |
|
| 43 |
+
|
| 44 |
+
The non-merged `decoder_model.onnx` is deliberately omitted β the merged form supersedes it for runtime use; keeping both would double the repo size for no benefit.
|
| 45 |
+
|
| 46 |
+
### Root β shared tokenizer + config files
|
| 47 |
+
|
| 48 |
+
| File | Role |
|
| 49 |
+
|---|---|
|
| 50 |
+
| `config.json` | Model architecture config |
|
| 51 |
+
| `generation_config.json` | Decoder generation defaults (max_length, EOS token, etc.) |
|
| 52 |
+
| `preprocessor_config.json` | Image preprocessing β `TrOCRProcessor` settings (resize, normalize) |
|
| 53 |
+
| `tokenizer.json` + `vocab.json` + `merges.txt` | BPE tokenizer files |
|
| 54 |
+
| `tokenizer_config.json` + `special_tokens_map.json` | Tokenizer metadata |
|
| 55 |
+
|
| 56 |
+
## Input / output (both variants)
|
| 57 |
+
|
| 58 |
+
| Stage | Input | Output |
|
| 59 |
+
|---|---|---|
|
| 60 |
+
| Encoder | `pixel_values` β NCHW **float32** (yes, even for the fp16 variant β IO types are kept fp32), preprocessed RGB image (typically 384Γ384) | `last_hidden_state` β encoder features (fp32 at the boundary; fp16 internally for the half-precision variant) |
|
| 61 |
+
| Decoder | `input_ids` (token sequence so far) + `encoder_hidden_states` + KV cache from prior step | next-token logits + updated KV cache |
|
| 62 |
+
|
| 63 |
+
The fp16 variant's `keep_io_types=True` setting means runtime code is **identical** between fp32 and fp16 β you don't have to cast inputs to `np.float16`. Only the on-disk weights and the internal compute differ.
|
| 64 |
+
|
| 65 |
+
## How to use
|
| 66 |
+
|
| 67 |
+
Greedy decoding orchestrated outside the ONNX graph β same encoder-decoder shape as Whisper, T5, BART, and friends:
|
| 68 |
+
|
| 69 |
+
```python
|
| 70 |
+
import onnxruntime as ort
|
| 71 |
+
import numpy as np
|
| 72 |
+
from PIL import Image
|
| 73 |
+
from transformers import TrOCRProcessor
|
| 74 |
+
|
| 75 |
+
# Pick a precision β same runtime code either way thanks to keep_io_types.
|
| 76 |
+
PRECISION_SUFFIX = "" # "" for fp32, "_fp16" for fp16
|
| 77 |
+
|
| 78 |
+
proc = TrOCRProcessor.from_pretrained(".")
|
| 79 |
+
encoder = ort.InferenceSession(f"onnx/encoder_model{PRECISION_SUFFIX}.onnx")
|
| 80 |
+
decoder = ort.InferenceSession(f"onnx/decoder_model_merged{PRECISION_SUFFIX}.onnx")
|
| 81 |
+
|
| 82 |
+
img = Image.open("cropped_text_line.jpg").convert("RGB")
|
| 83 |
+
pixel_values = proc(images=img, return_tensors="np").pixel_values # float32 in both cases
|
| 84 |
+
encoder_hidden = encoder.run(None, {"pixel_values": pixel_values})[0]
|
| 85 |
+
|
| 86 |
+
BOS = proc.tokenizer.cls_token_id
|
| 87 |
+
EOS = proc.tokenizer.eos_token_id
|
| 88 |
+
input_ids = np.array([[BOS]], dtype=np.int64)
|
| 89 |
+
generated, past_kv = [], None
|
| 90 |
+
|
| 91 |
+
for _ in range(64):
|
| 92 |
+
decoder_inputs = {"input_ids": input_ids, "encoder_hidden_states": encoder_hidden}
|
| 93 |
+
if past_kv is not None:
|
| 94 |
+
decoder_inputs.update(past_kv_to_inputs(past_kv))
|
| 95 |
+
outputs = decoder.run(None, decoder_inputs)
|
| 96 |
+
next_token = outputs[0][:, -1, :].argmax(-1)
|
| 97 |
+
if next_token.item() == EOS: break
|
| 98 |
+
generated.append(next_token.item())
|
| 99 |
+
input_ids = next_token.reshape(1, 1)
|
| 100 |
+
past_kv = outputs_to_past_kv(outputs[1:])
|
| 101 |
+
|
| 102 |
+
text = proc.tokenizer.decode(generated, skip_special_tokens=True)
|
| 103 |
+
```
|
| 104 |
+
|
| 105 |
+
The exact past-KV input/output names are Optimum-version-specific; inspect with Netron once after export to pin them down.
|
| 106 |
+
|
| 107 |
+
## Which precision should I use?
|
| 108 |
+
|
| 109 |
+
- **fp32** β full precision, identical numerics to upstream PyTorch reference. Default for accuracy-sensitive scientific work, OCR-accuracy benchmarks.
|
| 110 |
+
- **fp16** β ~half the disk footprint (~700 MB vs ~1.4 GB total) and half the model-load memory. On GPU / NPU with native fp16: modest speedup (typically 1.5-2Γ over fp32 on consumer GPUs). On CPU runtimes that upcast fp16 β fp32 internally, runtime speed is identical to fp32 but you save the memory.
|
| 111 |
+
|
| 112 |
+
The `keep_io_types=True` setting means switching between them is a single file-path change β no code changes needed.
|
| 113 |
+
|
| 114 |
+
## Related variants (not in this repo)
|
| 115 |
+
|
| 116 |
+
Microsoft publishes a small variant family of TrOCR β same architecture, different size / training corpus:
|
| 117 |
+
|
| 118 |
+
- `microsoft/trocr-small-printed` β smaller (~5Γ less), less accurate but faster.
|
| 119 |
+
- `microsoft/trocr-large-printed` β bigger, better quality.
|
| 120 |
+
- `microsoft/trocr-base-handwritten` β same size as this, trained on IAM handwritten dataset instead of SROIE.
|
| 121 |
+
|
| 122 |
+
All MIT-licensed, all re-exportable via the same script with a swapped `--model` arg.
|
| 123 |
+
|
| 124 |
+
## License
|
| 125 |
+
|
| 126 |
+
**MIT** β same as upstream [microsoft/trocr-base-printed](https://huggingface.co/microsoft/trocr-base-printed). `LICENSE` file included. Optimum ONNX export + fp16 conversion are numerical transformations only β no relicensing implication.
|