steeldb/

ocr.rs

//! OCR fallback for scanned/image PDFs with no text layer — the bundled PP-OCRv6 pipeline, natively in
//! Rust. Two ONNX stages: DB **detection** (image → text-region probability map) and CRNN **recognition**
//! (each cropped region → text via CTC). Pages are rasterized with poppler `pdftoppm`. Gated behind the
//! `ocr` feature (implies `docs`/`onnx`).
//!
//! Detection post-processing is a pragmatic axis-aligned variant: binarize the prob map, take connected
//! components' bounding boxes (documents are near-horizontal), unclip slightly, crop, recognize.

use image::{imageops::FilterType, RgbImage};
use ort::session::Session;
use ort::value::Tensor;
use std::path::Path;

type Res<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;

const DET_MAX_SIDE: u32 = 960;
const DET_THRESH: f32 = 0.3;
const REC_H: u32 = 48;
const REC_MAX_W: u32 = 640;
// ImageNet normalization for detection; rec uses (x/255-0.5)/0.5.
const MEAN: [f32; 3] = [0.485, 0.456, 0.406];
const STD: [f32; 3] = [0.229, 0.224, 0.225];

pub struct OcrEngine {
    det: Session,
    rec: Session,
    dict: Vec<String>,
}

impl OcrEngine {
    /// `dir` holds `pp-ocrv6_small_det.onnx`, `pp-ocrv6_small_rec.onnx`, `ppocrv6_dict.txt`.
    pub fn load(dir: &Path) -> Res<OcrEngine> {
        let det = Session::builder()?.commit_from_file(dir.join("pp-ocrv6_small_det.onnx"))?;
        let rec = Session::builder()?.commit_from_file(dir.join("pp-ocrv6_small_rec.onnx"))?;
        let dict = std::fs::read_to_string(dir.join("ppocrv6_dict.txt"))?.lines().map(|l| l.to_string()).collect();
        Ok(OcrEngine { det, rec, dict })
    }

    /// Resolves via `crate::paths` (env `STEELDB_OCR_DIR` → cache → next-to-exe → `./models/ocr`); None
    /// if the models aren't present.
    pub fn default_engine() -> Option<OcrEngine> {
        let dir = crate::paths::model_dir("ocr", "STEELDB_OCR_DIR", "pp-ocrv6_small_det.onnx")?;
        OcrEngine::load(&dir).ok()
    }

    /// Load an image file and OCR it (convenience for callers without the `image` crate in scope).
    pub fn ocr_path(&mut self, path: &Path) -> Res<String> {
        let img = image::open(path)?.to_rgb8();
        self.ocr_image(&img)
    }

    /// OCR one page image → recognized text (reading order: top-to-bottom, left-to-right).
    pub fn ocr_image(&mut self, img: &RgbImage) -> Res<String> {
        let boxes = self.detect(img)?;
        let mut lines: Vec<(i32, i32, String)> = Vec::new();
        for (x0, y0, x1, y1) in boxes {
            let crop = crop_rgb(img, x0, y0, x1, y1);
            if crop.width() < 4 || crop.height() < 4 {
                continue;
            }
            let text = self.recognize(&crop)?;
            if !text.trim().is_empty() {
                lines.push((y0, x0, text));
            }
        }
        // group into reading order: sort by y, then x; newline when the vertical gap is large.
        lines.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)));
        let mut out = String::new();
        let mut last_y = i32::MIN;
        for (y, _x, text) in lines {
            if last_y != i32::MIN && (y - last_y).abs() > 12 {
                out.push('\n');
            } else if !out.is_empty() {
                out.push(' ');
            }
            out.push_str(text.trim());
            last_y = y;
        }
        Ok(out)
    }

    /// DB detection → axis-aligned region boxes in original-image coordinates.
    fn detect(&mut self, img: &RgbImage) -> Res<Vec<(i32, i32, i32, i32)>> {
        let (ow, oh) = (img.width(), img.height());
        let scale = (DET_MAX_SIDE as f32 / ow.max(oh) as f32).min(1.0);
        let rw = (((ow as f32 * scale) as u32).max(32) + 31) / 32 * 32;
        let rh = (((oh as f32 * scale) as u32).max(32) + 31) / 32 * 32;
        let resized = image::imageops::resize(img, rw, rh, FilterType::Triangle);

        let (rhu, rwu) = (rh as usize, rw as usize);
        let mut data = vec![0f32; 3 * rhu * rwu];
        for y in 0..rh {
            for x in 0..rw {
                let p = resized.get_pixel(x, y).0;
                for c in 0..3usize {
                    data[c * rhu * rwu + (y as usize) * rwu + x as usize] = (p[c] as f32 / 255.0 - MEAN[c]) / STD[c];
                }
            }
        }
        let t = Tensor::from_array(([1usize, 3, rhu, rwu], data))?;
        let out = self.det.run(ort::inputs!["x" => t])?;
        let (shape, prob) = out["fetch_name_0"].try_extract_tensor::<f32>()?;
        let (ph, pw) = (shape[shape.len() - 2] as u32, shape[shape.len() - 1] as u32);

        // binarize → connected components → bounding boxes
        let bin: Vec<bool> = prob.iter().map(|&v| v > DET_THRESH).collect();
        let boxes = components_bboxes(&bin, pw, ph);
        let (sx, sy) = (ow as f32 / pw as f32, oh as f32 / ph as f32);
        Ok(boxes
            .into_iter()
            .filter(|&(x0, y0, x1, y1)| x1 - x0 >= 2 && y1 - y0 >= 2)
            .map(|(x0, y0, x1, y1)| {
                // scale to original + unclip (small margin proportional to line height)
                let bh = (y1 - y0) as f32 * sy;
                let mx = (0.15 * bh + 1.0) as i32;
                (
                    ((x0 as f32 * sx) as i32 - mx).max(0),
                    ((y0 as f32 * sy) as i32 - mx).max(0),
                    (((x1 as f32 * sx) as i32 + mx).min(ow as i32)),
                    (((y1 as f32 * sy) as i32 + mx).min(oh as i32)),
                )
            })
            .filter(|&(x0, y0, x1, y1)| (x1 - x0) as f32 > 3.0 && (y1 - y0) as f32 > 3.0)
            .collect())
    }

    /// CRNN recognition of a single text-line crop → CTC-decoded string.
    fn recognize(&mut self, crop: &RgbImage) -> Res<String> {
        let ratio = crop.width() as f32 / crop.height() as f32;
        let w = ((REC_H as f32 * ratio).round() as u32).clamp(16, REC_MAX_W);
        let resized = image::imageops::resize(crop, w, REC_H, FilterType::Triangle);
        let (hu, wu) = (REC_H as usize, w as usize);
        let mut data = vec![0f32; 3 * hu * wu];
        for y in 0..REC_H {
            for x in 0..w {
                let p = resized.get_pixel(x, y).0;
                for c in 0..3usize {
                    data[c * hu * wu + (y as usize) * wu + x as usize] = (p[c] as f32 / 255.0 - 0.5) / 0.5;
                }
            }
        }
        let t = Tensor::from_array(([1usize, 3, hu, wu], data))?;
        // scope the session output so its &mut self.rec borrow releases before we read self.dict
        let (steps, classes, logits) = {
            let out = self.rec.run(ort::inputs!["x" => t])?;
            let (shape, logits) = out["fetch_name_0"].try_extract_tensor::<f32>()?; // [1, T, C]
            (shape[shape.len() - 2] as usize, shape[shape.len() - 1] as usize, logits.to_vec())
        };
        Ok(self.ctc_decode(&logits, steps, classes))
    }

    /// Greedy CTC: per-timestep argmax, collapse repeats, drop blanks; map index → dict char.
    /// PaddleOCR convention: index 0 = blank; 1..=dict.len() → dict[i-1]; dict.len()+1 → space.
    fn ctc_decode(&self, logits: &[f32], steps: usize, classes: usize) -> String {
        let mut out = String::new();
        let mut prev = usize::MAX;
        for t in 0..steps {
            let base = t * classes;
            let mut arg = 0usize;
            let mut best = f32::NEG_INFINITY;
            for c in 0..classes {
                let v = logits[base + c];
                if v > best {
                    best = v;
                    arg = c;
                }
            }
            if arg != prev && arg != 0 {
                if arg <= self.dict.len() {
                    out.push_str(&self.dict[arg - 1]);
                } else {
                    out.push(' ');
                }
            }
            prev = arg;
        }
        out
    }
}

fn crop_rgb(img: &RgbImage, x0: i32, y0: i32, x1: i32, y1: i32) -> RgbImage {
    let x = x0.max(0) as u32;
    let y = y0.max(0) as u32;
    let w = (x1 - x0).max(1) as u32;
    let h = (y1 - y0).max(1) as u32;
    image::imageops::crop_imm(img, x, y, w.min(img.width().saturating_sub(x)), h.min(img.height().saturating_sub(y))).to_image()
}

/// 4-connected components of a binary mask → each component's (x0,y0,x1,y1) bounding box (exclusive max).
fn components_bboxes(bin: &[bool], w: u32, h: u32) -> Vec<(i32, i32, i32, i32)> {
    let (w, h) = (w as usize, h as usize);
    let mut seen = vec![false; bin.len()];
    let mut out = Vec::new();
    let mut stack: Vec<(usize, usize)> = Vec::new();
    for sy in 0..h {
        for sx in 0..w {
            let idx = sy * w + sx;
            if !bin[idx] || seen[idx] {
                continue;
            }
            let (mut x0, mut y0, mut x1, mut y1) = (sx, sy, sx, sy);
            let mut count = 0usize;
            seen[idx] = true;
            stack.push((sx, sy));
            while let Some((x, y)) = stack.pop() {
                count += 1;
                x0 = x0.min(x);
                y0 = y0.min(y);
                x1 = x1.max(x);
                y1 = y1.max(y);
                let push = |nx: usize, ny: usize, stack: &mut Vec<(usize, usize)>, seen: &mut [bool]| {
                    let ni = ny * w + nx;
                    if bin[ni] && !seen[ni] {
                        seen[ni] = true;
                        stack.push((nx, ny));
                    }
                };
                if x > 0 {
                    push(x - 1, y, &mut stack, &mut seen);
                }
                if x + 1 < w {
                    push(x + 1, y, &mut stack, &mut seen);
                }
                if y > 0 {
                    push(x, y - 1, &mut stack, &mut seen);
                }
                if y + 1 < h {
                    push(x, y + 1, &mut stack, &mut seen);
                }
            }
            if count >= 6 {
                out.push((x0 as i32, y0 as i32, x1 as i32 + 1, y1 as i32 + 1));
            }
        }
    }
    out
}

/// OCR a PDF by rasterizing its pages (poppler `pdftoppm`, 150 dpi) and running the PP-OCRv6 pipeline.
/// Returns concatenated page text, or None if rasterization/models are unavailable.
pub fn ocr_pdf(path: &Path) -> Option<String> {
    let mut engine = OcrEngine::default_engine()?;
    let tmp = std::env::temp_dir().join(format!("steeldb_ocr_{}", std::process::id()));
    let _ = std::fs::create_dir_all(&tmp);
    let prefix = tmp.join("page");
    let status = std::process::Command::new("pdftoppm")
        .args(["-png", "-r", "150"])
        .arg(path)
        .arg(&prefix)
        .status()
        .ok()?;
    if !status.success() {
        return None;
    }
    let mut pages: Vec<std::path::PathBuf> = std::fs::read_dir(&tmp).ok()?.flatten().map(|e| e.path()).filter(|p| p.extension().map(|e| e == "png").unwrap_or(false)).collect();
    pages.sort();
    let mut out = String::new();
    for page in &pages {
        if let Ok(img) = image::open(page) {
            if let Ok(text) = engine.ocr_image(&img.to_rgb8()) {
                out.push_str(&text);
                out.push('\n');
            }
        }
    }
    let _ = std::fs::remove_dir_all(&tmp);
    let out = out.trim().to_string();
    (!out.is_empty()).then_some(out)
}