edgeparse_core/pdf/ocr/
tesseract_cli.rs

1//! Tesseract / RapidOCR CLI-backed [`super::OcrEngine`].
2
3use std::env;
4use std::fs;
5use std::path::PathBuf;
6use std::process::Command;
7use std::sync::OnceLock;
8use std::time::{SystemTime, UNIX_EPOCH};
9
10use image::GrayImage;
11use serde::Deserialize;
12
13use super::{OcrEngine, OcrWord};
14
15const TESSERACT_EFFECTIVE_DPI: u32 = 300;
16const MIN_OCR_WORD_CONFIDENCE: f64 = 6.0;
17const MAX_OCR_WORD_CONFIDENCE: f64 = 101.0;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20enum CliOcrKind {
21    Tesseract,
22    RapidOcr,
23}
24
25#[derive(Debug, Deserialize)]
26struct RapidOcrLine {
27    left: u32,
28    top: u32,
29    width: u32,
30    height: u32,
31    text: String,
32    confidence: f64,
33}
34
35static CLI_KIND: OnceLock<CliOcrKind> = OnceLock::new();
36static RAPIDOCR_PYTHON: OnceLock<Option<String>> = OnceLock::new();
37
38const RAPIDOCR_RUNNER: &str = r#"
39import json, sys
40from rapidocr import RapidOCR
41
42engine = RapidOCR()
43result = engine(sys.argv[1], use_det=True, use_cls=True, use_rec=True)
44
45if result is None:
46    print('[]')
47    raise SystemExit(0)
48
49boxes = getattr(result, 'boxes', []) or []
50txts = getattr(result, 'txts', []) or []
51scores = getattr(result, 'scores', []) or []
52out = []
53for box, text, score in zip(boxes, txts, scores):
54    if not text or not str(text).strip():
55        continue
56    xs = [pt[0] for pt in box]
57    ys = [pt[1] for pt in box]
58    out.append({
59        'left': int(min(xs)),
60        'top': int(min(ys)),
61        'width': max(1, int(max(xs) - min(xs))),
62        'height': max(1, int(max(ys) - min(ys))),
63        'text': str(text),
64        'confidence': float(score),
65    })
66print(json.dumps(out, ensure_ascii=False))
67"#;
68
69/// Shells out to Tesseract (or RapidOCR via Python when `EDGEPARSE_OCR_ENGINE=rapidocr`).
70#[derive(Debug, Clone)]
71pub struct TesseractCliEngine {
72    /// Tesseract page-segmentation mode (ignored for RapidOCR).
73    pub psm: String,
74    /// Tesseract OCR engine mode.
75    pub oem: String,
76}
77
78impl Default for TesseractCliEngine {
79    fn default() -> Self {
80        Self {
81            psm: "6".into(),
82            oem: "3".into(),
83        }
84    }
85}
86
87impl TesseractCliEngine {
88    /// Engine with an explicit PSM mode.
89    pub fn with_psm(psm: impl Into<String>) -> Self {
90        Self {
91            psm: psm.into(),
92            oem: "3".into(),
93        }
94    }
95}
96
97impl OcrEngine for TesseractCliEngine {
98    fn recognize(&self, gray: &GrayImage) -> Vec<OcrWord> {
99        self.recognize_psm(gray, &self.psm)
100    }
101
102    fn recognize_psm(&self, gray: &GrayImage, psm: &str) -> Vec<OcrWord> {
103        match selected_cli_kind() {
104            CliOcrKind::RapidOcr => run_rapidocr_words(gray).unwrap_or_default(),
105            CliOcrKind::Tesseract => run_tesseract_tsv(gray, psm, &self.oem).unwrap_or_default(),
106        }
107    }
108}
109
110fn selected_cli_kind() -> CliOcrKind {
111    *CLI_KIND.get_or_init(|| match env::var("EDGEPARSE_OCR_ENGINE") {
112        Ok(value) => match value.to_ascii_lowercase().as_str() {
113            "rapidocr" if rapidocr_python_command().is_some() => CliOcrKind::RapidOcr,
114            _ => CliOcrKind::Tesseract,
115        },
116        Err(_) => CliOcrKind::Tesseract,
117    })
118}
119
120fn rapidocr_python_command() -> Option<&'static str> {
121    RAPIDOCR_PYTHON
122        .get_or_init(|| {
123            let preferred = env::var("EDGEPARSE_OCR_PYTHON").ok();
124            let mut candidates = Vec::new();
125            if let Some(cmd) = preferred {
126                candidates.push(cmd);
127            }
128            candidates.push("python3".to_string());
129            candidates.push("python".to_string());
130
131            for candidate in candidates {
132                let ok = Command::new(&candidate)
133                    .arg("-c")
134                    .arg("import rapidocr")
135                    .output()
136                    .ok()
137                    .is_some_and(|out| out.status.success());
138                if ok {
139                    return Some(candidate);
140                }
141            }
142            None
143        })
144        .as_deref()
145}
146
147fn create_temp_dir() -> std::io::Result<PathBuf> {
148    let unique = SystemTime::now()
149        .duration_since(UNIX_EPOCH)
150        .unwrap_or_default()
151        .as_nanos();
152    let dir = std::env::temp_dir().join(format!(
153        "edgeparse-ocr-cli-{}-{}",
154        std::process::id(),
155        unique
156    ));
157    fs::create_dir_all(&dir)?;
158    Ok(dir)
159}
160
161fn run_tesseract_tsv(image: &GrayImage, psm: &str, oem: &str) -> Option<Vec<OcrWord>> {
162    let temp_dir = create_temp_dir().ok()?;
163    let image_path = temp_dir.join("ocr.png");
164    if image.save(&image_path).is_err() {
165        let _ = fs::remove_dir_all(&temp_dir);
166        return None;
167    }
168
169    let dpi = TESSERACT_EFFECTIVE_DPI.to_string();
170    let output = Command::new("tesseract")
171        .current_dir(&temp_dir)
172        .arg("ocr.png")
173        .arg("stdout")
174        .arg("--dpi")
175        .arg(&dpi)
176        .arg("--oem")
177        .arg(oem)
178        .arg("--psm")
179        .arg(psm)
180        .arg("-c")
181        .arg("load_system_dawg=0")
182        .arg("-c")
183        .arg("load_freq_dawg=0")
184        .arg("tsv")
185        .output()
186        .ok()?;
187    let _ = fs::remove_dir_all(&temp_dir);
188    if !output.status.success() {
189        return None;
190    }
191
192    let tsv = String::from_utf8_lossy(&output.stdout);
193    Some(parse_tesseract_tsv(&tsv))
194}
195
196pub(crate) fn parse_tesseract_tsv(tsv: &str) -> Vec<OcrWord> {
197    let mut words = Vec::new();
198    for line in tsv.lines().skip(1) {
199        let mut cols = line.splitn(12, '\t');
200        let level = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
201        if level != 5 {
202            continue;
203        }
204        let _page_num = cols.next();
205        let block_num = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
206        let par_num = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
207        let line_num = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
208        let _word_num = cols.next();
209        let left = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
210        let top = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
211        let width = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
212        let height = cols.next().and_then(|s| s.parse::<u32>().ok()).unwrap_or(0);
213        let confidence = cols
214            .next()
215            .and_then(|s| s.parse::<f64>().ok())
216            .unwrap_or(-1.0);
217        let text = cols.next().unwrap_or("").trim().to_string();
218        if !(MIN_OCR_WORD_CONFIDENCE..=MAX_OCR_WORD_CONFIDENCE).contains(&confidence)
219            || text.is_empty()
220            || width == 0
221            || height == 0
222            || !text.chars().any(|ch| ch.is_alphanumeric())
223        {
224            continue;
225        }
226        words.push(OcrWord {
227            line_key: (block_num, par_num, line_num),
228            left,
229            top,
230            width,
231            height,
232            text,
233            confidence,
234        });
235    }
236    words
237}
238
239fn rapidocr_lines_to_words(lines: Vec<RapidOcrLine>) -> Vec<OcrWord> {
240    let mut words = Vec::new();
241
242    for (line_idx, line) in lines.into_iter().enumerate() {
243        let tokens: Vec<&str> = line.text.split_whitespace().collect();
244        if tokens.is_empty() {
245            continue;
246        }
247
248        let total_chars: u32 = tokens
249            .iter()
250            .map(|token| token.chars().count() as u32)
251            .sum();
252        if total_chars == 0 {
253            continue;
254        }
255
256        let mut cursor = line.left;
257        let mut remaining_width = line.width.max(tokens.len() as u32);
258        let mut remaining_chars = total_chars;
259
260        for (token_idx, token) in tokens.iter().enumerate() {
261            let token_chars = token.chars().count() as u32;
262            let width = if token_idx == tokens.len() - 1 || remaining_chars <= token_chars {
263                remaining_width.max(1)
264            } else {
265                let proportional = ((remaining_width as f64) * (token_chars as f64)
266                    / (remaining_chars as f64))
267                    .round() as u32;
268                proportional.max(1).min(remaining_width)
269            };
270
271            words.push(OcrWord {
272                line_key: (0, line_idx as u32, 0),
273                left: cursor,
274                top: line.top,
275                width,
276                height: line.height.max(1),
277                text: (*token).to_string(),
278                confidence: line.confidence,
279            });
280
281            cursor = cursor.saturating_add(width);
282            remaining_width = remaining_width.saturating_sub(width);
283            remaining_chars = remaining_chars.saturating_sub(token_chars);
284        }
285    }
286
287    words
288}
289
290fn run_rapidocr_words(image: &GrayImage) -> Option<Vec<OcrWord>> {
291    let python = rapidocr_python_command()?;
292    let temp_dir = create_temp_dir().ok()?;
293    let image_path = temp_dir.join("ocr.png");
294    if image.save(&image_path).is_err() {
295        let _ = fs::remove_dir_all(&temp_dir);
296        return None;
297    }
298
299    let output = Command::new(python)
300        .current_dir(&temp_dir)
301        .arg("-c")
302        .arg(RAPIDOCR_RUNNER)
303        .arg("ocr.png")
304        .output()
305        .ok()?;
306    let _ = fs::remove_dir_all(&temp_dir);
307    if !output.status.success() {
308        return None;
309    }
310
311    let json = String::from_utf8_lossy(&output.stdout);
312    let lines: Vec<RapidOcrLine> = serde_json::from_str(&json).ok()?;
313    let words = rapidocr_lines_to_words(lines);
314    (!words.is_empty()).then_some(words)
315}
edgeparse_core/pdf/ocr/tesseract_cli.rs

edgeparse_core/pdf/ocr/
tesseract_cli.rs
