use std::sync::mpsc;
use std::sync::{Arc, Mutex, OnceLock};
use pdfium_render::prelude::*;
use docparser_core::error::{DocError, ErrorCode, Result};
use docparser_core::model::{BlockKind, Document, Page, Rect};
use docparser_core::plugin::{
BackendKind, FormatPlugin, ParseCapability, RawDoc, RawImage, RawPage, RawTextItem,
};
use docparser_core::registry::Registry;
type Resp = Result<RawDoc>;
pub fn find_library() -> Option<std::path::PathBuf> {
if let Ok(p) = std::env::var("DOCPARSER_PDFIUM_LIB") {
let p = std::path::PathBuf::from(p);
if p.exists() {
return Some(p);
}
}
let home = std::env::var("HOME").unwrap_or_default();
#[cfg(target_os = "macos")]
let name = "libpdfium.dylib";
#[cfg(target_os = "linux")]
let name = "libpdfium.so";
#[cfg(target_os = "windows")]
let name = "pdfium.dll";
if let Ok(exe) = std::env::current_exe() {
if let Some(dir) = exe.parent() {
let p = dir.join("..").join("lib").join(name);
if p.exists() {
return Some(p);
}
}
}
for base in [
format!("{home}/miniconda3/lib/python3.13/site-packages/liteparse"),
format!("{home}/miniconda3/lib/python3.13/site-packages/pypdfium2_raw"),
format!("{home}/miniconda3/envs/kb/lib/python3.13/site-packages/pypdfium2_raw"),
] {
let p = std::path::Path::new(&base).join(name);
if p.exists() {
return Some(p);
}
}
None
}
fn classify(e: &PdfiumError) -> DocError {
let msg = format!("{e}").to_lowercase();
if msg.contains("password") || msg.contains("encrypt") || msg.contains("permission") {
DocError::encrypted(format!("PDF 加密或权限受限: {e}"))
} else {
DocError::corrupt(format!("PDF 打开/解析失败: {e}"))
}
}
static RENDER_SEQ: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
pub fn find_ocr_binary() -> Option<std::path::PathBuf> {
if let Ok(p) = std::env::var("DOCPARSER_OCR_BIN") {
let p = std::path::PathBuf::from(p);
if p.exists() {
return Some(p);
}
}
if let Ok(exe) = std::env::current_exe() {
if let Some(dir) = exe.parent() {
let p = dir.join("docparser-ocr");
if p.exists() {
return Some(p);
}
}
}
let cands = [
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../target/release/docparser-ocr"),
std::path::PathBuf::from("rust/target/release/docparser-ocr"),
];
cands.into_iter().find(|c| c.exists())
}
pub fn find_ocr_models() -> Option<std::path::PathBuf> {
if let Ok(p) = std::env::var("DOCPARSER_OCR_MODELS") {
let p = std::path::PathBuf::from(p);
if p.is_dir() {
return Some(p);
}
}
if let Ok(exe) = std::env::current_exe() {
if let Some(dir) = exe.parent() {
let p = dir.join("..").join("models/ocr/ppocrv5");
if p.is_dir() {
return Some(p);
}
}
}
let cands = [
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models/ocr/ppocrv5"),
std::path::PathBuf::from("models/ocr/ppocrv5"),
];
cands.into_iter().find(|c| c.is_dir())
}
fn is_garbled_char(ch: char) -> bool {
let cp = ch as u32;
ch == '\u{fffd}' || (0xE000..=0xF8FF).contains(&cp)
}
fn page_needs_ocr(texts: &[String]) -> bool {
let chars: usize = texts.iter().map(|t| t.trim().chars().count()).sum();
if chars < 10 {
return true;
}
let garbled: usize = texts
.iter()
.flat_map(|t| t.chars())
.filter(|c| is_garbled_char(*c))
.count();
garbled as f32 / chars as f32 > 0.3
}
fn sniff_ext(data: &[u8]) -> Option<&'static str> {
if data.len() >= 8 && data.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
return Some("png");
}
if data.len() >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
return Some("jpg");
}
if data.len() >= 2 && data.starts_with(b"BM") {
return Some("bmp");
}
if data.len() >= 4 && data.starts_with(b"GIF8") {
return Some("gif");
}
if data.len() >= 12 && data.starts_with(b"RIFF") && &data[8..12] == b"WEBP" {
return Some("webp");
}
None
}
fn ocrable_image(im: &RawImage) -> bool {
if im.data.is_empty() {
return false;
}
if !matches!(
im.ext.as_str(),
"png" | "jpg" | "jpeg" | "bmp" | "gif" | "webp"
) {
return false;
}
if let Some((w, h)) = im.px {
return w.min(h) >= 50;
}
false
}
fn collect_page_images(
doc: &PdfDocument,
page: &PdfPage,
h: f32,
extract_data: bool,
) -> Vec<RawImage> {
let mut out: Vec<RawImage> = Vec::new();
descend(doc, page.objects().iter(), h, 0, extract_data, &mut out);
out
}
fn descend<'a>(
doc: &'a PdfDocument,
objs: impl Iterator<Item = PdfPageObject<'a>>,
h: f32,
depth: usize,
extract_data: bool,
out: &mut Vec<RawImage>,
) {
if depth > 2 {
return;
}
for obj in objs {
if let Some(img) = obj.as_image_object() {
let bbox = obj.bounds().ok().map(|q| Rect {
x0: q.left().value,
y0: h - q.top().value,
x1: q.right().value,
y1: h - q.bottom().value,
});
let (px_w, px_h) = (
img.width().ok().unwrap_or(0) as u32,
img.height().ok().unwrap_or(0) as u32,
);
let mut data: Vec<u8> = Vec::new();
let mut ext: String = String::new();
if extract_data {
if let Ok(raw) = img.get_raw_image_data() {
if let Some(e) = sniff_ext(&raw) {
data = raw;
ext = e.to_string();
}
}
if data.is_empty() {
if let Ok(dyn_img) = img.get_processed_image(doc) {
let rgba = dyn_img.to_rgba8();
let (bw, bh) = rgba.dimensions();
let mut buf: Vec<u8> = Vec::new();
{
use image::ImageEncoder as _;
let enc = image::codecs::png::PngEncoder::new(std::io::BufWriter::new(
&mut buf,
));
if enc
.write_image(&rgba, bw, bh, image::ExtendedColorType::Rgba8)
.is_ok()
{
data = buf;
ext = "png".to_string();
}
}
}
}
}
out.push(RawImage {
data,
bbox,
ext,
px: Some((px_w, px_h)),
ocr_text: None,
});
} else if let Some(form) = obj.as_x_object_form_object() {
descend(doc, form.iter(), h, depth + 1, extract_data, out);
}
}
}
fn ocr_images_enabled() -> bool {
std::env::var("DOCPARSER_OCR_IMAGES").as_deref() == Ok("1")
}
fn ocr_enabled() -> bool {
!matches!(
std::env::var("DOCPARSER_OCR_ENABLED").as_deref(),
Ok("0") | Ok("false")
)
}
fn parse_one(pdfium: &Pdfium, path: &str) -> Resp {
let doc = pdfium
.load_pdf_from_file(path, None)
.map_err(|e| classify(&e))?;
let dpi = 200.0_f32;
let mut pages = Vec::new();
let mut render_dir: Option<std::path::PathBuf> = None;
let mut render_files: Vec<(usize, std::path::PathBuf)> = Vec::new();
let mut img_dir: Option<std::path::PathBuf> = None;
let mut img_seq = 0usize;
let mut img_jobs: Vec<(usize, usize, usize)> = Vec::new();
let mut img_files: Vec<(usize, std::path::PathBuf)> = Vec::new();
let ocr_images = ocr_images_enabled();
let ocr_on = ocr_enabled();
for (pi, page) in doc.pages().iter().enumerate() {
let w = page.width().value;
let h = page.height().value;
let mut items: Vec<RawTextItem> = Vec::new();
let mut texts: Vec<String> = Vec::new();
for obj in page.objects().iter() {
if obj.as_text_object().is_none() {
continue;
}
let s = obj.as_text_object().unwrap().text();
if s.trim().is_empty() {
continue;
}
texts.push(s.clone());
let bbox = obj.bounds().ok().map(|q| Rect {
x0: q.left().value,
y0: h - q.top().value, x1: q.right().value,
y1: h - q.bottom().value,
});
items.push(RawTextItem { text: s, bbox });
}
let images = collect_page_images(&doc, &page, h, ocr_images);
if ocr_on && page_needs_ocr(&texts) {
if render_dir.is_none() {
let d = std::env::temp_dir().join(format!(
"docparser_pdf_render_{}_{}",
std::process::id(),
RENDER_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
));
let _ = std::fs::create_dir_all(&d);
render_dir = Some(d);
}
let dir = render_dir.as_ref().unwrap();
let fn_name = format!("page-{pi:04}.png");
let out_path = dir.join(&fn_name);
let wpx = (w * dpi / 72.0).round() as i32;
let hpx = (h * dpi / 72.0).round() as i32;
let bitmap = page.render(wpx, hpx, None).map_err(|e| classify(&e));
if let Ok(bitmap) = bitmap {
let bw = bitmap.width() as u32;
let bh = bitmap.height() as u32;
let rgba = bitmap.as_rgba_bytes();
use image::ImageEncoder as _;
let enc = image::codecs::png::PngEncoder::new(std::io::BufWriter::new(
std::fs::File::create(&out_path).unwrap(),
));
if enc
.write_image(&rgba, bw, bh, image::ExtendedColorType::Rgba8)
.is_ok()
{
render_files.push((pi, out_path));
}
}
} else if ocr_on && ocr_images {
for (ki, im) in images.iter().enumerate() {
if !ocrable_image(im) {
continue;
}
if img_dir.is_none() {
let d = std::env::temp_dir().join(format!(
"docparser_pdf_img_{}_{}",
std::process::id(),
RENDER_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
));
let _ = std::fs::create_dir_all(&d);
img_dir = Some(d);
}
let dir = img_dir.as_ref().unwrap();
let fn_name = format!("img-{img_seq:05}.{}", im.ext);
let out_path = dir.join(&fn_name);
if std::fs::write(&out_path, &im.data).is_ok() {
img_files.push((img_seq, out_path));
img_jobs.push((pi, ki, img_seq));
img_seq += 1;
}
}
}
pages.push(RawPage {
name: None,
items,
images,
width: Some(w),
height: Some(h),
rotation: None,
..Default::default()
});
}
let mut any_ocr = false;
if !render_files.is_empty() {
if let Some(lines) = run_ocr(render_dir.as_ref().unwrap(), &render_files) {
let scale = 72.0 / dpi;
for (pi, items) in lines {
if let Some(pg) = pages.get_mut(pi) {
pg.items = items
.into_iter()
.map(|(text, quad)| RawTextItem {
text,
bbox: Some(Rect {
x0: quad.iter().map(|p| p[0]).fold(f32::MAX, f32::min) * scale,
y0: quad.iter().map(|p| p[1]).fold(f32::MAX, f32::min) * scale,
x1: quad.iter().map(|p| p[0]).fold(f32::MIN, f32::max) * scale,
y1: quad.iter().map(|p| p[1]).fold(f32::MIN, f32::max) * scale,
}),
})
.collect();
any_ocr = true;
}
}
}
}
if !img_files.is_empty() {
if let Some(lines) = run_ocr(img_dir.as_ref().unwrap(), &img_files) {
for (gi, line_items) in lines {
let Some(&(pi, ki, _)) = img_jobs.iter().find(|(_, _, g)| *g == gi) else {
continue;
};
let Some(pg) = pages.get_mut(pi) else {
continue;
};
let Some(im) = pg.images.get_mut(ki) else {
continue;
};
let ocr: Vec<&str> = line_items
.iter()
.map(|(s, _)| s.as_str())
.filter(|s| !s.trim().is_empty())
.collect();
if !ocr.is_empty() {
im.ocr_text = Some(ocr.join("\n"));
any_ocr = true;
}
}
}
}
let hf = strip_headers_footers(&mut pages);
let mut raw = RawDoc {
pages,
meta: Default::default(),
needs_ocr: false,
};
if !hf.is_empty() {
raw.meta.insert("headers_footers".into(), hf.join(" | "));
}
raw.needs_ocr = !any_ocr
&& raw.pages.iter().all(|p| {
p.items
.iter()
.map(|i| i.text.chars().count())
.sum::<usize>()
< 10
});
Ok(raw)
}
fn _norm_repeat(text: &str) -> String {
let mut out = String::with_capacity(text.len());
let mut last_digit = false;
for ch in text.chars() {
if ch.is_ascii_digit() {
if !last_digit {
out.push('#');
last_digit = true;
}
} else {
last_digit = false;
if !ch.is_whitespace() {
out.push(ch);
}
}
}
out
}
fn strip_headers_footers(pages: &mut [RawPage]) -> Vec<String> {
if pages.len() < 2 {
return Vec::new();
}
let n_with_h = pages.iter().filter(|p| p.height.is_some()).count();
if n_with_h == 0 {
return Vec::new();
}
let mut removed: Vec<String> = Vec::new();
for zone in ["top", "bottom"] {
let mut occ: std::collections::HashMap<String, Vec<(usize, usize)>> =
std::collections::HashMap::new(); for (pi, page) in pages.iter().enumerate() {
let Some(h) = page.height else { continue };
let band = if zone == "top" { h * 0.07 } else { h * 0.93 };
for (ii, it) in page.items.iter().enumerate() {
let Some(bb) = it.bbox else { continue };
let hit = if zone == "top" {
bb.y1 <= band
} else {
bb.y0 >= band
};
if hit {
occ.entry(_norm_repeat(&it.text))
.or_default()
.push((pi, ii));
}
}
}
for (norm, hits) in occ.iter() {
if norm.is_empty() {
continue;
}
let mut hit_pages: Vec<usize> = hits.iter().map(|(p, _)| *p).collect();
hit_pages.sort_unstable();
hit_pages.dedup();
if hit_pages.len() >= 2 && hit_pages.len() as f32 >= 0.5 * n_with_h as f32 {
for (pi, ii) in hits.iter().rev() {
pages[*pi].items.remove(*ii);
}
if !removed.iter().any(|r| r == norm) {
removed.push(norm.clone());
}
}
}
}
removed
}
struct ResidentOcr {
child: std::process::Child,
stdin: std::process::ChildStdin,
stdout: std::io::BufReader<std::process::ChildStdout>,
}
impl ResidentOcr {
fn spawn() -> Option<Self> {
let bin = find_ocr_binary()?;
let mut cmd = std::process::Command::new(&bin);
cmd.arg("serve")
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::null());
if let Some(m) = find_ocr_models() {
cmd.env("DOCPARSER_OCR_MODELS", m);
}
let mut child = cmd.spawn().ok()?;
let stdin = child.stdin.take()?;
let stdout = std::io::BufReader::new(child.stdout.take()?);
Some(Self {
child,
stdin,
stdout,
})
}
fn recognize(&mut self, path: &std::path::Path) -> Option<Vec<serde_json::Value>> {
use std::io::{BufRead as _, Write as _};
writeln!(self.stdin, "{}", path.display()).ok()?;
self.stdin.flush().ok()?;
let mut line = String::new();
if self.stdout.read_line(&mut line).ok()? == 0 {
return None;
}
let v: serde_json::Value = serde_json::from_str(&line).ok()?;
if v.get("error").is_some() {
return Some(Vec::new());
}
Some(
v.get("lines")
.and_then(|l| l.as_array())
.cloned()
.unwrap_or_default(),
)
}
}
static OCR_SERVE: OnceLock<Mutex<Option<ResidentOcr>>> = OnceLock::new();
fn serve_ocr(
files: &[(usize, std::path::PathBuf)],
) -> Option<Vec<(usize, Vec<serde_json::Value>)>> {
let lock = OCR_SERVE.get_or_init(|| Mutex::new(None));
let mut guard = lock.lock().ok()?;
if guard.is_none() {
*guard = ResidentOcr::spawn();
}
let srv = guard.as_mut()?;
let mut out = Vec::with_capacity(files.len());
for (key, path) in files {
match srv.recognize(path) {
Some(lines) => out.push((*key, lines)),
None => {
if let Some(mut dead) = guard.take() {
let _ = dead.child.kill();
}
return None;
}
}
}
Some(out)
}
fn lines_from_entries(entries: &[serde_json::Value]) -> Vec<(String, Vec<[f32; 2]>)> {
let mut items = Vec::new();
for ent in entries {
if let Some(s) = ent.get("score") {
if s.as_f64().unwrap_or(0.0) < 0.8 {
continue;
}
}
let text = ent
.get("text")
.and_then(|t| t.as_str())
.unwrap_or("")
.to_string();
if text.trim().is_empty() {
continue;
}
let mut quad: Vec<[f32; 2]> = Vec::new();
if let Some(boxes) = ent.get("box").and_then(|b| b.as_array()) {
for p in boxes {
let mut pt = [0.0_f32; 2];
if let (Some(x), Some(y)) = (
p.get(0).and_then(|x| x.as_f64()),
p.get(1).and_then(|y| y.as_f64()),
) {
pt = [x as f32, y as f32];
}
quad.push(pt);
}
}
items.push((text, quad));
}
items.sort_by(|a, b| {
let ay = a.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
let by = b.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
ay.total_cmp(&by)
});
items
}
type OcrLine = (String, Vec<[f32; 2]>);
type OcrPageLines = Vec<(usize, Vec<OcrLine>)>;
fn run_ocr(dir: &std::path::Path, files: &[(usize, std::path::PathBuf)]) -> Option<OcrPageLines> {
if let Some(served) = serve_ocr(files) {
let mut out = Vec::new();
for (key, entries) in served {
let items = lines_from_entries(&entries);
if !items.is_empty() {
out.push((key, items));
}
}
return Some(out);
}
batch_ocr(dir, files)
}
fn batch_ocr(dir: &std::path::Path, files: &[(usize, std::path::PathBuf)]) -> Option<OcrPageLines> {
let bin = find_ocr_binary()?;
let out_json = dir.join("ocr.json");
let mut cmd = std::process::Command::new(&bin);
cmd.arg(dir).arg("--out").arg(&out_json);
if let Some(m) = find_ocr_models() {
cmd.env("DOCPARSER_OCR_MODELS", m);
}
let ok = cmd.output().map(|o| o.status.success()).unwrap_or(false);
if !ok {
return None;
}
let text = std::fs::read_to_string(&out_json).ok()?;
let v: serde_json::Value = serde_json::from_str(&text).ok()?;
let arr = v.as_array()?;
let mut by_file: std::collections::HashMap<String, Vec<serde_json::Value>> =
std::collections::HashMap::new();
for ent in arr {
if let (Some(f), Some(s)) = (ent.get("file"), ent.get("score")) {
if s.as_f64().unwrap_or(0.0) >= 0.8 {
by_file
.entry(f.as_str().unwrap_or("").to_string())
.or_default()
.push(ent.clone());
}
}
}
let mut out = Vec::new();
for (pi, fp) in files {
let key = fp
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("")
.to_string();
let Some(entries) = by_file.get(&key) else {
continue;
};
let mut items = Vec::new();
for ent in entries {
let text = ent
.get("text")
.and_then(|t| t.as_str())
.unwrap_or("")
.to_string();
if text.trim().is_empty() {
continue;
}
let mut quad: Vec<[f32; 2]> = Vec::new();
if let Some(boxes) = ent.get("box").and_then(|b| b.as_array()) {
for p in boxes {
let mut pt = [0.0_f32; 2];
if let (Some(x), Some(y)) = (
p.get(0).and_then(|x| x.as_f64()),
p.get(1).and_then(|y| y.as_f64()),
) {
pt = [x as f32, y as f32];
}
quad.push(pt);
}
}
items.push((text, quad));
}
items.sort_by(|a, b| {
let ay = a.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
let by = b.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
ay.total_cmp(&by)
});
if !items.is_empty() {
out.push((*pi, items));
}
}
Some(out)
}
type WorkerReq = (String, mpsc::Sender<Resp>);
type WorkerTx = Arc<Mutex<mpsc::Sender<WorkerReq>>>;
fn worker_sender() -> Result<WorkerTx> {
static WORKER: OnceLock<std::result::Result<WorkerTx, String>> = OnceLock::new();
WORKER
.get_or_init(|| {
let file = find_library().ok_or_else(|| {
"未找到 PDFium 动态库(设 DOCPARSER_PDFIUM_LIB 或安装 pypdfium2)".to_string()
})?;
let (tx, rx) = mpsc::channel::<(String, mpsc::Sender<Resp>)>();
std::thread::Builder::new()
.name("docparser-pdfium".into())
.spawn(move || {
let lib = file.to_str().ok_or("PDFium 路径非 UTF-8")?.to_string();
let bindings = Pdfium::bind_to_library(&lib).map_err(|e| format!("{e}"))?;
let pdfium = Pdfium::new(bindings);
for (path, resp_tx) in rx.iter() {
let _ = resp_tx.send(parse_one(&pdfium, &path));
}
Ok::<(), String>(())
})
.map_err(|e| format!("{e}"))?;
Ok(Arc::new(Mutex::new(tx)))
})
.as_ref()
.map(Arc::clone)
.map_err(|m| DocError::corrupt(m.clone()))
}
pub struct PdfiumWorkerBackend;
impl ParseCapability for PdfiumWorkerBackend {
fn id(&self) -> &str {
"pdfium-worker"
}
fn formats(&self) -> &[&str] {
&["pdf"]
}
fn kind(&self) -> BackendKind {
BackendKind::Bundled
}
fn is_available(&self) -> bool {
find_library().is_some()
}
fn parse(&self, path: &str) -> Result<RawDoc> {
let tx = worker_sender()?
.lock()
.map_err(|_| DocError::corrupt("worker 锁"))?
.clone();
let (resp_tx, resp_rx) = mpsc::channel();
tx.send((path.to_string(), resp_tx))
.map_err(|_| DocError::corrupt("PDFium worker 已退出"))?;
resp_rx
.recv()
.map_err(|_| DocError::corrupt("PDFium worker 无响应"))?
}
}
pub struct PdfTextBackend;
impl ParseCapability for PdfTextBackend {
fn id(&self) -> &str {
"pdfium-text-v1"
}
fn formats(&self) -> &[&str] {
&["pdf"]
}
fn kind(&self) -> BackendKind {
BackendKind::Bundled
}
fn parse(&self, path: &str) -> Result<RawDoc> {
let per_page = std::panic::catch_unwind(|| pdf_extract::extract_text_by_pages(path))
.map_err(|_| {
DocError::corrupt(
"pdf-extract v1 无法解析该 PDF(CJK CID/无 ToUnicode,需 pdfium-worker)",
)
})?
.map_err(|e| {
let msg = format!("{e}").to_lowercase();
if msg.contains("password") || msg.contains("encrypt") || msg.contains("permission")
{
DocError::encrypted(format!("PDF 加密或权限受限: {e}"))
} else {
DocError::corrupt(format!("PDF 解析失败: {e}"))
}
})?;
let pages: Vec<RawPage> = per_page
.into_iter()
.map(|page_text| RawPage {
name: None,
items: page_text
.lines()
.map(|l| RawTextItem {
text: l.to_string(),
bbox: None,
})
.filter(|it| !it.text.trim().is_empty())
.collect(),
..Default::default()
})
.filter(|p| !p.items.is_empty())
.collect();
if pages.is_empty() {
return Err(DocError::corrupt("PDF 无有效文本层"));
}
let mut raw = RawDoc {
pages,
meta: Default::default(),
needs_ocr: false,
};
raw.needs_ocr = raw.pages.iter().all(|p| {
p.items
.iter()
.map(|i| i.text.chars().count())
.sum::<usize>()
< 10
});
Ok(raw)
}
}
fn merge_lines(items: Vec<RawTextItem>) -> Vec<RawTextItem> {
let no_bb: Vec<RawTextItem> = items
.iter()
.filter(|it| it.bbox.is_none())
.cloned()
.collect();
let with_bb: Vec<(usize, RawTextItem)> = items
.into_iter()
.enumerate()
.filter(|(_, it)| it.bbox.is_some())
.collect();
let mut lines: Vec<Vec<(usize, RawTextItem)>> = Vec::new();
let mut sorted = with_bb.clone();
sorted.sort_by(|a, b| {
a.1.bbox
.unwrap()
.y0
.total_cmp(&b.1.bbox.unwrap().y0)
.then(a.1.bbox.unwrap().x0.total_cmp(&b.1.bbox.unwrap().x0))
});
for ent in &sorted {
let bb = ent.1.bbox.unwrap();
let mut placed = false;
for line in lines.iter_mut() {
let y0 = line
.iter()
.map(|e| e.1.bbox.unwrap().y0)
.fold(f32::MAX, f32::min);
let y1 = line
.iter()
.map(|e| e.1.bbox.unwrap().y1)
.fold(f32::MIN, f32::max);
let ov = (y1.min(bb.y1) - y0.max(bb.y0)).max(0.0);
if ov > 0.5 * (y1 - y0).min(bb.y1 - bb.y0) {
line.push(ent.clone());
placed = true;
break;
}
}
if !placed {
lines.push(vec![ent.clone()]);
}
}
let mut out: Vec<RawTextItem> = Vec::new();
for mut line in lines {
line.sort_by(|a, b| a.1.bbox.unwrap().x0.total_cmp(&b.1.bbox.unwrap().x0));
let mut cur: Option<(String, Rect)> = None;
for (_i, it) in line {
let bb = it.bbox.unwrap();
match &mut cur {
None => cur = Some((it.text.clone(), bb)),
Some((t, cb)) => {
let gap = bb.x0 - cb.x1;
let h = (cb.y1 - cb.y0).min(bb.y1 - bb.y0);
if gap <= (6.0_f32).max(0.6 * h) {
let sep = if gap > 0.2 * h { " " } else { "" };
*t = format!("{t}{sep}{}", it.text);
*cb = Rect {
x0: cb.x0.min(bb.x0),
y0: cb.y0.min(bb.y0),
x1: cb.x1.max(bb.x1),
y1: cb.y1.max(bb.y1),
};
} else {
out.push(RawTextItem {
text: std::mem::take(t),
bbox: Some(*cb),
});
cur = Some((it.text.clone(), bb));
}
}
}
}
if let Some((t, cb)) = cur {
out.push(RawTextItem {
text: t,
bbox: Some(cb),
});
}
}
out.splice(0..0, no_bb); out
}
fn detect_table(lines: &[RawTextItem]) -> Option<(Vec<Vec<String>>, Vec<bool>)> {
let mut rows: Vec<Vec<RawTextItem>> = Vec::new();
for it in lines {
let bb = it.bbox.unwrap();
let mut placed = false;
for row in rows.iter_mut() {
let ry0 = row
.iter()
.map(|r| r.bbox.unwrap().y0)
.fold(f32::MAX, f32::min);
let ry1 = row
.iter()
.map(|r| r.bbox.unwrap().y1)
.fold(f32::MIN, f32::max);
let ov = (ry1.min(bb.y1) - ry0.max(bb.y0)).max(0.0);
if ov > 0.5 * (ry1 - ry0).min(bb.y1 - bb.y0) {
row.push(it.clone());
placed = true;
break;
}
}
if !placed {
rows.push(vec![it.clone()]);
}
}
rows.sort_by(|a, b| {
a.iter()
.map(|r| r.bbox.unwrap().y0)
.fold(f32::MAX, f32::min)
.total_cmp(
&b.iter()
.map(|r| r.bbox.unwrap().y0)
.fold(f32::MAX, f32::min),
)
});
if rows.len() < 2 {
return None;
}
let mut xs: Vec<f32> = lines.iter().map(|it| it.bbox.unwrap().x0).collect();
xs.sort_by(f32::total_cmp);
let mut anchors: Vec<f32> = Vec::new();
for x in xs {
if anchors.is_empty() || x - *anchors.last().unwrap() > 8.0 {
anchors.push(x);
}
}
if anchors.len() < 3 {
return None;
}
let col_of = |x: f32| -> usize {
let mut best = 0usize;
let mut bd = f32::MAX;
for (i, a) in anchors.iter().enumerate() {
let d = (x - a).abs();
if d < bd {
bd = d;
best = i;
}
}
best
};
let mut grid: Vec<Vec<RawTextItem>> = Vec::new();
for mut row in rows {
row.sort_by(|a, b| a.bbox.unwrap().x0.total_cmp(&b.bbox.unwrap().x0));
let mut merged: Vec<RawTextItem> = Vec::new();
for it in row {
let bb = it.bbox.unwrap();
match merged.last_mut() {
Some(prev) if col_of(prev.bbox.unwrap().x0) == col_of(bb.x0) => {
prev.text = format!("{}{}", prev.text, it.text);
let pb = prev.bbox.unwrap();
prev.bbox = Some(Rect {
x0: pb.x0.min(bb.x0),
y0: pb.y0.min(bb.y0),
x1: pb.x1.max(bb.x1),
y1: pb.y1.max(bb.y1),
});
}
_ => merged.push(it),
}
}
if merged.len() >= 3 {
grid.push(merged);
}
}
if grid.len() < 2 {
return None;
}
let counts: Vec<usize> = grid.iter().map(|r| r.len()).collect();
let modal = counts
.iter()
.max_by_key(|c| counts.iter().filter(|x| x == c).count())
.copied()
.unwrap_or(0);
if modal < 3 {
return None;
}
let consistent = counts.iter().filter(|c| **c == modal).count() as f32 / grid.len() as f32;
if consistent < 0.6 {
return None;
}
let used_keys: Vec<(f32, f32)> = grid
.iter()
.flat_map(|r| r.iter())
.map(|it| {
let b = it.bbox.unwrap();
(b.x0, b.y0)
})
.collect();
let used: Vec<bool> = lines
.iter()
.map(|it| {
let b = it.bbox.unwrap();
used_keys.contains(&(b.x0, b.y0))
})
.collect();
let rows_out: Vec<Vec<String>> = grid
.into_iter()
.map(|r| r.into_iter().map(|it| it.text).collect())
.collect();
Some((rows_out, used))
}
fn column_order(items: Vec<RawTextItem>) -> Vec<RawTextItem> {
let mut sorted = items.clone();
sorted.sort_by(|a, b| {
let (ax, _ay, ax1) = {
let r = a.bbox.unwrap();
(r.x0, r.y0, r.x1)
};
let (bx, _by, bx1) = {
let r = b.bbox.unwrap();
(r.x0, r.y0, r.x1)
};
ax.total_cmp(&bx).then(ax1.total_cmp(&bx1))
});
let mut cols: Vec<(f32, f32, Vec<RawTextItem>)> = Vec::new();
for it in sorted {
let bb = it.bbox.unwrap();
if let Some(last) = cols.last_mut() {
if bb.x0 <= last.1 + 6.0 {
last.1 = last.1.max(bb.x1);
last.2.push(it);
continue;
}
}
cols.push((bb.x0, bb.x1, vec![it]));
}
if cols.len() < 2 {
return items; }
let mut out = Vec::with_capacity(items.len());
for (_x0, _x1, mut col_items) in cols {
col_items.sort_by(|a, b| a.bbox.unwrap().y0.total_cmp(&b.bbox.unwrap().y0));
out.extend(col_items);
}
out
}
fn assemble(raw: &RawDoc) -> Document {
let pages: Vec<Page> = raw
.pages
.iter()
.map(|rp| {
let lines = merge_lines(rp.items.clone());
let all_bbox = lines.iter().all(|it| it.bbox.is_some());
let (table_rows, used, table_rect) =
match if all_bbox { detect_table(&lines) } else { None } {
Some((rows, used)) => {
let rect = {
let hits: Vec<Rect> = lines
.iter()
.zip(&used)
.filter(|(_, u)| **u)
.map(|(it, _)| it.bbox.unwrap())
.collect();
if hits.is_empty() {
None
} else {
Some(Rect {
x0: hits.iter().map(|r| r.x0).fold(f32::MAX, f32::min),
y0: hits.iter().map(|r| r.y0).fold(f32::MAX, f32::min),
x1: hits.iter().map(|r| r.x1).fold(f32::MIN, f32::max),
y1: hits.iter().map(|r| r.y1).fold(f32::MIN, f32::max),
})
}
};
(Some(rows), used, rect)
}
None => (None, vec![false; lines.len()], None),
};
let paras: Vec<RawTextItem> = lines
.iter()
.zip(&used)
.filter(|(_, u)| !**u)
.map(|(it, _)| it.clone())
.collect();
let ordered: Vec<RawTextItem> = if all_bbox { column_order(paras) } else { paras };
let mut blocks: Vec<docparser_core::model::Block> = Vec::new();
if let Some(rows) = &table_rows {
blocks.push(docparser_core::model::Block {
kind: BlockKind::Table,
table: Some(docparser_core::model::TableData { rows: rows.clone() }),
bbox: table_rect,
..Default::default()
});
}
for it in &ordered {
blocks.push(docparser_core::model::Block {
kind: BlockKind::Paragraph,
content: it.text.clone(),
bbox: it.bbox,
..Default::default()
});
}
let mut parts: Vec<String> = ordered.iter().map(|it| it.text.clone()).collect();
if let Some(rows) = &table_rows {
let ty = table_rect.map(|r| r.y0).unwrap_or(0.0);
let table_text = rows
.iter()
.map(|r| r.join(","))
.collect::<Vec<_>>()
.join("\n");
let mut idx = parts.len();
for (k, it) in ordered.iter().enumerate() {
if it.bbox.map(|r| r.y0).unwrap_or(f32::MAX) > ty {
idx = k;
break;
}
}
parts.insert(idx, table_text);
}
let imgs: Vec<docparser_core::model::Image> = rp
.images
.iter()
.map(|im| docparser_core::model::Image {
bbox: im.bbox,
ocr_text: im.ocr_text.clone(),
..Default::default()
})
.collect();
for im in &imgs {
if let Some(t) = &im.ocr_text {
if !t.trim().is_empty() {
parts.push(t.clone());
}
}
}
let text = parts
.into_iter()
.filter(|s| !s.trim().is_empty())
.collect::<Vec<_>>()
.join("\n");
Page {
blocks,
text,
images: imgs,
width: rp.width,
height: rp.height,
rotation: rp.rotation,
..Default::default()
}
})
.collect();
Document {
pages,
..Default::default()
}
}
pub struct PdfPlugin;
impl FormatPlugin for PdfPlugin {
fn id(&self) -> &str {
"pdf"
}
fn formats(&self) -> &[&str] {
&["pdf"]
}
fn default_parse_chain(&self) -> Vec<&'static str> {
vec!["pdfium-worker", "pdfium-text-v1"]
}
fn default_process_chain(&self) -> Vec<&'static str> {
vec![]
}
fn assemble(&self, raw: &RawDoc) -> Document {
assemble(raw)
}
}
pub fn register_defaults(reg: &mut Registry) {
reg.register_parse(Arc::new(PdfTextBackend), None).unwrap();
if find_library().is_some() {
reg.register_parse(Arc::new(PdfiumWorkerBackend), None)
.unwrap();
}
reg.register_plugin(Arc::new(PdfPlugin));
}
pub fn is_terminal(e: &DocError) -> bool {
matches!(e.code, ErrorCode::Encrypted | ErrorCode::Unsupported)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[ignore = "需本机 PDFium 动态库"]
fn probe_worker_cjk() {
let p = "/Users/zerojzheng/CODE/hi-agents/documents/docparser/python/bench/corpus/text_plain.pdf";
match worker_sender() {
Err(e) => eprintln!("worker init: {e}"),
Ok(tx) => {
let tx = tx.lock().unwrap().clone();
let (rtx, rrx) = mpsc::channel();
tx.send((p.to_string(), rtx)).unwrap();
match rrx.recv() {
Ok(Ok(raw)) => {
let n: usize = raw.pages.iter().map(|pg| pg.items.len()).sum();
let s: usize = raw
.pages
.iter()
.flat_map(|pg| &pg.items)
.map(|i| i.text.chars().count())
.sum();
eprintln!("ok: pages={} items={} chars={}", raw.pages.len(), n, s);
}
Ok(Err(e)) => eprintln!("worker parse err: code={} msg={}", e.code, e.message),
Err(e) => eprintln!("channel: {e}"),
}
}
}
}
}