pdf_ocr_dir/
pdf_ocr_dir.rs

1//! 带 OCR 的批量转换：把目录（递归）里的每个 PDF 转成 Markdown，扫描页用 PP-OCRv6 识别；
2//! 按 olmOCR-bench 的命名写出 `<输出目录>/<相对路径去掉 .pdf>_pg1_repeat1.md`。
3//!
4//!     pdf_ocr_dir <pdf 目录> <输出目录> <检测模型目录> <识别模型目录> [<版面模型目录> <公式模型目录>]
5//!
6//! 环境变量：`OCR_DEVICE=cpu|gpu`（默认 cpu）、`OCR_JOBS=N`（同时处理的 PDF 数，默认 1）。
7//! 结尾打印实际用的设备、显存峰值和平均批大小。
8use std::path::{Path, PathBuf};
9use std::sync::atomic::{AtomicUsize, Ordering};
10use std::sync::{Arc, Mutex};
11use std::time::Instant;
12
13use pagesift_ocr::device::Device;
14use pagesift_ocr::pdf::PdfOcr;
15
16fn walk(dir: &Path, out: &mut Vec<PathBuf>) {
17    for e in std::fs::read_dir(dir).unwrap().flatten() {
18        let p = e.path();
19        if p.is_dir() {
20            walk(&p, out);
21        } else if p.extension().is_some_and(|x| x == "pdf") {
22            out.push(p);
23        }
24    }
25}
26
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}
pdf_ocr_dir/pdf_ocr_dir.rs

pdf_ocr_dir/
pdf_ocr_dir.rs
