OCR (Optical Character Recognition)

OCR (Optical Character Recognition)

Overview

This example demonstrates TrOCR-based optical character recognition using Candle and Hugging Face transformers.

Source Code

• examples/OCR/src/main.rs
• examples/OCR/src/image_processor.rs

Status

Work in progress

src/main.rs currently defines only the model-download helpers — it has no fn main and never declares mod image_processor, so the crate does not build as a binary yet and cargo run fails. This page documents the parts that exist and the intended inference flow.

Project Structure

examples/OCR/
├── Cargo.toml                 # Crate manifest and dependencies
├── src/
│   ├── main.rs                # Model download helpers (no main() yet)
│   └── image_processor.rs     # Image load + normalize to a TrOCR input tensor
└── models/
    └── recognizer/            # Populated by ensure_model_files() (contains a stray 'ddd' file)

Path	Role
Cargo.toml	Defines the ocr crate and its dependencies
src/main.rs	ensure_model_files / download_model_file only
src/image_processor.rs	ImageProcessor::load and preprocess (not yet referenced from main)
models/recognizer/	Destination for config.json, model.safetensors, tokenizer.json

Dependencies

Crate	Version	Purpose
anyhow	1	Error handling
candle-core	0.9	Tensors, Device, DType
candle-nn	0.9	Neural-network modules
candle-transformers	0.9	Transformer model implementations
image	0.25	Image loading and resizing
ureq	2	Blocking HTTP download of model files (tls feature)
serde / serde_json	1	Config/tokenizer parsing
kornia	0.1	Declared but not used (preprocessing is written by hand)
tokenizers	0.21	Declared but not used yet (needed for decoding)

Model

TrOCR (Transformer-based Optical Character Recognition) is a state-of-the-art OCR model that uses a transformer architecture (Vision Transformer encoder + Transformer decoder) for text recognition.

Implementation Details

• Library: Candle (candle-core, candle-nn, candle-transformers)
• Model: TrOCR (pre-trained on printed/handwritten text)
• Task: Text recognition from images
• Model Source: Hugging Face Hub (microsoft/trocr-base-printed)
• Tokenizer: WordPiece tokenizer (loaded from tokenizer.json)

Running the Example

cd examples/OCR
cargo run

Code Walkthrough (Line by Line)

main.rs - Model Download (Lines 1-85)

use anyhow::{Context, Result};
use std::{
    fs::{self, File},
    io::{self, BufReader, Read, Write},
    path::Path,
};

const MODEL_DIR: &str = "models/recognizer";

const MODEL_FILES: &[&str] = &["config.json", "model.safetensors", "tokenizer.json"];

const MODEL_BASE_URL: &str = "https://huggingface.co/microsoft/trocr-base-printed/resolve/main/";

fn ensure_model_files() -> Result<()> {
    fs::create_dir_all(MODEL_DIR).context("Failed to create model directory")?;

    for filename in MODEL_FILES {
        let destination = Path::new(MODEL_DIR).join(filename);

        if destination.is_file() {
            let metadata = fs::metadata(&destination)?;

            if metadata.len() > 0 {
                println!("Model file exists: {}", destination.display());
                continue;
            }
        }

        download_model_file(filename, &destination)?;
    }

    Ok(())
}

fn download_model_file(filename: &str, destination: &Path) -> Result<()> {
    let url = format!("{MODEL_BASE_URL}{filename}");

    println!("Downloading {filename}...");
    println!("URL: {url}");

    let response = ureq::get(&url).call().with_context(|| {
        format!(
            "Failed to download {filename}. \
                 Check your internet connection and the Hugging Face URL."
        )
    })?;

    let temporary_path = destination.with_extension("download");

    let result = (|| -> Result<()> {
        let mut output =
            File::create(&temporary_path).context("Failed to create temporary model file")?;

        let mut reader = response.into_reader();

        io::copy(&mut reader, &mut output).context("Failed to write downloaded model file")?;

        output.flush()?;

        let size = output.metadata()?.len();

        if size == 0 {
            anyhow::bail!("Downloaded file is empty: {filename}");
        }

        drop(output);

        fs::rename(&temporary_path, destination)
            .context("Failed to move downloaded file into place")?;

        println!(
            "Downloaded {} ({:.2} MB)",
            destination.display(),
            size as f64 / (1024.0 * 1024.0)
        );

        Ok(())
    })();

    if result.is_err() {
        let _ = fs::remove_file(&temporary_path);
    }

    result
}

Model Download Details:

• Lines 8-10: Model directory and required files (config, weights, tokenizer)
• Line 12: Hugging Face Hub URL for microsoft/trocr-base-printed
• Lines 14-33: Check existing files, skip if present and non-empty
• Lines 35-84: Download with ureq, write to temp file, atomic rename
• Lines 50-56: Stream download directly to file (memory efficient)
• Lines 60-64: Validate non-empty download
• Lines 68-69: Atomic rename from .download to final name

image_processor.rs - Image Preprocessing (Lines 1-54)

use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use image::{DynamicImage, ImageReader};

pub struct ImageProcessor {
    width: u32,
    height: u32,
}

impl ImageProcessor {
    pub fn new() -> Self {
        Self {
            width: 384,
            height: 384,
        }
    }

    pub fn load(&self, path: &str) -> Result<DynamicImage> {
        let image = ImageReader::open(path)
            .with_context(|| format!("Cannot open image: {path}"))?
            .decode()
            .context("Cannot decode image")?;

        Ok(image)
    }

    pub fn preprocess(&self, image: &DynamicImage, device: &Device) -> Result<Tensor> {
        let resized = image.resize_exact(
            self.width,
            self.height,
            image::imageops::FilterType::Triangle,
        );

        let rgb = resized.to_rgb8();
        let raw = rgb.into_raw();

        let tensor = Tensor::from_vec(raw, (self.height as usize, self.width as usize, 3), device)?
            .permute((2, 0, 1))?
            .to_dtype(DType::F32)?;

        // Convert pixels from [0, 255] to [0, 1].
        let tensor = (tensor / 255.0)?;

        // Normalize using mean=0.5 and std=0.5 per RGB channel.
        let mean = Tensor::from_vec(vec![0.5_f32, 0.5, 0.5], (3, 1, 1), device)?;

        let std = Tensor::from_vec(vec![0.5_f32, 0.5, 0.5], (3, 1, 1), device)?;

        let normalized = tensor.broadcast_sub(&mean)?.broadcast_div(&std)?;

        // TrOCR expects a batch dimension: [1, 3, 384, 384].
        Ok(normalized.unsqueeze(0)?)
    }
}

Preprocessing Pipeline (TrOCR Specific):

• Lines 11-16: Fixed input size 384x384 (TrOCR base model requirement)
• Lines 18-25: Load and decode image
• Lines 27-38: Resize to exact 384x384
• Lines 40-41: Convert to RGB, get raw bytes (HWC layout)
• Line 42: Create tensor [H, W, 3] → permute to [3, H, W] (CHW)
• Line 43: Convert to F32
• Line 46: Scale [0, 255] → [0, 1]
• Lines 49-50: Normalize with mean=0.5, std=0.5 (TrOCR uses this normalization)
• Line 52: Add batch dimension → [1, 3, 384, 384]

Missing Inference Code

Note: The current implementation only includes model downloading and image preprocessing. The actual TrOCR inference (Vision Transformer encoder + Transformer decoder with autoregressive token generation) would be implemented in additional modules. The typical TrOCR inference flow would be:

// Pseudocode for TrOCR inference
pub struct TrOCR {
    encoder: VisionTransformer,
    decoder: TransformerDecoder,
    tokenizer: Tokenizer,
}

impl TrOCR {
    pub fn recognize(&self, image_tensor: &Tensor) -> Result<String> {
        // 1. Encode image: [1, 3, 384, 384] -> [1, seq_len, hidden_dim]
        let encoder_outputs = self.encoder.forward(image_tensor)?;

        // 2. Autoregressive decoding
        let mut decoder_input_ids = vec![self.tokenizer.get_vocab()["[CLS]"]];
        let mut generated = Vec::new();

        for _ in 0..max_length {
            let logits = self.decoder.forward(
                &Tensor::new(&[decoder_input_ids], device)?,
                &encoder_outputs
            )?;

            let next_token = logits.argmax(-1)?;
            if next_token == self.tokenizer.get_vocab()["[SEP]"] {
                break;
            }
            generated.push(next_token);
            decoder_input_ids.push(next_token);
        }

        // 3. Decode tokens to text
        self.tokenizer.decode(&generated, true)
    }
}

Model Files Required

Downloaded automatically to models/recognizer/ by ensure_model_files():

• config.json - TrOCR model configuration
• model.safetensors - Model weights (~150MB for the base model)
• tokenizer.json - WordPiece tokenizer (needed once decoding is implemented)

Expected Output

When fully implemented, the program will output the recognized text from the input image(s):

Recognized text: "Hello World"

Architecture Notes

TrOCR consists of:

1. Encoder: Vision Transformer (ViT) - processes image patches into embeddings
2. Decoder: Transformer decoder - generates text tokens autoregressively
3. Tokenizer: WordPiece vocabulary (~50k tokens)

The model is pre-trained on large-scale synthetic and real printed/handwritten text datasets.