Adaptive Optical Compression for DeepSeek OCR

Adaptive OCR is designed to automatically select optimal compression level for DeepSeek-OCR based on document characteristics in order to reduce token usage without compromising reconstruction quality.

Repository Layout

├─ download_omnidocbench.py        # Dataset fetcher (images + annotations)
├─ evaluate/                       # Evaluation CLI entrypoints
│  ├─ evaluate.py                  # Main evaluation script
│  └─ evaluate_all_models.py       # Batch runner across compression presets
├─ pipeline/
│  ├─ document_processor.py        # Orchestrates OCR + metrics
│  ├─ ocr_inference.py             # DeepSeek-OCR HF wrapper
│  ├─ metrics.py                   # Character/word/edit-distance metrics
│  └─ omni_doc_bench_loader.py     # Dataset loader w/ embedded GT text
├─ label_optimal_models_threshold.py  # Data labelling: threshold-based optimal model selection
├─ train_simple_cnn.py             # Classifier training: CNN model training
├─ evaluate_simple_cnn.py          # Classifier evaluation: token savings calculation
├─ plot.py                         # Visualization: Pareto frontier plots
├─ calc_summary_stats.py           # Data-source counts (OmniDocBench)
├─ calc_reconstruction_acc.py      # Normalized edit distance vs compression/data_source table
├─ results/                        # Evaluation outputs (gitignored)
├─ simple_cnn_output/             # Classifier outputs (gitignored)
├─ OmniDocBench/                   # OmniDocBench Dataset (gitignored)
└─ requirements_vm.txt             # Python dependencies

Environment Setup

VM Setup

1. Python & CUDA

   • Python ≥3.10 (tested with 3.12.8).
   • CUDA-capable GPU recommended

2. Install dependencies

   conda create -n adaptive-ocr python=3.12 -y
   conda activate adaptive-ocr
   pip install -r requirements_vm.txt

3. Hugging Face authentication

   • Obtain a personal access token with dataset read permissions.
   • Store it in a .env file or export HF_TOKEN before running download/eval scripts:

     echo "HF_TOKEN=hf_xxx" > .env

Dataset Preparation

python download_omnidocbench.py

• Images land in OmniDocBench/images/, annotations in OmniDocBench/annotations/OmniDocBench.json.
• pipeline/omni_doc_bench_loader.py expects the above layout; adjust paths via CLI flags if you relocate the dataset.

Pipeline Overview

The adaptive OCR pipeline consists of two main stages:

1. Self-Supervised Data Labelling: Evaluate documents at multiple compression levels and label optimal models
2. Classifier Training & Evaluation: Train a CNN to predict optimal compression, then evaluate token savings

Stage 1: Self-Supervised Data Labelling Pipeline

Step 1: Evaluate Documents at Multiple Compression Levels

First, evaluate all documents across different compression presets to collect accuracy metrics.

Single Document Evaluation

python evaluate/evaluate.py \
  --doc_path /path/to/page.png \
  --gt_path /path/to/ground_truth.txt \
  --output_dir results/ \
  --image_size 640 --base_size 640 --prompt "<image>\nFree OCR."

OmniDocBench Subset Evaluation

python evaluate/evaluate.py \
  --dataset_path OmniDocBench \
  --output_dir results/tiny \
  --image_size 512 \
  --base_size 512 \
  --max_docs 100

Key flags:

• --prompt customizes the grounding text passed to DeepSeek-OCR.
• --model_name lets you point to alternative checkpoints.
• --max_docs (optional) randomly samples documents for quick experiments.
• --crop_mode (optional) uses gundam mode for dynamic tiling.
• --no_save_individual skips writing individual_results.json if storage is a concern.

Batch Evaluation Across All Compression Levels

Preset image sizes live in evaluate/evaluate_all_models.py.

python evaluate/evaluate_all_models.py \
  --dataset_path OmniDocBench \
  --models nano micro tiny small base large gundam \
  --max_docs 200

Each run writes to results/<model_key>/ with aggregated_metrics.json and individual_results.json.

Step 2: Label Optimal Models (Threshold Method)

After evaluating all compression levels, label each document with its optimal model using a threshold-based approach. This selects the smallest model that achieves within X% of the best accuracy.

python label_optimal_models_threshold.py \
  --results-root results \
  --output results/optimal_model_labels_threshold.json \
  --threshold 0.04 \
  --model-order nano micro tiny small base large gundam

Parameters:

• --threshold: Accuracy threshold as decimal (default: 0.04 = 4%). Picks smallest model within 4% of best accuracy.
• --model-order: Model names in size order (smallest to largest).
• --results-root: Root directory containing model result folders (e.g., results/tiny/, results/small/, etc.).

Output: optimal_model_labels_threshold.json containing document IDs mapped to their optimal model labels.

Stage 2: Classifier Training & Evaluation Pipeline

Step 3: Train CNN Classifier

Train a simple CNN from scratch to predict the optimal compression level for each document based on image features.

python train_simple_cnn.py \
  --labels_file results/optimal_model_labels_threshold.json \
  --images_dir OmniDocBench/images \
  --output_dir simple_cnn_output \
  --epochs 50 \
  --batch_size 32 \
  --lr 0.001 \
  --image_size 640

Key Parameters:

• --labels_file: Path to optimal model labels from Step 2.
• --epochs: Number of training epochs (default: 50).
• --batch_size: Batch size for training (default: 32).
• --lr: Learning rate (default: 0.001).
• --scheduler: Learning rate scheduler (plateau, cosine, onecycle).
• --label_smoothing: Label smoothing factor (0.0-0.1).
• --mixed_precision: Enable mixed precision training for faster training.

Output:

• simple_cnn_output/best_model.pth: Best model checkpoint
• simple_cnn_output/training_history.json: Training metrics
• simple_cnn_output/confusion_matrix.png: Classification confusion matrix

Step 4: Evaluate Classifier & Calculate Token Savings

Evaluate the trained classifier on the test set and calculate token usage savings compared to always using the largest (most expensive) model.

python evaluate_simple_cnn.py \
  --model_path simple_cnn_output/best_model.pth \
  --images_dir OmniDocBench/images \
  --results_dir results \
  --output_dir simple_cnn_output \
  --batch_size 32 \
  --image_size 640

What it calculates:

1. Token Savings: Total tokens saved by using predicted compression vs always using largest model
2. Accuracy Drop: Quality impact when classifier predicts smaller model than optimal
3. Per-Model Statistics: Breakdown of predictions and accuracy by model class

Output: simple_cnn_output/evaluation_results.json containing:

• Token savings statistics
• Accuracy metrics (normalized edit distance)
• Per-model prediction counts
• Comparison with single-resolution baselines

Step 5: Visualize Results

Generate Pareto frontier plots comparing adaptive-OCR against single-resolution approaches, showing the trade-off between tokens (cost) and normalized edit distance (quality).

python plot.py \
  --eval_results simple_cnn_output/evaluation_results.json \
  --output figs/pareto_plot.png \
  --include_gundam

Output: figs/pareto_plot.png showing:

• Adaptive-OCR performance (tokens vs accuracy)
• Single-resolution baselines (nano, micro, tiny, small, base, large, gundam)
• Pareto frontier highlighting optimal trade-offs

Outputs & Logs

Data Labelling Stage

• results/<model>/aggregated_metrics.json – macro averages (edit distance, accuracies, lengths, error counts).
• results/<model>/individual_results.json – per-document OCR text, metadata, and metrics.
• results/<model>/errors.json – only emitted when a document fails.
• results/optimal_model_labels_threshold.json – optimal model labels for each document (from Step 2).
• results/ocr_output/images/ – raw model dumps when DeepSeek-OCR saves intermediates.

Classifier Training Stage

• simple_cnn_output/best_model.pth – trained CNN model checkpoint.
• simple_cnn_output/training_history.json – training/validation metrics over epochs.
• simple_cnn_output/confusion_matrix.png – classification confusion matrix.
• simple_cnn_output/evaluation_results.json – token savings and accuracy metrics (from Step 4).
• figs/pareto_plot.png – visualization comparing adaptive vs single-resolution methods (from Step 5).

Analysis Utilities

These utilities provide additional insights into the dataset and evaluation results. They can be run at any point after the data labelling stage (Step 1-2).

OmniDocBench Data-Source Breakdown

python calc_summary_stats.py \
  --dataset-dir OmniDocBench

Prints the total document count plus counts/percentages for each page_info.page_attribute.data_source. Useful for understanding dataset composition.

Normalized Edit Distance vs Compression/Data Source

python calc_reconstruction_acc.py \
  --dataset-dir OmniDocBench \
  --results-root results

Collects every individual_results.json, joins with annotations by document_id, and emits a table with compression levels as rows and document types (data sources) as columns (values are mean normalized edit distance). Skipped entries (missing GT, doc IDs, etc.) are summarized below the table. Useful for analyzing which document types benefit most from different compression levels.

Troubleshooting

Data Labelling Stage

• Model load errors – ensure transformers==4.46 and that flash-attn dependencies are either installed or let pipeline/ocr_inference.py patch LlamaFlashAttention2.
• Missing HF token – both dataset download and model pulls may require authentication; set HF_TOKEN or run huggingface-cli login.
• OOM during inference – reduce image_size/base_size, enable --crop_mode, or switch to the tiny preset.
• Empty OCR outputs – check GPU logs and results/errors.json; DeepSeek-OCR sometimes returns empty strings if prompts or max lengths are invalid.
• Missing results for labelling – ensure all model evaluations completed successfully before running label_optimal_models_threshold.py.

Classifier Training Stage

• Class imbalance – if certain models are rarely selected, consider adjusting --threshold in Step 2 or using --label_smoothing in training.
• Low training accuracy – try increasing --epochs, adjusting --lr, or using a different --scheduler.
• OOM during training – reduce --batch_size or --image_size.
• Evaluation errors – ensure --results_dir points to the same results directory used in Step 1, and that all model result folders exist.

Contributing

1. Fork & branch.
2. Keep scripts runnable via CLI and documented here.
3. Run linting/tests relevant to your changes.
4. Open a PR detailing motivation, setup, and sample outputs.