Struct Engine

Source

pub struct Engine { /* private fields */ }

Expand description

检测加识别。

Implementations§

Source §

impl Engine

Source

pub fn load(det_dir: &Path, rec_dir: &Path) -> Result<Engine, String>

在 CPU 上跑，同 load_with(det_dir, rec_dir, Device::Cpu)。

Source

pub fn load_with( det_dir: &Path, rec_dir: &Path, device: Device, ) -> Result<Engine, String>

载入模型并打开 device。GPU 打不开（没有驱动、PTX 编译失败、PAGESIFT_GPU=off 等）不算错误：退回 CPU，原因见 Engine::device_status。只有模型文件的问题才返回 Err。

Source

pub fn with_math(self, math: Math) -> Engine

同时识别行间公式，输出 $$…$$。

Source

pub fn device_status(&self) -> DeviceStatus

实际用的设备、是否退回了 CPU、退回原因。

Examples found in repository ?

examples/pdf_ocr_dir.rs (line 93)

27fn main() {
28    let a: Vec<PathBuf> = std::env::args().skip(1).map(PathBuf::from).collect();
29    let (src, dst) = (&a[0], &a[1]);
30    let math = (a.len() >= 6).then(|| (a[4].clone(), a[5].clone()));
31    let device = match std::env::var("OCR_DEVICE").as_deref() {
32        Ok("gpu") => Device::Gpu { index: 0 },
33        _ => Device::Cpu,
34    };
35    let jobs: usize = std::env::var("OCR_JOBS")
36        .ok()
37        .and_then(|v| v.parse().ok())
38        .unwrap_or(1);
39    let engine =
40        Arc::new(PdfOcr::from_dirs(a[2].clone(), a[3].clone(), math, device).expect("载入模型"));
41    // 按字节数估的时限没算 OCR 的耗时：关掉，只靠逐页的识别
42    let opts = pagesift_pdf::Options {
43        ocr: Some(engine.clone()),
44        adaptive_timeout: false,
45        // 所有带图的页都识别（Firecrawl 的 mode: "ocr"）：MRC 页有文字层，不判为扫描页
46        force_ocr: std::env::var_os("PDF_OCR_FORCE").is_some(),
47        ..Default::default()
48    };
49    let mut files = Vec::new();
50    walk(src, &mut files);
51    files.sort();
52    let next = AtomicUsize::new(0);
53    let totals = Mutex::new((0.0f64, 0usize));
54    let wall = Instant::now();
55    std::thread::scope(|s| {
56        for _ in 0..jobs {
57            s.spawn(|| {
58                while let Some(f) = files.get(next.fetch_add(1, Ordering::Relaxed)) {
59                    let bytes = std::fs::read(f).unwrap();
60                    let t = Instant::now();
61                    let (md, bad) = match pagesift_pdf::extract_with(&bytes, &opts) {
62                        Ok(doc) => (doc.markdown, 0),
63                        Err(e) => {
64                            eprintln!("{}: {e}", f.display());
65                            (String::new(), 1)
66                        }
67                    };
68                    let secs = t.elapsed().as_secs_f64();
69                    {
70                        let mut g = totals.lock().unwrap();
71                        g.0 += secs;
72                        g.1 += bad;
73                    }
74                    eprintln!("{}\t{:.0} ms", f.display(), secs * 1e3);
75                    let rel = f.strip_prefix(src).unwrap().with_extension("");
76                    let out = dst.join(format!("{}_pg1_repeat1.md", rel.display()));
77                    std::fs::create_dir_all(out.parent().unwrap()).unwrap();
78                    std::fs::write(out, md).unwrap();
79                }
80            });
81        }
82    });
83    let (total, failed) = *totals.lock().unwrap();
84    let wall = wall.elapsed().as_secs_f64();
85    println!(
86        "pagesift+ocr: {} files, {failed} failed, {:.0} ms/file, wall {:.1} s, {:.2} pages/s",
87        files.len(),
88        total * 1000.0 / files.len() as f64,
89        wall,
90        files.len() as f64 / wall
91    );
92    let Some(e) = engine.engine() else { return };
93    let st = e.device_status();
94    println!("device: {:?}, fallback: {:?}", st.actual, st.fallback);
95    if let Some(g) = e.gpu_stats() {
96        println!(
97            "gpu: peak {} MiB, needed {} MiB, {} batches, {} lines, {:.1} lines/batch, {} splits, {} offloaded ops",
98            g.peak_bytes >> 20,
99            g.needed_bytes >> 20,
100            g.batches,
101            g.lines,
102            g.lines as f64 / g.batches.max(1) as f64,
103            g.splits,
104            g.offloaded_ops
105        );
106    }
107}

Source

pub fn gpu_stats(&self) -> Option<GpuStats>

GPU 的显存和批次统计；没在用 GPU 时为 None。

Examples found in repository ?

examples/pdf_ocr_dir.rs (line 95)

27fn main() {
28    let a: Vec<PathBuf> = std::env::args().skip(1).map(PathBuf::from).collect();
29    let (src, dst) = (&a[0], &a[1]);
30    let math = (a.len() >= 6).then(|| (a[4].clone(), a[5].clone()));
31    let device = match std::env::var("OCR_DEVICE").as_deref() {
32        Ok("gpu") => Device::Gpu { index: 0 },
33        _ => Device::Cpu,
34    };
35    let jobs: usize = std::env::var("OCR_JOBS")
36        .ok()
37        .and_then(|v| v.parse().ok())
38        .unwrap_or(1);
39    let engine =
40        Arc::new(PdfOcr::from_dirs(a[2].clone(), a[3].clone(), math, device).expect("载入模型"));
41    // 按字节数估的时限没算 OCR 的耗时：关掉，只靠逐页的识别
42    let opts = pagesift_pdf::Options {
43        ocr: Some(engine.clone()),
44        adaptive_timeout: false,
45        // 所有带图的页都识别（Firecrawl 的 mode: "ocr"）：MRC 页有文字层，不判为扫描页
46        force_ocr: std::env::var_os("PDF_OCR_FORCE").is_some(),
47        ..Default::default()
48    };
49    let mut files = Vec::new();
50    walk(src, &mut files);
51    files.sort();
52    let next = AtomicUsize::new(0);
53    let totals = Mutex::new((0.0f64, 0usize));
54    let wall = Instant::now();
55    std::thread::scope(|s| {
56        for _ in 0..jobs {
57            s.spawn(|| {
58                while let Some(f) = files.get(next.fetch_add(1, Ordering::Relaxed)) {
59                    let bytes = std::fs::read(f).unwrap();
60                    let t = Instant::now();
61                    let (md, bad) = match pagesift_pdf::extract_with(&bytes, &opts) {
62                        Ok(doc) => (doc.markdown, 0),
63                        Err(e) => {
64                            eprintln!("{}: {e}", f.display());
65                            (String::new(), 1)
66                        }
67                    };
68                    let secs = t.elapsed().as_secs_f64();
69                    {
70                        let mut g = totals.lock().unwrap();
71                        g.0 += secs;
72                        g.1 += bad;
73                    }
74                    eprintln!("{}\t{:.0} ms", f.display(), secs * 1e3);
75                    let rel = f.strip_prefix(src).unwrap().with_extension("");
76                    let out = dst.join(format!("{}_pg1_repeat1.md", rel.display()));
77                    std::fs::create_dir_all(out.parent().unwrap()).unwrap();
78                    std::fs::write(out, md).unwrap();
79                }
80            });
81        }
82    });
83    let (total, failed) = *totals.lock().unwrap();
84    let wall = wall.elapsed().as_secs_f64();
85    println!(
86        "pagesift+ocr: {} files, {failed} failed, {:.0} ms/file, wall {:.1} s, {:.2} pages/s",
87        files.len(),
88        total * 1000.0 / files.len() as f64,
89        wall,
90        files.len() as f64 / wall
91    );
92    let Some(e) = engine.engine() else { return };
93    let st = e.device_status();
94    println!("device: {:?}, fallback: {:?}", st.actual, st.fallback);
95    if let Some(g) = e.gpu_stats() {
96        println!(
97            "gpu: peak {} MiB, needed {} MiB, {} batches, {} lines, {:.1} lines/batch, {} splits, {} offloaded ops",
98            g.peak_bytes >> 20,
99            g.needed_bytes >> 20,
100            g.batches,
101            g.lines,
102            g.lines as f64 / g.batches.max(1) as f64,
103            g.splits,
104            g.offloaded_ops
105        );
106    }
107}

Source

pub fn with_gpu_ops<R>(&self, f: impl FnOnce() -> R) -> R

在 f 里，当前线程直接调 pagesift_infer::ops 和 nn 层的模型（公式、版面）按实测开销把大算子分流到 GPU。没有 GPU 时就是直接调 f。第一次调用会校准几十毫秒。流水线目前没有用它：真实的公式模型上结果和 CPU 有差异（见 Engine::run）。

Source

pub fn run(&self, page: &Image) -> Result<Vec<Line>, String>

识别整页图，按检测给的顺序（从上到下、从左到右）返回各行；有公式模型时公式排在最后。

Source

pub fn run_until( &self, page: &Image, deadline: Option<Instant>, ) -> Result<PageRun, String>

同 Engine::run，另外报告这一页用的设备和退回、部分失败的情况，到 deadline 就尽快返回。版面检测失败时退回没有公式的文字 OCR，个别公式失败时跳过（原因见 PageRun::partial_error）。版面和公式模型留在 CPU：B5 上经 Engine::with_gpu_ops 分流后更慢、分数更低，原因还没查明。