Document Understanding & OCR

1. Modern OCR with TrOCR

Traditional OCR systems use convolutional neural networks for character recognition, often combined with recurrent layers (CRNN) for sequence modeling. TrOCR (Transformer-based OCR) replaces this entire pipeline with an encoder-decoder transformer. The encoder is a vision transformer (ViT or BEiT) pre-trained on images, and the decoder is a language model pre-trained on text. This architecture benefits from large-scale pre-training on both visual and textual data, achieving state-of-the-art results on handwritten and printed text recognition.

from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image

# Load TrOCR for printed text recognition
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-printed")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-printed")
model = model.to("cuda")

# OCR on a cropped text line image
image = Image.open("text_line.png").convert("RGB")
pixel_values = processor(images=image, return_tensors="pt").pixel_values.to("cuda")

generated_ids = model.generate(pixel_values, max_new_tokens=128)
text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(f"Recognized text: {text}")

2. The LayoutLM Family

The LayoutLM family of models (LayoutLM, LayoutLMv2, LayoutLMv3, LayoutXLM) pioneered the idea of jointly modeling text content, visual features, and 2D positional information in a single transformer. These models treat document understanding as a multimodal problem where the spatial arrangement of text is as informative as the text itself.

LayoutLMv3 Architecture

LayoutLMv3 unifies text, layout, and image pre-training with a single multimodal transformer. Text tokens receive both word embeddings and 2D position embeddings (bounding box coordinates on the page). Image patches are embedded alongside text tokens. The model is pre-trained with three objectives: masked language modeling, masked image modeling, and word-patch alignment. This design allows LayoutLMv3 to understand that text at the top-right of an invoice is likely a date, while numbers in a right-aligned column are likely prices.

from transformers import AutoProcessor, AutoModelForTokenClassification
from PIL import Image

# Load LayoutLMv3 fine-tuned for document entity extraction
processor = AutoProcessor.from_pretrained(
    "microsoft/layoutlmv3-base",
    apply_ocr=True,  # Built-in Tesseract OCR
)
model = AutoModelForTokenClassification.from_pretrained(
    "microsoft/layoutlmv3-base",
    num_labels=7,    # e.g., HEADER, QUESTION, ANSWER, etc.
)

# Process a document image
image = Image.open("invoice.png").convert("RGB")
encoding = processor(image, return_tensors="pt")

# Run inference
outputs = model(**encoding)
predictions = outputs.logits.argmax(-1).squeeze().tolist()

# Map predictions to words
words = encoding["input_ids"].squeeze().tolist()
tokens = processor.tokenizer.convert_ids_to_tokens(words)
for token, pred in zip(tokens, predictions):
    print(f"{token}: label_{pred}")

3. Document AI Pipelines

Production document understanding typically involves a multi-stage pipeline: document classification (what type of document is this?), OCR (extract text with bounding boxes), layout analysis (identify regions like headers, tables, paragraphs), entity extraction (find specific fields like dates, amounts, names), and validation (check extracted values for consistency). Each stage can use specialized models or a single end-to-end model.

Building a Document Processing Pipeline

import pytesseract
from PIL import Image
from transformers import pipeline

# Stage 1: OCR with Tesseract
image = Image.open("receipt.png")
ocr_data = pytesseract.image_to_data(
    image, output_type=pytesseract.Output.DICT
)

# Extract words and bounding boxes
words, boxes = [], []
for i in range(len(ocr_data["text"])):
    if ocr_data["conf"][i] > 50:  # Confidence threshold
        words.append(ocr_data["text"][i])
        boxes.append([
            ocr_data["left"][i],
            ocr_data["top"][i],
            ocr_data["left"][i] + ocr_data["width"][i],
            ocr_data["top"][i] + ocr_data["height"][i],
        ])

# Stage 2: Document question answering with LayoutLM
doc_qa = pipeline(
    "document-question-answering",
    model="impira/layoutlm-document-qa",
)

result = doc_qa(image, "What is the total amount?")
print(f"Total: {result[0]['answer']} (confidence: {result[0]['score']:.2f})")

VLM-Based Document Understanding

Vision-language models like GPT-4V, Gemini, and Qwen-VL offer a fundamentally different approach to document understanding. Instead of specialized OCR and layout models, you simply pass the document image to a VLM and ask questions in natural language. This approach requires no OCR preprocessing, handles diverse document types without task-specific fine-tuning, and can reason about complex layouts, charts, and tables. The tradeoff is higher latency, higher cost per document, and less predictable structured outputs compared to specialized pipelines.

from openai import OpenAI
import base64

client = OpenAI()

# Encode the document image
with open("invoice.png", "rb") as f:
    img_b64 = base64.b64encode(f.read()).decode()

# Extract structured data using a VLM
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{
        "role": "user",
        "content": [
            {"type": "text", "text": """Extract the following fields from this invoice as JSON:
- vendor_name, invoice_number, date, line_items (description, qty, price), subtotal, tax, total"""},
            {"type": "image_url", "image_url": {
                "url": f"data:image/png;base64,{img_b64}"
            }},
        ],
    }],
    response_format={"type": "json_object"},
)

print(response.choices[0].message.content)

4. Comparing Document Understanding Approaches