use std::io::Write;
use std::process::{Command, Stdio};
use image::{imageops, RgbImage};
use crate::layout::Region;
use crate::ocr_prep::is_text_label;
use crate::pdfium_backend::TextCell;
use docling_core::debug_log;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TesseractOptions {
pub cmd: String,
pub lang: Option<String>,
pub psm: Option<u8>,
pub tessdata_dir: Option<String>,
}
impl TesseractOptions {
pub fn from_env(lang: Option<String>) -> Self {
let psm = docling_core::env::nonempty("DOCLING_RS_TESSERACT_PSM").and_then(|raw| match raw
.trim()
.parse::<u8>()
{
Ok(n) if n <= 13 => Some(n),
_ => {
eprintln!(
"docling-pdf: DOCLING_RS_TESSERACT_PSM={raw:?} is not a page \
segmentation mode 0-13; using Tesseract's default"
);
None
}
});
Self {
cmd: docling_core::env::nonempty("DOCLING_TESSERACT")
.unwrap_or_else(|| "tesseract".to_string()),
lang,
psm,
tessdata_dir: docling_core::env::nonempty("DOCLING_RS_TESSDATA_DIR"),
}
}
}
pub fn lang_arg(raw: &str) -> Result<String, String> {
let mut stems = Vec::new();
for part in raw.split('+') {
let part = part.trim();
let part = part
.strip_prefix("iso:")
.or_else(|| part.strip_prefix("ISO:"))
.map_or(part, |tag| tag.trim());
if part.is_empty() {
return Err(format!("ocr_lang {raw:?} has an empty language entry"));
}
let mut subtags = part.split(['-', '_']);
let primary = subtags.next().unwrap_or_default();
let stem = if primary.len() == 2 && primary.bytes().all(|b| b.is_ascii_alphabetic()) {
let rest: Vec<String> = subtags.map(str::to_ascii_lowercase).collect();
bcp47_stem(&primary.to_ascii_lowercase(), &rest).ok_or_else(|| {
format!(
"ocr_lang {part:?}: no Tesseract traineddata name is known for that \
BCP-47 tag; give the tessdata stem directly (e.g. `deu`, `chi_sim`, \
`script/Latin` — see `tesseract --list-langs`)"
)
})?
} else {
match part.to_ascii_lowercase().as_str() {
"english" => "eng".to_string(),
"chinese" => "chi_sim".to_string(),
"chinese_cht" => "chi_tra".to_string(),
_ => {
let mut chars = part.chars();
let head_ok = chars
.next()
.is_some_and(|c| c.is_ascii_alphanumeric() || c == '_' || c == '/');
let tail_ok =
chars.all(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | '/' | '-'));
if !(head_ok && tail_ok) {
return Err(format!(
"ocr_lang {part:?} is not a Tesseract language identifier \
(letters, digits, `_`, `/`, `-`; e.g. `eng`, `deu+fra`, \
`script/Cyrillic`)"
));
}
part.to_string()
}
}
};
if !stems.contains(&stem) {
stems.push(stem);
}
}
Ok(stems.join("+"))
}
fn bcp47_stem(primary: &str, subtags: &[String]) -> Option<String> {
let has = |s: &str| subtags.iter().any(|t| t == s);
let stem = match primary {
"zh" | "ch" => {
if has("hant") || has("tw") || has("hk") || has("mo") || has("tra") {
"chi_tra"
} else {
"chi_sim"
}
}
"sr" => {
if has("latn") {
"srp_latn"
} else {
"srp"
}
}
"az" => {
if has("cyrl") {
"aze_cyrl"
} else {
"aze"
}
}
"uz" => {
if has("cyrl") {
"uzb_cyrl"
} else {
"uzb"
}
}
"de" => {
if has("latf") {
"deu_latf"
} else {
"deu"
}
}
"ku" => "kmr",
"no" | "nb" | "nn" => "nor",
"af" => "afr",
"am" => "amh",
"ar" => "ara",
"as" => "asm",
"be" => "bel",
"bn" => "ben",
"bo" => "bod",
"bs" => "bos",
"br" => "bre",
"bg" => "bul",
"ca" => "cat",
"cs" => "ces",
"co" => "cos",
"cy" => "cym",
"da" => "dan",
"dv" => "div",
"dz" => "dzo",
"el" => "ell",
"en" => "eng",
"eo" => "epo",
"et" => "est",
"eu" => "eus",
"fo" => "fao",
"fa" => "fas",
"fi" => "fin",
"fr" => "fra",
"fy" => "fry",
"gd" => "gla",
"ga" => "gle",
"gl" => "glg",
"gu" => "guj",
"ht" => "hat",
"he" => "heb",
"hi" => "hin",
"hr" => "hrv",
"hu" => "hun",
"hy" => "hye",
"iu" => "iku",
"id" => "ind",
"is" => "isl",
"it" => "ita",
"jv" => "jav",
"ja" => "jpn",
"kn" => "kan",
"ka" => "kat",
"kk" => "kaz",
"km" => "khm",
"ky" => "kir",
"ko" => "kor",
"lo" => "lao",
"la" => "lat",
"lv" => "lav",
"lt" => "lit",
"lb" => "ltz",
"ml" => "mal",
"mr" => "mar",
"mk" => "mkd",
"mt" => "mlt",
"mn" => "mon",
"mi" => "mri",
"ms" => "msa",
"my" => "mya",
"ne" => "nep",
"nl" => "nld",
"oc" => "oci",
"or" => "ori",
"pa" => "pan",
"pl" => "pol",
"pt" => "por",
"ps" => "pus",
"qu" => "que",
"ro" => "ron",
"ru" => "rus",
"sa" => "san",
"si" => "sin",
"sk" => "slk",
"sl" => "slv",
"sd" => "snd",
"es" => "spa",
"sq" => "sqi",
"su" => "sun",
"sw" => "swa",
"sv" => "swe",
"ta" => "tam",
"tt" => "tat",
"te" => "tel",
"tg" => "tgk",
"th" => "tha",
"ti" => "tir",
"to" => "ton",
"tr" => "tur",
"ug" => "uig",
"uk" => "ukr",
"ur" => "urd",
"vi" => "vie",
"yi" => "yid",
"yo" => "yor",
_ => return None,
};
Some(stem.to_string())
}
#[derive(Debug, Clone, PartialEq)]
struct Word {
block: u32,
par: u32,
line: u32,
l: f32,
t: f32,
r: f32,
b: f32,
conf: f32,
text: String,
}
fn parse_tsv(tsv: &str) -> Vec<Word> {
let mut lines = tsv.lines();
let Some(header) = lines.next() else {
return Vec::new();
};
let cols: Vec<&str> = header.split('\t').collect();
let col = |name: &str| cols.iter().position(|c| *c == name);
let (
Some(level),
Some(block),
Some(par),
Some(line),
Some(left),
Some(top),
Some(w),
Some(h),
Some(conf),
Some(text),
) = (
col("level"),
col("block_num"),
col("par_num"),
col("line_num"),
col("left"),
col("top"),
col("width"),
col("height"),
col("conf"),
col("text"),
)
else {
return Vec::new();
};
let mut words = Vec::new();
for row in lines {
let f: Vec<&str> = row.split('\t').collect();
if f.len() <= text || f[level] != "5" {
continue;
}
let text = f[text].trim();
if text.is_empty() {
continue;
}
let num = |i: usize| f[i].trim().parse::<f32>().unwrap_or(0.0);
let idx = |i: usize| f[i].trim().parse::<u32>().unwrap_or(0);
let (l, t) = (num(left), num(top));
words.push(Word {
block: idx(block),
par: idx(par),
line: idx(line),
l,
t,
r: l + num(w),
b: t + num(h),
conf: num(conf).clamp(0.0, 100.0),
text: text.to_string(),
});
}
words
}
type Unit = (f32, f32, f32, f32, String, f32);
fn group_lines(words: &[Word]) -> Vec<Unit> {
let mut out: Vec<Unit> = Vec::new();
let mut key: Option<(u32, u32, u32)> = None;
let mut n = 0.0f32;
for w in words {
let k = (w.block, w.par, w.line);
if key == Some(k) {
let last = out.last_mut().expect("a line is open");
last.0 = last.0.min(w.l);
last.1 = last.1.min(w.t);
last.2 = last.2.max(w.r);
last.3 = last.3.max(w.b);
last.4.push(' ');
last.4.push_str(&w.text);
last.5 = (last.5 * n + w.conf / 100.0) / (n + 1.0);
n += 1.0;
} else {
key = Some(k);
n = 1.0;
out.push((w.l, w.t, w.r, w.b, w.text.clone(), w.conf / 100.0));
}
}
out
}
fn word_units(words: &[Word]) -> Vec<Unit> {
words
.iter()
.map(|w| (w.l, w.t, w.r, w.b, w.text.clone(), w.conf / 100.0))
.collect()
}
const CROP_PAD: u32 = 6;
const CROP_PAD_RATIO: f32 = 0.3;
const CROP_PAD_MAX: u32 = 48;
fn crop_pad(h_px: f32) -> u32 {
(h_px * CROP_PAD_RATIO)
.round()
.max(CROP_PAD as f32)
.min(CROP_PAD_MAX as f32) as u32
}
struct Crop {
ox: u32,
oy: u32,
png: Vec<u8>,
bounds: (f32, f32, f32, f32),
}
fn crop_png(img: &RgbImage, region: &Region, scale: f32) -> Option<Crop> {
let (iw, ih) = img.dimensions();
let bounds = (
region.l * scale,
region.t * scale,
region.r * scale,
region.b * scale,
);
let pad = crop_pad(bounds.3 - bounds.1);
let l = (bounds.0.max(0.0) as u32).saturating_sub(pad);
let t = (bounds.1.max(0.0) as u32).saturating_sub(pad);
let r = ((bounds.2.max(0.0) as u32).saturating_add(pad)).min(iw);
let b = ((bounds.3.max(0.0) as u32).saturating_add(pad)).min(ih);
if r <= l || b <= t {
return None;
}
let crop = imageops::crop_imm(img, l, t, r - l, b - t).to_image();
Some(Crop {
ox: l,
oy: t,
png: encode_png(&crop)?,
bounds,
})
}
fn unit_in_region(crop: &Crop, unit: &Unit) -> bool {
let (l, t, r, b, _, _) = unit;
let cx = crop.ox as f32 + (l + r) / 2.0;
let cy = crop.oy as f32 + (t + b) / 2.0;
let slack = CROP_PAD as f32;
let (bl, bt, br, bb) = crop.bounds;
cx >= bl - slack && cx <= br + slack && cy >= bt - slack && cy <= bb + slack
}
fn encode_png(img: &RgbImage) -> Option<Vec<u8>> {
let mut buf = std::io::Cursor::new(Vec::new());
img.write_to(&mut buf, image::ImageFormat::Png).ok()?;
Some(buf.into_inner())
}
pub struct TesseractOcr {
opts: TesseractOptions,
lanes: usize,
has_osd: bool,
}
impl TesseractOcr {
pub fn load(opts: TesseractOptions, lanes: usize) -> Result<Self, String> {
let lanes = docling_core::env::parse::<usize>("DOCLING_RS_OCR_SESSIONS")
.filter(|&n| n > 0)
.unwrap_or(lanes)
.clamp(1, 16);
let version = Command::new(&opts.cmd)
.arg("--version")
.stdin(Stdio::null())
.output()
.map_err(|e| {
format!(
"tesseract binary {:?} not runnable ({e}); install tesseract-ocr or point \
DOCLING_TESSERACT at it",
opts.cmd
)
})?;
let banner = [version.stdout, version.stderr]
.iter()
.map(|b| String::from_utf8_lossy(b).trim().to_string())
.find(|s| !s.is_empty())
.unwrap_or_default();
let banner = banner.lines().next().unwrap_or_default().to_string();
let mut list = Command::new(&opts.cmd);
list.arg("--list-langs").stdin(Stdio::null());
if let Some(dir) = &opts.tessdata_dir {
list.arg("--tessdata-dir").arg(dir);
}
let list = list
.output()
.map_err(|e| format!("tesseract --list-langs failed to run: {e}"))?;
let installed: Vec<String> = String::from_utf8_lossy(&list.stdout)
.lines()
.skip(1)
.map(|l| l.trim().replace('\\', "/"))
.filter(|l| !l.is_empty())
.collect();
if installed.is_empty() {
return Err(format!(
"{banner}: no traineddata found (tesseract --list-langs printed nothing); \
install a language pack (e.g. tesseract-ocr-eng) or set \
DOCLING_RS_TESSDATA_DIR / TESSDATA_PREFIX"
));
}
let wanted: Vec<&str> = match &opts.lang {
Some(l) => l.split('+').collect(),
None => vec!["eng"],
};
let missing: Vec<&str> = wanted
.iter()
.copied()
.filter(|w| !installed.iter().any(|i| i == w))
.collect();
if !missing.is_empty() {
return Err(format!(
"{banner}: no traineddata for {} (installed: {}); install the language \
pack or pick an installed one with ocr_lang",
missing.join(", "),
installed.join(", ")
));
}
let has_osd = installed.iter().any(|i| i == "osd");
debug_log!(
"docling-pdf: OCR engine {banner} (lang {}, psm {}, {} lane(s), osd {})",
opts.lang.as_deref().unwrap_or("eng"),
opts.psm.map_or("default".to_string(), |p| p.to_string()),
lanes,
if has_osd { "yes" } else { "no" }
);
Ok(Self {
opts,
lanes,
has_osd,
})
}
fn command(&self) -> Command {
let mut cmd = Command::new(&self.opts.cmd);
cmd.env("OMP_THREAD_LIMIT", "1")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null());
if let Some(dir) = &self.opts.tessdata_dir {
cmd.arg("--tessdata-dir").arg(dir);
}
cmd
}
fn run(&self, mut cmd: Command, png: &[u8]) -> Result<String, String> {
let mut child = cmd
.spawn()
.map_err(|e| format!("tesseract: spawn {:?}: {e}", self.opts.cmd))?;
let mut stdin = child.stdin.take().expect("piped stdin");
let png = png.to_vec();
let writer = std::thread::spawn(move || stdin.write_all(&png));
let out = child
.wait_with_output()
.map_err(|e| format!("tesseract: wait: {e}"))?;
let _ = writer.join();
if !out.status.success() {
return Err(format!("tesseract exited with {}", out.status));
}
Ok(String::from_utf8_lossy(&out.stdout).into_owned())
}
fn recognize(&self, png: &[u8], dpi: u32) -> Result<Vec<Word>, String> {
let mut cmd = self.command();
if let Some(lang) = &self.opts.lang {
cmd.arg("-l").arg(lang);
}
if let Some(psm) = self.opts.psm {
cmd.arg("--psm").arg(psm.to_string());
}
cmd.arg("--dpi").arg(dpi.to_string());
cmd.args(["stdin", "stdout", "tsv"]);
Ok(parse_tsv(&self.run(cmd, png)?))
}
fn recognize_all(
&self,
jobs: &[Crop],
scale: f32,
units: fn(&[Word]) -> Vec<Unit>,
) -> Vec<(TextCell, f32)> {
let dpi = (scale * 72.0).round().max(1.0) as u32;
let lanes = self.lanes.min(jobs.len()).max(1);
let next = std::sync::atomic::AtomicUsize::new(0);
let results: Vec<std::sync::Mutex<Option<Vec<Word>>>> =
jobs.iter().map(|_| std::sync::Mutex::new(None)).collect();
std::thread::scope(|s| {
for _ in 0..lanes {
s.spawn(|| loop {
let i = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let Some(job) = jobs.get(i) else {
break;
};
let words = match self.recognize(&job.png, dpi) {
Ok(words) => words,
Err(e) => {
debug_log!("docling-pdf: tesseract: crop {i}: {e}; no text");
Vec::new()
}
};
*results[i].lock().expect("crop result slot") = Some(words);
});
}
});
let mut cells = Vec::new();
for (job, slot) in jobs.iter().zip(results) {
let words = slot
.into_inner()
.expect("crop result slot")
.unwrap_or_default();
let (ox, oy) = (job.ox as f32, job.oy as f32);
for unit in units(&words) {
if !unit_in_region(job, &unit) {
continue;
}
let (l, t, r, b, text, conf) = unit;
let text = text.trim().to_string();
if text.is_empty() {
continue;
}
cells.push((
TextCell {
text,
l: (ox + l) / scale,
t: (oy + t) / scale,
r: (ox + r) / scale,
b: (oy + b) / scale,
},
conf,
));
}
}
cells
}
fn crops(img: &RgbImage, regions: &[Region], scale: f32, keep: fn(&str) -> bool) -> Vec<Crop> {
regions
.iter()
.filter(|r| keep(r.label))
.filter_map(|r| crop_png(img, r, scale))
.collect()
}
pub fn ocr_page(
&mut self,
img: &RgbImage,
regions: &[Region],
scale: f32,
) -> Result<Vec<(TextCell, f32)>, String> {
let jobs = crate::timing::timed("ocr.prep", || {
Self::crops(img, regions, scale, is_text_label)
});
Ok(crate::timing::timed("ocr.rec", || {
self.recognize_all(&jobs, scale, group_lines)
}))
}
pub fn ocr_table_words(
&mut self,
img: &RgbImage,
regions: &[Region],
scale: f32,
) -> Result<Vec<(TextCell, f32)>, String> {
let jobs = Self::crops(img, regions, scale, crate::assemble::is_table_like);
Ok(self.recognize_all(&jobs, scale, word_units))
}
pub fn detect_orientation(&self, img: &RgbImage, scale: f32) -> Option<u16> {
if !self.has_osd {
debug_log!("docling-pdf: tesseract: no `osd` traineddata; orientation not probed");
return None;
}
let png = encode_png(img)?;
let mut cmd = self.command();
cmd.args(["--psm", "0", "-l", "osd", "--dpi"])
.arg(((scale * 72.0).round().max(1.0) as u32).to_string())
.args(["stdin", "stdout"]);
match self.run(cmd, &png) {
Ok(out) => parse_osd(&out),
Err(e) => {
debug_log!("docling-pdf: tesseract OSD failed ({e}); assuming upright");
None
}
}
}
}
fn parse_osd(out: &str) -> Option<u16> {
let deg = out
.lines()
.find_map(|l| l.trim().strip_prefix("Orientation in degrees:"))?
.trim()
.parse::<u16>()
.ok()?;
matches!(deg, 0 | 90 | 180 | 270).then_some(deg)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn lang_arg_maps_tags_and_keeps_stems() {
for (raw, want) in [
("eng", "eng"),
("en", "eng"),
("EN-us", "eng"),
("iso:en", "eng"),
("english", "eng"),
("ch", "chi_sim"),
("ch_tra", "chi_tra"),
("chinese_cht", "chi_tra"),
("zh", "chi_sim"),
("zh-Hant", "chi_tra"),
("zh-TW", "chi_tra"),
("iso:zh-CN", "chi_sim"),
("de", "deu"),
("de-Latf", "deu_latf"),
("sr-Latn", "srp_latn"),
("sr", "srp"),
("nb", "nor"),
("deu+fra", "deu+fra"),
("iso:de + fr", "deu+fra"),
("eng+eng", "eng"),
("script/Cyrillic", "script/Cyrillic"),
("chi_sim", "chi_sim"),
("srp_latn", "srp_latn"),
("custom_model", "custom_model"),
] {
assert_eq!(lang_arg(raw).as_deref(), Ok(want), "{raw:?}");
}
for raw in ["", "xx", "en+", "de;u", "a b", "iso:", "deu fra"] {
assert!(lang_arg(raw).is_err(), "{raw:?} should be rejected");
}
}
#[test]
fn crop_pad_grows_with_the_region_and_is_capped() {
assert_eq!(crop_pad(20.0), CROP_PAD);
assert_eq!(crop_pad(48.0), 14);
assert_eq!(crop_pad(1000.0), CROP_PAD_MAX);
assert_eq!(crop_pad(f32::NAN), CROP_PAD);
let region = Region {
label: "text",
score: 0.9,
l: 100.0,
t: 100.0,
r: 300.0,
b: 148.0,
};
let img = RgbImage::from_pixel(400, 400, image::Rgb([255, 255, 255]));
let crop = crop_png(&img, ®ion, 1.0).expect("crop");
assert_eq!((crop.ox, crop.oy), (86, 86), "14 px pad on a 48 px line");
assert_eq!(crop.bounds, (100.0, 100.0, 300.0, 148.0));
let crop = crop_png(&img, ®ion, 0.4).expect("crop");
assert_eq!((crop.ox, crop.oy), (34, 34), "6 px floor on a 19 px line");
}
#[test]
fn units_outside_the_region_are_dropped() {
let crop = Crop {
ox: 86,
oy: 86,
png: Vec::new(),
bounds: (100.0, 100.0, 300.0, 148.0),
};
let own: Unit = (14.0, 14.0, 214.0, 62.0, "Ação".into(), 0.9);
let jittered: Unit = (14.0, 6.0, 214.0, 58.0, "Ação".into(), 0.9);
let neighbour: Unit = (14.0, 62.0, 214.0, 110.0, "próxima".into(), 0.9);
assert!(unit_in_region(&crop, &own));
assert!(unit_in_region(&crop, &jittered));
assert!(!unit_in_region(&crop, &neighbour));
}
#[test]
fn tsv_words_group_into_lines() {
let tsv = "level\tpage_num\tblock_num\tpar_num\tline_num\tword_num\tleft\ttop\twidth\theight\tconf\ttext\n\
1\t1\t0\t0\t0\t0\t0\t0\t600\t160\t-1\t\n\
4\t1\t1\t1\t1\t0\t23\t28\t242\t25\t-1\t\n\
5\t1\t1\t1\t1\t1\t23\t28\t64\t20\t93.2\tHello\n\
5\t1\t1\t1\t1\t2\t94\t28\t93\t25\t92.4\tdocling\n\
5\t1\t1\t1\t1\t3\t195\t28\t70\t20\t96.0\tworld\n\
5\t1\t2\t1\t1\t1\t21\t98\t94\t20\t96.8\tSecond\n\
5\t1\t2\t1\t1\t2\t125\t98\t44\t20\t50\t \n\
5\t1\t2\t1\t1\t3\t179\t99\t44\t19\t96.8\t123\n";
let words = parse_tsv(tsv);
assert_eq!(words.len(), 5);
let lines = group_lines(&words);
assert_eq!(lines.len(), 2);
let (l, t, r, b, text, conf) = &lines[0];
assert_eq!(text, "Hello docling world");
assert_eq!((*l, *t, *r, *b), (23.0, 28.0, 265.0, 53.0));
assert!((conf - 0.9387).abs() < 1e-3, "{conf}");
assert_eq!(lines[1].4, "Second 123");
assert_eq!(word_units(&words).len(), 5);
assert!(parse_tsv("").is_empty());
assert!(parse_tsv("garbage\n5\t1\n").is_empty());
}
#[test]
fn osd_orientation_parses() {
let out = "Page number: 0\nOrientation in degrees: 270\nRotate: 90\n\
Orientation confidence: 6.47\nScript: Latin\nScript confidence: 4.05\n";
assert_eq!(parse_osd(out), Some(270));
assert_eq!(parse_osd("Orientation in degrees: 0\n"), Some(0));
assert_eq!(parse_osd("Orientation in degrees: 45\n"), None);
assert_eq!(parse_osd("Too few characters. Skipping this page\n"), None);
}
#[test]
fn psm_env_is_range_checked() {
let opts = TesseractOptions::from_env(Some("eng".into()));
assert_eq!(opts.lang.as_deref(), Some("eng"));
assert!(opts.psm.is_none() || opts.psm.is_some_and(|p| p <= 13));
}
}