1use std::sync::mpsc;
14use std::sync::{Arc, Mutex, OnceLock};
15
16use pdfium_render::prelude::*;
17
18use docparser_core::error::{DocError, ErrorCode, Result};
19use docparser_core::model::{BlockKind, Document, Page, Rect};
20use docparser_core::plugin::{
21 BackendKind, FormatPlugin, ParseCapability, RawDoc, RawImage, RawPage, RawTextItem,
22};
23use docparser_core::registry::Registry;
24
25type Resp = Result<RawDoc>;
26
27pub fn find_library() -> Option<std::path::PathBuf> {
29 if let Ok(p) = std::env::var("DOCPARSER_PDFIUM_LIB") {
30 let p = std::path::PathBuf::from(p);
31 if p.exists() {
32 return Some(p);
33 }
34 }
35 let home = std::env::var("HOME").unwrap_or_default();
36 #[cfg(target_os = "macos")]
37 let name = "libpdfium.dylib";
38 #[cfg(target_os = "linux")]
39 let name = "libpdfium.so";
40 #[cfg(target_os = "windows")]
41 let name = "pdfium.dll";
42 if let Ok(exe) = std::env::current_exe() {
44 if let Some(dir) = exe.parent() {
45 let p = dir.join("..").join("lib").join(name);
46 if p.exists() {
47 return Some(p);
48 }
49 }
50 }
51 for base in [
52 format!("{home}/miniconda3/lib/python3.13/site-packages/liteparse"),
54 format!("{home}/miniconda3/lib/python3.13/site-packages/pypdfium2_raw"),
56 format!("{home}/miniconda3/envs/kb/lib/python3.13/site-packages/pypdfium2_raw"),
57 ] {
58 let p = std::path::Path::new(&base).join(name);
59 if p.exists() {
60 return Some(p);
61 }
62 }
63 None
64}
65
66fn classify(e: &PdfiumError) -> DocError {
67 let msg = format!("{e}").to_lowercase();
68 if msg.contains("password") || msg.contains("encrypt") || msg.contains("permission") {
69 DocError::encrypted(format!("PDF 加密或权限受限: {e}"))
70 } else {
71 DocError::corrupt(format!("PDF 打开/解析失败: {e}"))
72 }
73}
74
75static RENDER_SEQ: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
76
77pub fn find_ocr_binary() -> Option<std::path::PathBuf> {
78 if let Ok(p) = std::env::var("DOCPARSER_OCR_BIN") {
79 let p = std::path::PathBuf::from(p);
80 if p.exists() {
81 return Some(p);
82 }
83 }
84 if let Ok(exe) = std::env::current_exe() {
86 if let Some(dir) = exe.parent() {
87 let p = dir.join("docparser-ocr");
88 if p.exists() {
89 return Some(p);
90 }
91 }
92 }
93 let cands = [
95 std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../target/release/docparser-ocr"),
96 std::path::PathBuf::from("rust/target/release/docparser-ocr"),
97 ];
98 cands.into_iter().find(|c| c.exists())
99}
100
101pub fn find_ocr_models() -> Option<std::path::PathBuf> {
102 if let Ok(p) = std::env::var("DOCPARSER_OCR_MODELS") {
103 let p = std::path::PathBuf::from(p);
104 if p.is_dir() {
105 return Some(p);
106 }
107 }
108 if let Ok(exe) = std::env::current_exe() {
110 if let Some(dir) = exe.parent() {
111 let p = dir.join("..").join("models/ocr/ppocrv5");
112 if p.is_dir() {
113 return Some(p);
114 }
115 }
116 }
117 let cands = [
118 std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models/ocr/ppocrv5"),
119 std::path::PathBuf::from("models/ocr/ppocrv5"),
120 ];
121 cands.into_iter().find(|c| c.is_dir())
122}
123
124fn is_garbled_char(ch: char) -> bool {
126 let cp = ch as u32;
127 ch == '\u{fffd}' || (0xE000..=0xF8FF).contains(&cp)
128}
129
130fn page_needs_ocr(texts: &[String]) -> bool {
133 let chars: usize = texts.iter().map(|t| t.trim().chars().count()).sum();
134 if chars < 10 {
135 return true;
136 }
137 let garbled: usize = texts
138 .iter()
139 .flat_map(|t| t.chars())
140 .filter(|c| is_garbled_char(*c))
141 .count();
142 garbled as f32 / chars as f32 > 0.3
143}
144
145fn sniff_ext(data: &[u8]) -> Option<&'static str> {
147 if data.len() >= 8 && data.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
148 return Some("png");
149 }
150 if data.len() >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
151 return Some("jpg");
152 }
153 if data.len() >= 2 && data.starts_with(b"BM") {
154 return Some("bmp");
155 }
156 if data.len() >= 4 && data.starts_with(b"GIF8") {
157 return Some("gif");
158 }
159 if data.len() >= 12 && data.starts_with(b"RIFF") && &data[8..12] == b"WEBP" {
160 return Some("webp");
161 }
162 None
163}
164
165fn ocrable_image(im: &RawImage) -> bool {
167 if im.data.is_empty() {
168 return false;
169 }
170 if !matches!(
171 im.ext.as_str(),
172 "png" | "jpg" | "jpeg" | "bmp" | "gif" | "webp"
173 ) {
174 return false;
175 }
176 if let Some((w, h)) = im.px {
177 return w.min(h) >= 50;
178 }
179 false
180}
181
182fn collect_page_images(
189 doc: &PdfDocument,
190 page: &PdfPage,
191 h: f32,
192 extract_data: bool,
193) -> Vec<RawImage> {
194 let mut out: Vec<RawImage> = Vec::new();
195 descend(doc, page.objects().iter(), h, 0, extract_data, &mut out);
196 out
197}
198
199fn descend<'a>(
200 doc: &'a PdfDocument,
201 objs: impl Iterator<Item = PdfPageObject<'a>>,
202 h: f32,
203 depth: usize,
204 extract_data: bool,
205 out: &mut Vec<RawImage>,
206) {
207 if depth > 2 {
208 return;
209 }
210 for obj in objs {
211 if let Some(img) = obj.as_image_object() {
212 let bbox = obj.bounds().ok().map(|q| Rect {
213 x0: q.left().value,
214 y0: h - q.top().value,
215 x1: q.right().value,
216 y1: h - q.bottom().value,
217 });
218 let (px_w, px_h) = (
219 img.width().ok().unwrap_or(0) as u32,
220 img.height().ok().unwrap_or(0) as u32,
221 );
222 let mut data: Vec<u8> = Vec::new();
223 let mut ext: String = String::new();
224 if extract_data {
225 if let Ok(raw) = img.get_raw_image_data() {
227 if let Some(e) = sniff_ext(&raw) {
228 data = raw;
229 ext = e.to_string();
230 }
231 }
232 if data.is_empty() {
233 if let Ok(dyn_img) = img.get_processed_image(doc) {
234 let rgba = dyn_img.to_rgba8();
235 let (bw, bh) = rgba.dimensions();
236 let mut buf: Vec<u8> = Vec::new();
237 {
238 use image::ImageEncoder as _;
239 let enc = image::codecs::png::PngEncoder::new(std::io::BufWriter::new(
240 &mut buf,
241 ));
242 if enc
243 .write_image(&rgba, bw, bh, image::ExtendedColorType::Rgba8)
244 .is_ok()
245 {
246 data = buf;
247 ext = "png".to_string();
248 }
249 }
250 }
251 }
252 }
253 out.push(RawImage {
254 data,
255 bbox,
256 ext,
257 px: Some((px_w, px_h)),
258 ocr_text: None,
259 });
260 } else if let Some(form) = obj.as_x_object_form_object() {
261 descend(doc, form.iter(), h, depth + 1, extract_data, out);
262 }
263 }
264}
265
266fn ocr_images_enabled() -> bool {
268 std::env::var("DOCPARSER_OCR_IMAGES").as_deref() == Ok("1")
269}
270
271fn ocr_enabled() -> bool {
274 !matches!(
275 std::env::var("DOCPARSER_OCR_ENABLED").as_deref(),
276 Ok("0") | Ok("false")
277 )
278}
279
280fn parse_one(pdfium: &Pdfium, path: &str) -> Resp {
281 let doc = pdfium
282 .load_pdf_from_file(path, None)
283 .map_err(|e| classify(&e))?;
284 let dpi = 200.0_f32;
285 let mut pages = Vec::new();
286 let mut render_dir: Option<std::path::PathBuf> = None;
287 let mut render_files: Vec<(usize, std::path::PathBuf)> = Vec::new();
288 let mut img_dir: Option<std::path::PathBuf> = None;
290 let mut img_seq = 0usize;
291 let mut img_jobs: Vec<(usize, usize, usize)> = Vec::new();
292 let mut img_files: Vec<(usize, std::path::PathBuf)> = Vec::new();
293 let ocr_images = ocr_images_enabled();
294 let ocr_on = ocr_enabled();
295 for (pi, page) in doc.pages().iter().enumerate() {
296 let w = page.width().value;
297 let h = page.height().value;
298 let mut items: Vec<RawTextItem> = Vec::new();
299 let mut texts: Vec<String> = Vec::new();
300 for obj in page.objects().iter() {
301 if obj.as_text_object().is_none() {
302 continue;
303 }
304 let s = obj.as_text_object().unwrap().text();
305 if s.trim().is_empty() {
306 continue;
307 }
308 texts.push(s.clone());
309 let bbox = obj.bounds().ok().map(|q| Rect {
310 x0: q.left().value,
311 y0: h - q.top().value, x1: q.right().value,
313 y1: h - q.bottom().value,
314 });
315 items.push(RawTextItem { text: s, bbox });
316 }
317 let images = collect_page_images(&doc, &page, h, ocr_images);
318 if ocr_on && page_needs_ocr(&texts) {
321 if render_dir.is_none() {
322 let d = std::env::temp_dir().join(format!(
323 "docparser_pdf_render_{}_{}",
324 std::process::id(),
325 RENDER_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
326 ));
327 let _ = std::fs::create_dir_all(&d);
328 render_dir = Some(d);
329 }
330 let dir = render_dir.as_ref().unwrap();
331 let fn_name = format!("page-{pi:04}.png");
332 let out_path = dir.join(&fn_name);
333 let wpx = (w * dpi / 72.0).round() as i32;
334 let hpx = (h * dpi / 72.0).round() as i32;
335 let bitmap = page.render(wpx, hpx, None).map_err(|e| classify(&e));
336 if let Ok(bitmap) = bitmap {
337 let bw = bitmap.width() as u32;
338 let bh = bitmap.height() as u32;
339 let rgba = bitmap.as_rgba_bytes();
340 use image::ImageEncoder as _;
341 let enc = image::codecs::png::PngEncoder::new(std::io::BufWriter::new(
342 std::fs::File::create(&out_path).unwrap(),
343 ));
344 if enc
345 .write_image(&rgba, bw, bh, image::ExtendedColorType::Rgba8)
346 .is_ok()
347 {
348 render_files.push((pi, out_path));
349 }
350 }
351 } else if ocr_on && ocr_images {
352 for (ki, im) in images.iter().enumerate() {
354 if !ocrable_image(im) {
355 continue;
356 }
357 if img_dir.is_none() {
358 let d = std::env::temp_dir().join(format!(
359 "docparser_pdf_img_{}_{}",
360 std::process::id(),
361 RENDER_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
362 ));
363 let _ = std::fs::create_dir_all(&d);
364 img_dir = Some(d);
365 }
366 let dir = img_dir.as_ref().unwrap();
367 let fn_name = format!("img-{img_seq:05}.{}", im.ext);
368 let out_path = dir.join(&fn_name);
369 if std::fs::write(&out_path, &im.data).is_ok() {
370 img_files.push((img_seq, out_path));
371 img_jobs.push((pi, ki, img_seq));
372 img_seq += 1;
373 }
374 }
375 }
376 pages.push(RawPage {
377 name: None,
378 items,
379 images,
380 width: Some(w),
381 height: Some(h),
382 rotation: None,
383 ..Default::default()
384 });
385 }
386
387 let mut any_ocr = false;
389 if !render_files.is_empty() {
390 if let Some(lines) = run_ocr(render_dir.as_ref().unwrap(), &render_files) {
391 let scale = 72.0 / dpi;
392 for (pi, items) in lines {
393 if let Some(pg) = pages.get_mut(pi) {
394 pg.items = items
395 .into_iter()
396 .map(|(text, quad)| RawTextItem {
397 text,
398 bbox: Some(Rect {
399 x0: quad.iter().map(|p| p[0]).fold(f32::MAX, f32::min) * scale,
400 y0: quad.iter().map(|p| p[1]).fold(f32::MAX, f32::min) * scale,
401 x1: quad.iter().map(|p| p[0]).fold(f32::MIN, f32::max) * scale,
402 y1: quad.iter().map(|p| p[1]).fold(f32::MIN, f32::max) * scale,
403 }),
404 })
405 .collect();
406 any_ocr = true;
407 }
408 }
409 }
410 }
411
412 if !img_files.is_empty() {
414 if let Some(lines) = run_ocr(img_dir.as_ref().unwrap(), &img_files) {
415 for (gi, line_items) in lines {
416 let Some(&(pi, ki, _)) = img_jobs.iter().find(|(_, _, g)| *g == gi) else {
417 continue;
418 };
419 let Some(pg) = pages.get_mut(pi) else {
420 continue;
421 };
422 let Some(im) = pg.images.get_mut(ki) else {
423 continue;
424 };
425 let ocr: Vec<&str> = line_items
426 .iter()
427 .map(|(s, _)| s.as_str())
428 .filter(|s| !s.trim().is_empty())
429 .collect();
430 if !ocr.is_empty() {
431 im.ocr_text = Some(ocr.join("\n"));
432 any_ocr = true;
433 }
434 }
435 }
436 }
437
438 let hf = strip_headers_footers(&mut pages);
440
441 let mut raw = RawDoc {
442 pages,
443 meta: Default::default(),
444 needs_ocr: false,
445 };
446 if !hf.is_empty() {
447 raw.meta.insert("headers_footers".into(), hf.join(" | "));
448 }
449 raw.needs_ocr = !any_ocr
450 && raw.pages.iter().all(|p| {
451 p.items
452 .iter()
453 .map(|i| i.text.chars().count())
454 .sum::<usize>()
455 < 10
456 });
457 Ok(raw)
458}
459
460fn _norm_repeat(text: &str) -> String {
461 let mut out = String::with_capacity(text.len());
462 let mut last_digit = false;
463 for ch in text.chars() {
464 if ch.is_ascii_digit() {
465 if !last_digit {
466 out.push('#');
467 last_digit = true;
468 }
469 } else {
470 last_digit = false;
471 if !ch.is_whitespace() {
472 out.push(ch);
473 }
474 }
475 }
476 out
477}
478
479fn strip_headers_footers(pages: &mut [RawPage]) -> Vec<String> {
482 if pages.len() < 2 {
483 return Vec::new();
484 }
485 let n_with_h = pages.iter().filter(|p| p.height.is_some()).count();
486 if n_with_h == 0 {
487 return Vec::new();
488 }
489 let mut removed: Vec<String> = Vec::new();
490 for zone in ["top", "bottom"] {
491 let mut occ: std::collections::HashMap<String, Vec<(usize, usize)>> =
492 std::collections::HashMap::new(); for (pi, page) in pages.iter().enumerate() {
494 let Some(h) = page.height else { continue };
495 let band = if zone == "top" { h * 0.07 } else { h * 0.93 };
496 for (ii, it) in page.items.iter().enumerate() {
497 let Some(bb) = it.bbox else { continue };
498 let hit = if zone == "top" {
499 bb.y1 <= band
500 } else {
501 bb.y0 >= band
502 };
503 if hit {
504 occ.entry(_norm_repeat(&it.text))
505 .or_default()
506 .push((pi, ii));
507 }
508 }
509 }
510 for (norm, hits) in occ.iter() {
511 if norm.is_empty() {
512 continue;
513 }
514 let mut hit_pages: Vec<usize> = hits.iter().map(|(p, _)| *p).collect();
515 hit_pages.sort_unstable();
516 hit_pages.dedup();
517 if hit_pages.len() >= 2 && hit_pages.len() as f32 >= 0.5 * n_with_h as f32 {
518 for (pi, ii) in hits.iter().rev() {
519 pages[*pi].items.remove(*ii);
520 }
521 if !removed.iter().any(|r| r == norm) {
522 removed.push(norm.clone());
523 }
524 }
525 }
526 }
527 removed
528}
529
530struct ResidentOcr {
534 child: std::process::Child,
535 stdin: std::process::ChildStdin,
536 stdout: std::io::BufReader<std::process::ChildStdout>,
537}
538
539impl ResidentOcr {
540 fn spawn() -> Option<Self> {
541 let bin = find_ocr_binary()?;
542 let mut cmd = std::process::Command::new(&bin);
543 cmd.arg("serve")
544 .stdin(std::process::Stdio::piped())
545 .stdout(std::process::Stdio::piped())
546 .stderr(std::process::Stdio::null());
547 if let Some(m) = find_ocr_models() {
548 cmd.env("DOCPARSER_OCR_MODELS", m);
549 }
550 let mut child = cmd.spawn().ok()?;
551 let stdin = child.stdin.take()?;
552 let stdout = std::io::BufReader::new(child.stdout.take()?);
553 Some(Self {
554 child,
555 stdin,
556 stdout,
557 })
558 }
559
560 fn recognize(&mut self, path: &std::path::Path) -> Option<Vec<serde_json::Value>> {
561 use std::io::{BufRead as _, Write as _};
562 writeln!(self.stdin, "{}", path.display()).ok()?;
563 self.stdin.flush().ok()?;
564 let mut line = String::new();
565 if self.stdout.read_line(&mut line).ok()? == 0 {
566 return None;
567 }
568 let v: serde_json::Value = serde_json::from_str(&line).ok()?;
569 if v.get("error").is_some() {
570 return Some(Vec::new());
571 }
572 Some(
573 v.get("lines")
574 .and_then(|l| l.as_array())
575 .cloned()
576 .unwrap_or_default(),
577 )
578 }
579}
580
581static OCR_SERVE: OnceLock<Mutex<Option<ResidentOcr>>> = OnceLock::new();
582
583fn serve_ocr(
585 files: &[(usize, std::path::PathBuf)],
586) -> Option<Vec<(usize, Vec<serde_json::Value>)>> {
587 let lock = OCR_SERVE.get_or_init(|| Mutex::new(None));
588 let mut guard = lock.lock().ok()?;
589 if guard.is_none() {
590 *guard = ResidentOcr::spawn();
591 }
592 let srv = guard.as_mut()?;
593 let mut out = Vec::with_capacity(files.len());
594 for (key, path) in files {
595 match srv.recognize(path) {
596 Some(lines) => out.push((*key, lines)),
597 None => {
598 if let Some(mut dead) = guard.take() {
600 let _ = dead.child.kill();
601 }
602 return None;
603 }
604 }
605 }
606 Some(out)
607}
608
609fn lines_from_entries(entries: &[serde_json::Value]) -> Vec<(String, Vec<[f32; 2]>)> {
611 let mut items = Vec::new();
612 for ent in entries {
613 if let Some(s) = ent.get("score") {
614 if s.as_f64().unwrap_or(0.0) < 0.8 {
615 continue;
616 }
617 }
618 let text = ent
619 .get("text")
620 .and_then(|t| t.as_str())
621 .unwrap_or("")
622 .to_string();
623 if text.trim().is_empty() {
624 continue;
625 }
626 let mut quad: Vec<[f32; 2]> = Vec::new();
627 if let Some(boxes) = ent.get("box").and_then(|b| b.as_array()) {
628 for p in boxes {
629 let mut pt = [0.0_f32; 2];
630 if let (Some(x), Some(y)) = (
631 p.get(0).and_then(|x| x.as_f64()),
632 p.get(1).and_then(|y| y.as_f64()),
633 ) {
634 pt = [x as f32, y as f32];
635 }
636 quad.push(pt);
637 }
638 }
639 items.push((text, quad));
640 }
641 items.sort_by(|a, b| {
642 let ay = a.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
643 let by = b.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
644 ay.total_cmp(&by)
645 });
646 items
647}
648
649type OcrLine = (String, Vec<[f32; 2]>);
651
652type OcrPageLines = Vec<(usize, Vec<OcrLine>)>;
654
655fn run_ocr(dir: &std::path::Path, files: &[(usize, std::path::PathBuf)]) -> Option<OcrPageLines> {
659 if let Some(served) = serve_ocr(files) {
660 let mut out = Vec::new();
661 for (key, entries) in served {
662 let items = lines_from_entries(&entries);
663 if !items.is_empty() {
664 out.push((key, items));
665 }
666 }
667 return Some(out);
668 }
669 batch_ocr(dir, files)
670}
671
672fn batch_ocr(dir: &std::path::Path, files: &[(usize, std::path::PathBuf)]) -> Option<OcrPageLines> {
673 let bin = find_ocr_binary()?;
674 let out_json = dir.join("ocr.json");
675 let mut cmd = std::process::Command::new(&bin);
676 cmd.arg(dir).arg("--out").arg(&out_json);
677 if let Some(m) = find_ocr_models() {
678 cmd.env("DOCPARSER_OCR_MODELS", m);
679 }
680 let ok = cmd.output().map(|o| o.status.success()).unwrap_or(false);
681 if !ok {
682 return None;
683 }
684 let text = std::fs::read_to_string(&out_json).ok()?;
685 let v: serde_json::Value = serde_json::from_str(&text).ok()?;
686 let arr = v.as_array()?;
687 let mut by_file: std::collections::HashMap<String, Vec<serde_json::Value>> =
688 std::collections::HashMap::new();
689 for ent in arr {
690 if let (Some(f), Some(s)) = (ent.get("file"), ent.get("score")) {
691 if s.as_f64().unwrap_or(0.0) >= 0.8 {
692 by_file
693 .entry(f.as_str().unwrap_or("").to_string())
694 .or_default()
695 .push(ent.clone());
696 }
697 }
698 }
699 let mut out = Vec::new();
700 for (pi, fp) in files {
701 let key = fp
702 .file_name()
703 .and_then(|n| n.to_str())
704 .unwrap_or("")
705 .to_string();
706 let Some(entries) = by_file.get(&key) else {
707 continue;
708 };
709 let mut items = Vec::new();
710 for ent in entries {
711 let text = ent
712 .get("text")
713 .and_then(|t| t.as_str())
714 .unwrap_or("")
715 .to_string();
716 if text.trim().is_empty() {
717 continue;
718 }
719 let mut quad: Vec<[f32; 2]> = Vec::new();
720 if let Some(boxes) = ent.get("box").and_then(|b| b.as_array()) {
721 for p in boxes {
722 let mut pt = [0.0_f32; 2];
723 if let (Some(x), Some(y)) = (
724 p.get(0).and_then(|x| x.as_f64()),
725 p.get(1).and_then(|y| y.as_f64()),
726 ) {
727 pt = [x as f32, y as f32];
728 }
729 quad.push(pt);
730 }
731 }
732 items.push((text, quad));
733 }
734 items.sort_by(|a, b| {
735 let ay = a.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
736 let by = b.1.iter().map(|p| p[1]).fold(f32::MAX, f32::min);
737 ay.total_cmp(&by)
738 });
739 if !items.is_empty() {
740 out.push((*pi, items));
741 }
742 }
743 Some(out)
744}
745
746type WorkerReq = (String, mpsc::Sender<Resp>);
748
749type WorkerTx = Arc<Mutex<mpsc::Sender<WorkerReq>>>;
751
752fn worker_sender() -> Result<WorkerTx> {
753 static WORKER: OnceLock<std::result::Result<WorkerTx, String>> = OnceLock::new();
754 WORKER
755 .get_or_init(|| {
756 let file = find_library().ok_or_else(|| {
757 "未找到 PDFium 动态库(设 DOCPARSER_PDFIUM_LIB 或安装 pypdfium2)".to_string()
758 })?;
759 let (tx, rx) = mpsc::channel::<(String, mpsc::Sender<Resp>)>();
760 std::thread::Builder::new()
761 .name("docparser-pdfium".into())
762 .spawn(move || {
763 let lib = file.to_str().ok_or("PDFium 路径非 UTF-8")?.to_string();
764 let bindings = Pdfium::bind_to_library(&lib).map_err(|e| format!("{e}"))?;
765 let pdfium = Pdfium::new(bindings);
766 for (path, resp_tx) in rx.iter() {
767 let _ = resp_tx.send(parse_one(&pdfium, &path));
768 }
769 Ok::<(), String>(())
770 })
771 .map_err(|e| format!("{e}"))?;
772 Ok(Arc::new(Mutex::new(tx)))
773 })
774 .as_ref()
775 .map(Arc::clone)
776 .map_err(|m| DocError::corrupt(m.clone()))
777}
778
779pub struct PdfiumWorkerBackend;
781
782impl ParseCapability for PdfiumWorkerBackend {
783 fn id(&self) -> &str {
784 "pdfium-worker"
785 }
786 fn formats(&self) -> &[&str] {
787 &["pdf"]
788 }
789 fn kind(&self) -> BackendKind {
790 BackendKind::Bundled
791 }
792 fn is_available(&self) -> bool {
793 find_library().is_some()
794 }
795 fn parse(&self, path: &str) -> Result<RawDoc> {
796 let tx = worker_sender()?
797 .lock()
798 .map_err(|_| DocError::corrupt("worker 锁"))?
799 .clone();
800 let (resp_tx, resp_rx) = mpsc::channel();
801 tx.send((path.to_string(), resp_tx))
802 .map_err(|_| DocError::corrupt("PDFium worker 已退出"))?;
803 resp_rx
804 .recv()
805 .map_err(|_| DocError::corrupt("PDFium worker 无响应"))?
806 }
807}
808
809pub struct PdfTextBackend;
811
812impl ParseCapability for PdfTextBackend {
813 fn id(&self) -> &str {
814 "pdfium-text-v1"
815 }
816 fn formats(&self) -> &[&str] {
817 &["pdf"]
818 }
819 fn kind(&self) -> BackendKind {
820 BackendKind::Bundled
821 }
822 fn parse(&self, path: &str) -> Result<RawDoc> {
823 let per_page = std::panic::catch_unwind(|| pdf_extract::extract_text_by_pages(path))
824 .map_err(|_| {
825 DocError::corrupt(
826 "pdf-extract v1 无法解析该 PDF(CJK CID/无 ToUnicode,需 pdfium-worker)",
827 )
828 })?
829 .map_err(|e| {
830 let msg = format!("{e}").to_lowercase();
831 if msg.contains("password") || msg.contains("encrypt") || msg.contains("permission")
832 {
833 DocError::encrypted(format!("PDF 加密或权限受限: {e}"))
834 } else {
835 DocError::corrupt(format!("PDF 解析失败: {e}"))
836 }
837 })?;
838 let pages: Vec<RawPage> = per_page
839 .into_iter()
840 .map(|page_text| RawPage {
841 name: None,
842 items: page_text
843 .lines()
844 .map(|l| RawTextItem {
845 text: l.to_string(),
846 bbox: None,
847 })
848 .filter(|it| !it.text.trim().is_empty())
849 .collect(),
850 ..Default::default()
851 })
852 .filter(|p| !p.items.is_empty())
853 .collect();
854 if pages.is_empty() {
855 return Err(DocError::corrupt("PDF 无有效文本层"));
856 }
857 let mut raw = RawDoc {
858 pages,
859 meta: Default::default(),
860 needs_ocr: false,
861 };
862 raw.needs_ocr = raw.pages.iter().all(|p| {
863 p.items
864 .iter()
865 .map(|i| i.text.chars().count())
866 .sum::<usize>()
867 < 10
868 });
869 Ok(raw)
870 }
871}
872
873fn merge_lines(items: Vec<RawTextItem>) -> Vec<RawTextItem> {
878 let no_bb: Vec<RawTextItem> = items
879 .iter()
880 .filter(|it| it.bbox.is_none())
881 .cloned()
882 .collect();
883 let with_bb: Vec<(usize, RawTextItem)> = items
884 .into_iter()
885 .enumerate()
886 .filter(|(_, it)| it.bbox.is_some())
887 .collect();
888 let mut lines: Vec<Vec<(usize, RawTextItem)>> = Vec::new();
889 let mut sorted = with_bb.clone();
890 sorted.sort_by(|a, b| {
891 a.1.bbox
892 .unwrap()
893 .y0
894 .total_cmp(&b.1.bbox.unwrap().y0)
895 .then(a.1.bbox.unwrap().x0.total_cmp(&b.1.bbox.unwrap().x0))
896 });
897 for ent in &sorted {
898 let bb = ent.1.bbox.unwrap();
899 let mut placed = false;
900 for line in lines.iter_mut() {
901 let y0 = line
902 .iter()
903 .map(|e| e.1.bbox.unwrap().y0)
904 .fold(f32::MAX, f32::min);
905 let y1 = line
906 .iter()
907 .map(|e| e.1.bbox.unwrap().y1)
908 .fold(f32::MIN, f32::max);
909 let ov = (y1.min(bb.y1) - y0.max(bb.y0)).max(0.0);
910 if ov > 0.5 * (y1 - y0).min(bb.y1 - bb.y0) {
911 line.push(ent.clone());
912 placed = true;
913 break;
914 }
915 }
916 if !placed {
917 lines.push(vec![ent.clone()]);
918 }
919 }
920 let mut out: Vec<RawTextItem> = Vec::new();
921 for mut line in lines {
922 line.sort_by(|a, b| a.1.bbox.unwrap().x0.total_cmp(&b.1.bbox.unwrap().x0));
923 let mut cur: Option<(String, Rect)> = None;
924 for (_i, it) in line {
925 let bb = it.bbox.unwrap();
926 match &mut cur {
927 None => cur = Some((it.text.clone(), bb)),
928 Some((t, cb)) => {
929 let gap = bb.x0 - cb.x1;
930 let h = (cb.y1 - cb.y0).min(bb.y1 - bb.y0);
931 if gap <= (6.0_f32).max(0.6 * h) {
932 let sep = if gap > 0.2 * h { " " } else { "" };
933 *t = format!("{t}{sep}{}", it.text);
934 *cb = Rect {
935 x0: cb.x0.min(bb.x0),
936 y0: cb.y0.min(bb.y0),
937 x1: cb.x1.max(bb.x1),
938 y1: cb.y1.max(bb.y1),
939 };
940 } else {
941 out.push(RawTextItem {
942 text: std::mem::take(t),
943 bbox: Some(*cb),
944 });
945 cur = Some((it.text.clone(), bb));
946 }
947 }
948 }
949 }
950 if let Some((t, cb)) = cur {
951 out.push(RawTextItem {
952 text: t,
953 bbox: Some(cb),
954 });
955 }
956 }
957 out.splice(0..0, no_bb); out
959}
960
961fn detect_table(lines: &[RawTextItem]) -> Option<(Vec<Vec<String>>, Vec<bool>)> {
962 let mut rows: Vec<Vec<RawTextItem>> = Vec::new();
964 for it in lines {
965 let bb = it.bbox.unwrap();
966 let mut placed = false;
967 for row in rows.iter_mut() {
968 let ry0 = row
969 .iter()
970 .map(|r| r.bbox.unwrap().y0)
971 .fold(f32::MAX, f32::min);
972 let ry1 = row
973 .iter()
974 .map(|r| r.bbox.unwrap().y1)
975 .fold(f32::MIN, f32::max);
976 let ov = (ry1.min(bb.y1) - ry0.max(bb.y0)).max(0.0);
977 if ov > 0.5 * (ry1 - ry0).min(bb.y1 - bb.y0) {
978 row.push(it.clone());
979 placed = true;
980 break;
981 }
982 }
983 if !placed {
984 rows.push(vec![it.clone()]);
985 }
986 }
987 rows.sort_by(|a, b| {
988 a.iter()
989 .map(|r| r.bbox.unwrap().y0)
990 .fold(f32::MAX, f32::min)
991 .total_cmp(
992 &b.iter()
993 .map(|r| r.bbox.unwrap().y0)
994 .fold(f32::MAX, f32::min),
995 )
996 });
997 if rows.len() < 2 {
998 return None;
999 }
1000 let mut xs: Vec<f32> = lines.iter().map(|it| it.bbox.unwrap().x0).collect();
1002 xs.sort_by(f32::total_cmp);
1003 let mut anchors: Vec<f32> = Vec::new();
1004 for x in xs {
1005 if anchors.is_empty() || x - *anchors.last().unwrap() > 8.0 {
1006 anchors.push(x);
1007 }
1008 }
1009 if anchors.len() < 3 {
1010 return None;
1011 }
1012 let col_of = |x: f32| -> usize {
1013 let mut best = 0usize;
1014 let mut bd = f32::MAX;
1015 for (i, a) in anchors.iter().enumerate() {
1016 let d = (x - a).abs();
1017 if d < bd {
1018 bd = d;
1019 best = i;
1020 }
1021 }
1022 best
1023 };
1024 let mut grid: Vec<Vec<RawTextItem>> = Vec::new();
1025 for mut row in rows {
1026 row.sort_by(|a, b| a.bbox.unwrap().x0.total_cmp(&b.bbox.unwrap().x0));
1027 let mut merged: Vec<RawTextItem> = Vec::new();
1028 for it in row {
1029 let bb = it.bbox.unwrap();
1030 match merged.last_mut() {
1031 Some(prev) if col_of(prev.bbox.unwrap().x0) == col_of(bb.x0) => {
1032 prev.text = format!("{}{}", prev.text, it.text);
1033 let pb = prev.bbox.unwrap();
1034 prev.bbox = Some(Rect {
1035 x0: pb.x0.min(bb.x0),
1036 y0: pb.y0.min(bb.y0),
1037 x1: pb.x1.max(bb.x1),
1038 y1: pb.y1.max(bb.y1),
1039 });
1040 }
1041 _ => merged.push(it),
1042 }
1043 }
1044 if merged.len() >= 3 {
1045 grid.push(merged);
1046 }
1047 }
1048 if grid.len() < 2 {
1049 return None;
1050 }
1051 let counts: Vec<usize> = grid.iter().map(|r| r.len()).collect();
1052 let modal = counts
1053 .iter()
1054 .max_by_key(|c| counts.iter().filter(|x| x == c).count())
1055 .copied()
1056 .unwrap_or(0);
1057 if modal < 3 {
1058 return None;
1059 }
1060 let consistent = counts.iter().filter(|c| **c == modal).count() as f32 / grid.len() as f32;
1061 if consistent < 0.6 {
1062 return None;
1063 }
1064 let used_keys: Vec<(f32, f32)> = grid
1065 .iter()
1066 .flat_map(|r| r.iter())
1067 .map(|it| {
1068 let b = it.bbox.unwrap();
1069 (b.x0, b.y0)
1070 })
1071 .collect();
1072 let used: Vec<bool> = lines
1073 .iter()
1074 .map(|it| {
1075 let b = it.bbox.unwrap();
1076 used_keys.contains(&(b.x0, b.y0))
1077 })
1078 .collect();
1079 let rows_out: Vec<Vec<String>> = grid
1080 .into_iter()
1081 .map(|r| r.into_iter().map(|it| it.text).collect())
1082 .collect();
1083 Some((rows_out, used))
1084}
1085
1086fn column_order(items: Vec<RawTextItem>) -> Vec<RawTextItem> {
1089 let mut sorted = items.clone();
1090 sorted.sort_by(|a, b| {
1091 let (ax, _ay, ax1) = {
1092 let r = a.bbox.unwrap();
1093 (r.x0, r.y0, r.x1)
1094 };
1095 let (bx, _by, bx1) = {
1096 let r = b.bbox.unwrap();
1097 (r.x0, r.y0, r.x1)
1098 };
1099 ax.total_cmp(&bx).then(ax1.total_cmp(&bx1))
1100 });
1101 let mut cols: Vec<(f32, f32, Vec<RawTextItem>)> = Vec::new();
1103 for it in sorted {
1104 let bb = it.bbox.unwrap();
1105 if let Some(last) = cols.last_mut() {
1106 if bb.x0 <= last.1 + 6.0 {
1107 last.1 = last.1.max(bb.x1);
1108 last.2.push(it);
1109 continue;
1110 }
1111 }
1112 cols.push((bb.x0, bb.x1, vec![it]));
1113 }
1114 if cols.len() < 2 {
1115 return items; }
1117 let mut out = Vec::with_capacity(items.len());
1118 for (_x0, _x1, mut col_items) in cols {
1119 col_items.sort_by(|a, b| a.bbox.unwrap().y0.total_cmp(&b.bbox.unwrap().y0));
1120 out.extend(col_items);
1121 }
1122 out
1123}
1124
1125fn assemble(raw: &RawDoc) -> Document {
1126 let pages: Vec<Page> = raw
1127 .pages
1128 .iter()
1129 .map(|rp| {
1130 let lines = merge_lines(rp.items.clone());
1132 let all_bbox = lines.iter().all(|it| it.bbox.is_some());
1134 let (table_rows, used, table_rect) =
1135 match if all_bbox { detect_table(&lines) } else { None } {
1136 Some((rows, used)) => {
1137 let rect = {
1138 let hits: Vec<Rect> = lines
1139 .iter()
1140 .zip(&used)
1141 .filter(|(_, u)| **u)
1142 .map(|(it, _)| it.bbox.unwrap())
1143 .collect();
1144 if hits.is_empty() {
1145 None
1146 } else {
1147 Some(Rect {
1148 x0: hits.iter().map(|r| r.x0).fold(f32::MAX, f32::min),
1149 y0: hits.iter().map(|r| r.y0).fold(f32::MAX, f32::min),
1150 x1: hits.iter().map(|r| r.x1).fold(f32::MIN, f32::max),
1151 y1: hits.iter().map(|r| r.y1).fold(f32::MIN, f32::max),
1152 })
1153 }
1154 };
1155 (Some(rows), used, rect)
1156 }
1157 None => (None, vec![false; lines.len()], None),
1158 };
1159 let paras: Vec<RawTextItem> = lines
1161 .iter()
1162 .zip(&used)
1163 .filter(|(_, u)| !**u)
1164 .map(|(it, _)| it.clone())
1165 .collect();
1166 let ordered: Vec<RawTextItem> = if all_bbox { column_order(paras) } else { paras };
1167
1168 let mut blocks: Vec<docparser_core::model::Block> = Vec::new();
1170 if let Some(rows) = &table_rows {
1171 blocks.push(docparser_core::model::Block {
1172 kind: BlockKind::Table,
1173 table: Some(docparser_core::model::TableData { rows: rows.clone() }),
1174 bbox: table_rect,
1175 ..Default::default()
1176 });
1177 }
1178 for it in &ordered {
1179 blocks.push(docparser_core::model::Block {
1180 kind: BlockKind::Paragraph,
1181 content: it.text.clone(),
1182 bbox: it.bbox,
1183 ..Default::default()
1184 });
1185 }
1186 let mut parts: Vec<String> = ordered.iter().map(|it| it.text.clone()).collect();
1187 if let Some(rows) = &table_rows {
1188 let ty = table_rect.map(|r| r.y0).unwrap_or(0.0);
1189 let table_text = rows
1190 .iter()
1191 .map(|r| r.join(","))
1192 .collect::<Vec<_>>()
1193 .join("\n");
1194 let mut idx = parts.len();
1196 for (k, it) in ordered.iter().enumerate() {
1197 if it.bbox.map(|r| r.y0).unwrap_or(f32::MAX) > ty {
1198 idx = k;
1199 break;
1200 }
1201 }
1202 parts.insert(idx, table_text);
1203 }
1204 let imgs: Vec<docparser_core::model::Image> = rp
1206 .images
1207 .iter()
1208 .map(|im| docparser_core::model::Image {
1209 bbox: im.bbox,
1210 ocr_text: im.ocr_text.clone(),
1211 ..Default::default()
1212 })
1213 .collect();
1214 for im in &imgs {
1215 if let Some(t) = &im.ocr_text {
1216 if !t.trim().is_empty() {
1217 parts.push(t.clone());
1218 }
1219 }
1220 }
1221 let text = parts
1222 .into_iter()
1223 .filter(|s| !s.trim().is_empty())
1224 .collect::<Vec<_>>()
1225 .join("\n");
1226 Page {
1227 blocks,
1228 text,
1229 images: imgs,
1230 width: rp.width,
1231 height: rp.height,
1232 rotation: rp.rotation,
1233 ..Default::default()
1234 }
1235 })
1236 .collect();
1237 Document {
1238 pages,
1239 ..Default::default()
1240 }
1241}
1242
1243pub struct PdfPlugin;
1244
1245impl FormatPlugin for PdfPlugin {
1246 fn id(&self) -> &str {
1247 "pdf"
1248 }
1249 fn formats(&self) -> &[&str] {
1250 &["pdf"]
1251 }
1252 fn default_parse_chain(&self) -> Vec<&'static str> {
1253 vec!["pdfium-worker", "pdfium-text-v1"]
1255 }
1256 fn default_process_chain(&self) -> Vec<&'static str> {
1257 vec![]
1258 }
1259 fn assemble(&self, raw: &RawDoc) -> Document {
1260 assemble(raw)
1261 }
1262}
1263
1264pub fn register_defaults(reg: &mut Registry) {
1266 reg.register_parse(Arc::new(PdfTextBackend), None).unwrap();
1267 if find_library().is_some() {
1268 reg.register_parse(Arc::new(PdfiumWorkerBackend), None)
1269 .unwrap();
1270 }
1271 reg.register_plugin(Arc::new(PdfPlugin));
1272}
1273
1274pub fn is_terminal(e: &DocError) -> bool {
1276 matches!(e.code, ErrorCode::Encrypted | ErrorCode::Unsupported)
1277}
1278
1279#[cfg(test)]
1280mod tests {
1281 use super::*;
1282
1283 #[test]
1284 #[ignore = "需本机 PDFium 动态库"]
1285 fn probe_worker_cjk() {
1286 let p = "/Users/zerojzheng/CODE/hi-agents/documents/docparser/python/bench/corpus/text_plain.pdf";
1287 match worker_sender() {
1288 Err(e) => eprintln!("worker init: {e}"),
1289 Ok(tx) => {
1290 let tx = tx.lock().unwrap().clone();
1291 let (rtx, rrx) = mpsc::channel();
1292 tx.send((p.to_string(), rtx)).unwrap();
1293 match rrx.recv() {
1294 Ok(Ok(raw)) => {
1295 let n: usize = raw.pages.iter().map(|pg| pg.items.len()).sum();
1296 let s: usize = raw
1297 .pages
1298 .iter()
1299 .flat_map(|pg| &pg.items)
1300 .map(|i| i.text.chars().count())
1301 .sum();
1302 eprintln!("ok: pages={} items={} chars={}", raw.pages.len(), n, s);
1303 }
1304 Ok(Err(e)) => eprintln!("worker parse err: code={} msg={}", e.code, e.message),
1305 Err(e) => eprintln!("channel: {e}"),
1306 }
1307 }
1308 }
1309 }
1310}