pagesift_ocr/
pdf.rs

1//! 把 [`Engine`] 接到 PDF 的扫描页接口上：第一张扫描图来时才准备模型（缺的先下载，由调用方提供
2//! 下载函数，本 crate 不碰网络）、才载入，没遇到扫描页的文档一点开销都没有。
3//!
4//! 准备模型不能堵住 CPU 池（它还负责正文提取和缓存编解码）：模型齐全时由第一个调用者载入，其他调用者等
5//! 几秒；缺模型时常驻服务在独立线程里下载，这期间识别立刻返回错误，文档照常返回并带警告。
6//! 失败不永久缓存，退避一段时间后再试。
7
8use std::panic::{AssertUnwindSafe, catch_unwind};
9use std::path::PathBuf;
10use std::sync::{Arc, Condvar, Mutex, MutexGuard, OnceLock};
11use std::time::{Duration, Instant};
12
13use pagesift_pdf::{OcrInfo, OcrLine, OcrOutput, PageOcr};
14
15use crate::device::{Device, ModelChoice, Quality, select_models};
16use crate::math::Math;
17use crate::models::{FORMULA, LAYOUT, Model, SMALL_REC};
18use crate::pipeline::Engine;
19use crate::{Image, Line};
20
21/// 把模型下载到 `<缓存根>/<仓库名>`（只在文件不全时调用），边下边校验。
22pub type Fetch = Arc<dyn Fn(&Model) -> Result<(), String> + Send + Sync>;
23
24/// 模型和设备怎么定。
25pub struct PdfOcrConfig {
26    pub root: PathBuf,
27    pub quality: Quality,
28    pub device: Device,
29    /// 调用方明确要了 GPU（`cuda`、`cuda:N`）。`auto` 时本机没有 N 卡不算退回，不报 `gpu_fallback`。
30    pub gpu_required: bool,
31    /// 也识别行间公式（多下载 860 MB 模型，每个公式要几秒）。
32    pub formulas: bool,
33    /// 没有时缺模型就报错（离线环境，提前用 `pagesift models fetch` 下好）。
34    pub fetch: Option<Fetch>,
35    /// 缺模型时放到独立线程下载，这期间识别立刻返回错误（常驻服务）；否则调用者等下载完（命令行、绑定）。
36    pub background_download: bool,
37}
38
39/// [策略] 下载或载入失败后隔这么久再试：网络抖动不该让 OCR 到重启前都不可用，也不该每张图都重试一次
40const RETRY_AFTER: Duration = Duration::from_secs(60);
41
42/// [`PdfOcr`] 准备模型的进度（模型载入成功后就不再看它）。
43enum Attempt {
44    /// 没有人在准备；上次失败退避完也回到这里。
45    Idle,
46    /// 有调用者在载入（或在前台下载），其他调用者等。
47    Loading,
48    /// 独立线程在下载，调用者不等。
49    Downloading,
50    Failed {
51        error: String,
52        at: Instant,
53    },
54}
55
56/// 准备进度和通知，也交给下载线程。
57struct Gate {
58    attempt: Mutex<Attempt>,
59    changed: Condvar,
60}
61
62impl Gate {
63    fn lock(&self) -> MutexGuard<'_, Attempt> {
64        self.attempt.lock().unwrap_or_else(|e| e.into_inner())
65    }
66
67    fn finish(&self, next: Attempt) {
68        *self.lock() = next;
69        self.changed.notify_all();
70    }
71}
72
73fn failed(error: String) -> Attempt {
74    Attempt::Failed {
75        error,
76        at: Instant::now(),
77    }
78}
79
80struct Loaded {
81    engine: Engine,
82    choice: ModelChoice,
83}
84
85/// 懒加载的 PDF OCR 引擎。
86pub struct PdfOcr {
87    config: PdfOcrConfig,
88    ready: OnceLock<Loaded>,
89    gate: Arc<Gate>,
90    retry_after: Duration,
91}
92
93impl PdfOcr {
94    pub fn new(config: PdfOcrConfig) -> Self {
95        Self {
96            config,
97            ready: OnceLock::new(),
98            gate: Arc::new(Gate {
99                attempt: Mutex::new(Attempt::Idle),
100                changed: Condvar::new(),
101            }),
102            retry_after: RETRY_AFTER,
103        }
104    }
105
106    /// 已载入的引擎（设备状态、GPU 统计从这里取）；还没遇到扫描页、没载入时为 None。
107    pub fn engine(&self) -> Option<&Engine> {
108        self.ready.get().map(|l| &l.engine)
109    }
110
111    /// 模型目录直接给（测评用）：立刻载入，不走缓存和下载。
112    pub fn from_dirs(
113        det_dir: PathBuf,
114        rec_dir: PathBuf,
115        math_dirs: Option<(PathBuf, PathBuf)>,
116        device: Device,
117    ) -> Result<Self, String> {
118        let this = Self::new(PdfOcrConfig {
119            root: PathBuf::new(),
120            quality: Quality::Fast,
121            device,
122            gpu_required: false,
123            formulas: false,
124            fetch: None,
125            background_download: false,
126        });
127        let mut engine = Engine::load_with(&det_dir, &rec_dir, device)?;
128        if let Some((layout, formula)) = math_dirs {
129            engine = engine.with_math(Math::load(&layout, &formula)?);
130        }
131        let choice = select_models(Quality::Fast, device);
132        let _ = this.ready.set(Loaded { engine, choice });
133        Ok(this)
134    }
135
136    /// 载入好的模型。别人在载入时等（最多等到 `deadline`）；还在下载、失败退避中时立刻返回原因。
137    fn loaded(&self, deadline: Option<Instant>) -> Result<&Loaded, String> {
138        if let Some(l) = self.ready.get() {
139            return Ok(l);
140        }
141        let mut attempt = self.gate.lock();
142        loop {
143            if let Some(l) = self.ready.get() {
144                return Ok(l);
145            }
146            match &*attempt {
147                Attempt::Idle => break,
148                Attempt::Downloading => {
149                    return Err("OCR models are still downloading".into());
150                }
151                Attempt::Failed { error, at } if at.elapsed() < self.retry_after => {
152                    return Err(error.clone());
153                }
154                Attempt::Failed { .. } => break,
155                Attempt::Loading => {}
156            }
157            attempt = match deadline {
158                None => self
159                    .gate
160                    .changed
161                    .wait(attempt)
162                    .unwrap_or_else(|e| e.into_inner()),
163                Some(d) => {
164                    let wait = d.saturating_duration_since(Instant::now());
165                    if wait.is_zero() {
166                        return Err("the OCR deadline passed while models were loading".into());
167                    }
168                    self.gate
169                        .changed
170                        .wait_timeout(attempt, wait)
171                        .unwrap_or_else(|e| e.into_inner())
172                        .0
173                }
174            };
175        }
176        *attempt = Attempt::Loading;
177        drop(attempt);
178        self.load_or_download()
179    }
180
181    /// 轮到这个调用者准备模型。模型文件上的问题（panic 也算）只记成失败：别的调用者还在等结果。
182    fn load_or_download(&self) -> Result<&Loaded, String> {
183        let c = &self.config;
184        let result = catch_unwind(AssertUnwindSafe(|| {
185            let choice = select_models(c.quality, c.device);
186            let math: &[Model] = if c.formulas { &[LAYOUT, FORMULA] } else { &[] };
187            let missing: Vec<Model> = [&choice.det, &choice.rec]
188                .into_iter()
189                .chain(math)
190                .filter(|m| !m.present(&c.root))
191                .copied()
192                .collect();
193            if !missing.is_empty() {
194                let Some(fetch) = &c.fetch else {
195                    return Err(format!(
196                        "model {} is missing from {}; run `pagesift models fetch`",
197                        missing[0].repo,
198                        c.root.display()
199                    ));
200                };
201                if c.background_download {
202                    return self
203                        .download_in_background(fetch.clone(), missing)
204                        .map(|()| None);
205                }
206                missing.iter().try_for_each(|m| fetch(m))?;
207            }
208            self.load(choice).map(Some)
209        }))
210        .unwrap_or_else(|_| Err("loading the OCR models panicked".into()));
211        match result {
212            Ok(Some(loaded)) => {
213                let _ = self.ready.set(loaded);
214                self.gate.finish(Attempt::Idle);
215                self.ready
216                    .get()
217                    .ok_or_else(|| "the OCR models are not loaded".to_owned())
218            }
219            // 下载线程已经接手，进度由它维护
220            Ok(None) => Err("OCR models are still downloading".into()),
221            Err(error) => {
222                self.gate.finish(failed(error.clone()));
223                Err(error)
224            }
225        }
226    }
227
228    /// 在独立线程里下载缺的模型。下载完回到空闲，下一个调用者来载入。
229    fn download_in_background(&self, fetch: Fetch, missing: Vec<Model>) -> Result<(), String> {
230        // 先改状态再起线程：线程可能很快结束
231        *self.gate.lock() = Attempt::Downloading;
232        let gate = self.gate.clone();
233        std::thread::Builder::new()
234            .name("pagesift-ocr-download".into())
235            .spawn(move || {
236                let done = catch_unwind(AssertUnwindSafe(|| {
237                    missing.iter().try_for_each(|m| fetch(m))
238                }))
239                .unwrap_or_else(|_| Err("downloading the OCR models panicked".into()));
240                gate.finish(match done {
241                    Ok(()) => Attempt::Idle,
242                    Err(error) => failed(error),
243                });
244            })
245            .map(|_| ())
246            .map_err(|e| format!("failed to start the model download thread: {e}"))
247    }
248
249    fn load(&self, choice: ModelChoice) -> Result<Loaded, String> {
250        let c = &self.config;
251        let mut engine =
252            Engine::load_with(&choice.det.dir(&c.root), &choice.rec.dir(&c.root), c.device)?;
253        if c.formulas {
254            engine = engine.with_math(Math::load(&LAYOUT.dir(&c.root), &FORMULA.dir(&c.root))?);
255        }
256        Ok(Loaded { engine, choice })
257    }
258}
259
260impl std::fmt::Debug for PdfOcr {
261    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
262        f.debug_struct("PdfOcr")
263            .field("quality", &self.config.quality)
264            .field("device", &self.config.device)
265            .finish_non_exhaustive()
266    }
267}
268
269/// 外框（各角 x、y 的最小、最大）。
270fn to_ocr_line(l: Line) -> OcrLine {
271    let xs = l.quad.map(|p| p.0);
272    let ys = l.quad.map(|p| p.1);
273    let min = |v: [f32; 4]| v.into_iter().fold(f32::MAX, f32::min);
274    let max = |v: [f32; 4]| v.into_iter().fold(f32::MIN, f32::max);
275    OcrLine {
276        x0: min(xs),
277        y0: min(ys),
278        x1: max(xs),
279        y1: max(ys),
280        formula: l.formula,
281        text: l.text,
282        score: l.score,
283    }
284}
285
286impl PageOcr for PdfOcr {
287    fn prepare(&self, deadline: Option<Instant>) -> Result<(), String> {
288        self.loaded(deadline).map(|_| ())
289    }
290
291    fn recognize(
292        &self,
293        width: u32,
294        height: u32,
295        gray: &[u8],
296        deadline: Option<Instant>,
297    ) -> Result<OcrOutput, String> {
298        let loaded = self.loaded(deadline)?;
299        let img = Image {
300            width: width as usize,
301            height: height as usize,
302            channels: 1,
303            data: gray.to_vec(),
304        };
305        let run = loaded.engine.run_until(&img, deadline)?;
306        // auto 模式在没有 N 卡的机器上用 CPU 是正常情况，不算退回
307        let expected = self.config.gpu_required || pagesift_gpu::device_count() > 0;
308        Ok(OcrOutput {
309            lines: run.lines.into_iter().map(to_ocr_line).collect(),
310            device: Some(device_name(run.device)),
311            fallback: run.fallback.filter(|_| expected),
312            partial_error: run.partial_error,
313        })
314    }
315
316    fn info(&self) -> OcrInfo {
317        let mut info = OcrInfo {
318            engine: ENGINE.into(),
319            model: SMALL_REC.repo.into(),
320            device: "cpu".into(),
321            ..Default::default()
322        };
323        match self.ready.get() {
324            None => {
325                if let Attempt::Failed { error, .. } = &*self.gate.lock() {
326                    info.error = Some(error.clone());
327                }
328            }
329            Some(l) => {
330                info.model = l.choice.rec.repo.into();
331                let status = l.engine.device_status();
332                info.device = device_name(status.actual);
333                // auto 模式在没有 N 卡的机器上用 CPU 是正常情况，不算退回
334                let expected = self.config.gpu_required || pagesift_gpu::device_count() > 0;
335                if status.fell_back() && expected {
336                    info.fallback = status.fallback;
337                }
338            }
339        }
340        info
341    }
342}
343
344/// [`OcrInfo::engine`]。
345const ENGINE: &str = "pp-ocrv6";
346
347/// `cpu` 或 `cuda:N`。
348fn device_name(d: Device) -> String {
349    match d {
350        Device::Cpu => "cpu".into(),
351        Device::Gpu { index } => format!("cuda:{index}"),
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use std::sync::atomic::{AtomicUsize, Ordering};
358    use std::sync::mpsc::{Sender, channel};
359
360    use super::*;
361
362    /// 一个不存在的模型目录：模型都算缺的。
363    fn config(fetch: Option<Fetch>, background_download: bool) -> PdfOcrConfig {
364        PdfOcrConfig {
365            root: PathBuf::from("/nonexistent/pagesift-ocr-test"),
366            quality: Quality::Fast,
367            device: Device::Cpu,
368            gpu_required: false,
369            formulas: false,
370            fetch,
371            background_download,
372        }
373    }
374
375    fn counting(calls: &Arc<AtomicUsize>) -> Fetch {
376        let calls = calls.clone();
377        Arc::new(move |_: &Model| {
378            calls.fetch_add(1, Ordering::SeqCst);
379            Err("network down".to_owned())
380        })
381    }
382
383    /// 下载失败不永久缓存：退避期内直接报原因，退避过后再试。
384    #[test]
385    fn a_failed_download_is_retried_after_the_backoff() {
386        let calls = Arc::new(AtomicUsize::new(0));
387        let mut ocr = PdfOcr::new(config(Some(counting(&calls)), false));
388        assert_eq!(ocr.prepare(None), Err("network down".to_owned()));
389        assert_eq!(calls.load(Ordering::SeqCst), 1);
390        // 退避期内不再下载，info 也如实报告
391        assert_eq!(ocr.prepare(None), Err("network down".to_owned()));
392        assert_eq!(calls.load(Ordering::SeqCst), 1);
393        assert_eq!(ocr.info().error.as_deref(), Some("network down"));
394        // 退避过后再试
395        ocr.retry_after = Duration::ZERO;
396        assert_eq!(ocr.prepare(None), Err("network down".to_owned()));
397        assert_eq!(calls.load(Ordering::SeqCst), 2);
398    }
399
400    /// 没有下载函数时报缺模型；之后补上文件也能恢复（不永久缓存），这里只看错误信息。
401    #[test]
402    fn missing_models_without_a_fetcher_say_what_to_do() {
403        let ocr = PdfOcr::new(config(None, true));
404        let e = ocr.prepare(None).unwrap_err();
405        assert!(e.contains("pagesift models fetch"), "{e}");
406    }
407
408    fn wait_for(ocr: &PdfOcr, f: impl Fn(&Attempt) -> bool) {
409        // wall-clock: 轮询的上界，慢机器只是多等
410        let end = Instant::now() + Duration::from_secs(30);
411        let mut a = ocr.gate.lock();
412        while !f(&a) {
413            assert!(Instant::now() < end, "等状态超时");
414            a = ocr
415                .gate
416                .changed
417                .wait_timeout(a, Duration::from_millis(50))
418                .unwrap()
419                .0;
420        }
421    }
422
423    /// 常驻服务：下载放到独立线程，调用者立刻拿到“还在下载”；下载失败后退避、再试。
424    #[test]
425    fn background_download_never_blocks_the_caller() {
426        let (go_tx, go_rx) = channel::<()>();
427        let (seen_tx, seen_rx) = channel::<()>();
428        let go = Mutex::new(go_rx);
429        let seen: Mutex<Sender<()>> = Mutex::new(seen_tx);
430        let fetch: Fetch = Arc::new(move |_: &Model| {
431            seen.lock().unwrap().send(()).unwrap();
432            // 卡住，直到测试放行
433            go.lock().unwrap().recv().unwrap();
434            Err("network down".to_owned())
435        });
436        let mut ocr = PdfOcr::new(config(Some(fetch), true));
437        let e = ocr.prepare(None).unwrap_err();
438        assert!(e.contains("still downloading"), "{e}");
439        seen_rx.recv_timeout(Duration::from_secs(30)).unwrap();
440        // 下载线程还卡着，调用者（包括 recognize）照样立刻返回
441        // wall-clock: 下载线程在等 go_tx，等它的 bug 会一直卡住；10 s 远大于正常的几微秒
442        let t = Instant::now();
443        assert!(ocr.prepare(None).unwrap_err().contains("still downloading"));
444        let e = ocr.recognize(1, 1, &[255], None).unwrap_err();
445        assert!(e.contains("still downloading"), "{e}");
446        assert!(t.elapsed() < Duration::from_secs(10));
447        go_tx.send(()).unwrap();
448        wait_for(&ocr, |a| matches!(a, Attempt::Failed { .. }));
449        assert_eq!(ocr.prepare(None), Err("network down".to_owned()));
450        // 退避过后重新起一个下载线程
451        ocr.retry_after = Duration::ZERO;
452        assert!(ocr.prepare(None).unwrap_err().contains("still downloading"));
453        seen_rx.recv_timeout(Duration::from_secs(30)).unwrap();
454        go_tx.send(()).unwrap();
455        wait_for(&ocr, |a| matches!(a, Attempt::Failed { .. }));
456    }
457
458    /// 有人在载入时其他调用者等，等到截止时间就报错，不永远占着工人。
459    #[test]
460    fn waiters_give_up_at_their_deadline() {
461        let ocr = PdfOcr::new(config(None, false));
462        *ocr.gate.lock() = Attempt::Loading;
463        let e = ocr
464            // wall-clock: 截止时间要等它过去才报错，机器慢只会过得更早
465            .prepare(Some(Instant::now() + Duration::from_millis(20)))
466            .unwrap_err();
467        assert!(e.contains("deadline"), "{e}");
468    }
469}
pagesift_ocr/pdf.rs

pagesift_ocr/
pdf.rs
