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>;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum LayoutMode {
#[default]
Auto,
Rebuild,
Raw,
}
impl LayoutMode {
pub fn as_str(self) -> &'static str {
match self {
LayoutMode::Auto => "auto",
LayoutMode::Rebuild => "rebuild",
LayoutMode::Raw => "raw",
}
}
pub fn parse(s: &str) -> Option<Self> {
match s.trim().to_ascii_lowercase().as_str() {
"auto" => Some(LayoutMode::Auto),
"rebuild" => Some(LayoutMode::Rebuild),
"raw" => Some(LayoutMode::Raw),
_ => None,
}
}
pub fn reflows(self) -> bool {
!matches!(self, LayoutMode::Raw)
}
}
#[derive(Debug, Clone, Default)]
pub struct PdfOptions {
pub layout_mode: LayoutMode,
pub max_pages: Option<usize>,
}
impl PdfOptions {
pub fn new() -> Self {
Self::default()
}
pub fn with_layout_mode(mut self, mode: LayoutMode) -> Self {
self.layout_mode = mode;
self
}
pub fn with_max_pages(mut self, pages: usize) -> Self {
self.max_pages = Some(pages);
self
}
}
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);
}
}
#[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 home in home_dirs() {
for base in [
home.join("miniconda3/lib/python3.13/site-packages/liteparse"),
home.join("miniconda3/lib/python3.13/site-packages/pypdfium2_raw"),
home.join("miniconda3/envs/kb/lib/python3.13/site-packages/pypdfium2_raw"),
home.join("anaconda3/lib/python3.13/site-packages/pypdfium2_raw"),
] {
let p = base.join(name);
if p.exists() {
return Some(p);
}
}
}
for sp in python_site_packages() {
for pkg in ["pypdfium2_raw", "liteparse"] {
let p = sp.join(pkg).join(name);
if p.exists() {
return Some(p);
}
}
}
None
}
fn home_dirs() -> Vec<std::path::PathBuf> {
let mut out: Vec<std::path::PathBuf> = Vec::new();
for key in ["HOME", "USERPROFILE"] {
if let Ok(v) = std::env::var(key) {
let p = std::path::PathBuf::from(v);
if p.is_dir() && !out.contains(&p) {
out.push(p);
}
}
}
if let (Ok(drive), Ok(path)) = (std::env::var("HOMEDRIVE"), std::env::var("HOMEPATH")) {
let p = std::path::PathBuf::from(format!("{drive}{path}"));
if p.is_dir() && !out.contains(&p) {
out.push(p);
}
}
out
}
fn python_site_packages() -> Vec<std::path::PathBuf> {
let mut roots: Vec<std::path::PathBuf> = Vec::new();
for home in home_dirs() {
for conda in ["miniconda3", "anaconda3", "miniforge3"] {
roots.push(home.join(conda).join("lib"));
roots.push(home.join(conda).join("Lib"));
}
roots.push(home.join("AppData/Local/Programs/Python"));
roots.push(home.join("AppData/Roaming/Python"));
roots.push(home.join(".local/lib"));
roots.push(home.join("Library/Python"));
}
#[cfg(target_os = "windows")]
for key in ["LOCALAPPDATA", "PROGRAMFILES", "PROGRAMFILES(X86)"] {
if let Ok(v) = std::env::var(key) {
roots.push(std::path::PathBuf::from(&v).join("Programs/Python"));
roots.push(std::path::PathBuf::from(&v));
}
}
roots.push(std::path::PathBuf::from("/usr/lib"));
roots.push(std::path::PathBuf::from("/usr/local/lib"));
let mut out: Vec<std::path::PathBuf> = Vec::new();
let mut push = |p: std::path::PathBuf| {
if p.is_dir() && !out.contains(&p) {
out.push(p);
}
};
for root in roots {
if !root.is_dir() {
continue;
}
push(root.join("site-packages"));
let Ok(entries) = std::fs::read_dir(&root) else {
continue;
};
for e in entries.flatten() {
let p = e.path();
if !p.is_dir() {
continue;
}
let name = p.file_name().and_then(|s| s.to_str()).unwrap_or_default();
if name.starts_with("python") || name.starts_with("Python") {
push(p.join("site-packages"));
push(p.join("Lib/site-packages"));
} else if name == "Lib" || name == "lib" {
push(p.join("site-packages"));
}
}
}
out
}
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);
}
}
let name = format!("docparser-ocr{}", std::env::consts::EXE_SUFFIX);
if let Ok(exe) = std::env::current_exe() {
if let Some(dir) = exe.parent() {
let p = dir.join(&name);
if p.exists() {
return Some(p);
}
}
}
let target = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../target");
let cands = [
target.join("release").join(&name),
target.join("debug").join(&name),
std::path::PathBuf::from("rust/target/release").join(&name),
std::path::PathBuf::from("rust/target/debug").join(&name),
];
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 image_dedup_key(im: &RawImage) -> Option<String> {
if im.data.is_empty() {
let (w, h) = im.px?;
return Some(format!("{w}x{h}@{}", im.ext));
}
use md5::Digest as _;
Some(format!("{:x}", md5::Md5::digest(&im.data)))
}
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 extract_text_items(page: &PdfPage<'_>, page_h: f32) -> (Vec<RawTextItem>, Vec<String>) {
let mut items: Vec<RawTextItem> = Vec::new();
let mut texts: Vec<String> = Vec::new();
if let Ok(text_page) = page.text() {
for seg in text_page.segments().iter() {
let s = seg.text();
if s.trim().is_empty() {
continue;
}
texts.push(s.clone());
let b = seg.bounds();
let height = (b.top().value - b.bottom().value).abs();
if page_h > 0.0 && height > page_h * 0.25 {
continue;
}
items.push(RawTextItem {
text: s,
bbox: Some(Rect {
x0: b.left().value,
y0: page_h - b.top().value,
x1: b.right().value,
y1: page_h - b.bottom().value,
}),
..Default::default()
});
}
if items.is_empty() {
let whole = text_page.all();
if !whole.trim().is_empty() {
if texts.is_empty() {
texts.push(whole.clone());
}
items.push(RawTextItem {
text: whole,
bbox: None,
..Default::default()
});
}
}
}
if texts.is_empty() {
for obj in page.objects().iter() {
let Some(t) = obj.as_text_object() else {
continue;
};
let s = t.text();
if s.trim().is_empty() {
continue;
}
texts.push(s.clone());
let bbox = obj.bounds().ok().map(|q| Rect {
x0: q.left().value,
y0: page_h - q.top().value,
x1: q.right().value,
y1: page_h - q.bottom().value,
});
items.push(RawTextItem {
text: s,
bbox,
..Default::default()
});
}
}
(items, texts)
}
fn parse_one(pdfium: &Pdfium, path: &str, options: &PdfOptions) -> Resp {
let doc = pdfium
.load_pdf_from_file(path, None)
.map_err(|e| classify(&e))?;
let dpi = 200.0_f32;
let mut meta: std::collections::HashMap<String, String> = std::collections::HashMap::new();
meta.insert("pdf.version".into(), pdf_version_str(doc.version()));
for (key, tag) in [
("pdf.title", PdfDocumentMetadataTagType::Title),
("pdf.author", PdfDocumentMetadataTagType::Author),
("pdf.producer", PdfDocumentMetadataTagType::Producer),
("pdf.creator", PdfDocumentMetadataTagType::Creator),
] {
if let Some(tag_value) = doc.metadata().get(tag) {
let v = tag_value.value().trim().to_string();
if !v.is_empty() {
meta.insert(key.into(), v);
}
}
}
let total_pages = doc.pages().len() as usize;
let limit = options.max_pages.unwrap_or(total_pages).min(total_pages);
meta.insert("pdf.page_count".into(), total_pages.to_string());
meta.insert("pdf.truncated".into(), (limit < total_pages).to_string());
let mut complexity: Vec<serde_json::Value> = Vec::new();
let mut page_errors: Vec<serde_json::Value> = Vec::new();
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().take(limit).enumerate() {
let w = page.width().value;
let h = page.height().value;
let (items, texts) = extract_text_items(&page, h);
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;
match page.render(wpx, hpx, None) {
Ok(bitmap) => {
let bw = bitmap.width() as u32;
let bh = bitmap.height() as u32;
let rgba = bitmap.as_rgba_bytes();
use image::ImageEncoder as _;
let Ok(file) = std::fs::File::create(&out_path) else {
page_errors.push(serde_json::json!({
"page": pi, "error": "无法创建渲染临时文件",
}));
continue;
};
let enc = image::codecs::png::PngEncoder::new(std::io::BufWriter::new(file));
if enc
.write_image(&rgba, bw, bh, image::ExtendedColorType::Rgba8)
.is_ok()
{
render_files.push((pi, out_path));
}
}
Err(e) => page_errors.push(serde_json::json!({
"page": pi, "error": format!("{e}"),
})),
}
} 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;
}
}
}
let chars: usize = texts.iter().map(|t| t.trim().chars().count()).sum();
let garbled = texts
.iter()
.flat_map(|t| t.chars())
.filter(|c| is_garbled_char(*c))
.count();
complexity.push(serde_json::json!({
"page": pi,
"chars": chars,
"garbled_ratio": if chars == 0 { 0.0 } else { garbled as f64 / chars as f64 },
"needs_ocr": page_needs_ocr(&texts),
}));
let rotation = page.rotation().ok().map(|r| match r {
PdfPageRenderRotation::None => 0,
PdfPageRenderRotation::Degrees90 => 90,
PdfPageRenderRotation::Degrees180 => 180,
PdfPageRenderRotation::Degrees270 => 270,
});
pages.push(RawPage {
name: None,
items,
images,
width: Some(w),
height: Some(h),
rotation,
..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,
}),
..Default::default()
})
.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,
needs_ocr: false,
};
if !hf.is_empty() {
raw.meta.insert("headers_footers".into(), hf.join(" | "));
}
if let Ok(s) = serde_json::to_string(&complexity) {
raw.meta.insert("pdf.complexity".into(), s);
}
if !page_errors.is_empty() {
if let Ok(s) = serde_json::to_string(&page_errors) {
raw.meta.insert("pdf.page_errors".into(), s);
}
}
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 pdf_version_str(v: PdfDocumentVersion) -> String {
match v {
PdfDocumentVersion::Unset => String::new(),
PdfDocumentVersion::Pdf1_0 => "1.0".into(),
PdfDocumentVersion::Pdf1_1 => "1.1".into(),
PdfDocumentVersion::Pdf1_2 => "1.2".into(),
PdfDocumentVersion::Pdf1_3 => "1.3".into(),
PdfDocumentVersion::Pdf1_4 => "1.4".into(),
PdfDocumentVersion::Pdf1_5 => "1.5".into(),
PdfDocumentVersion::Pdf1_6 => "1.6".into(),
PdfDocumentVersion::Pdf1_7 => "1.7".into(),
PdfDocumentVersion::Pdf2_0 => "2.0".into(),
PdfDocumentVersion::Other(n) => format!("{}.{}", n / 10, n % 10),
}
}
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));
}
}
}
let mut keys: Vec<(&String, &Vec<(usize, usize)>)> = occ.iter().collect();
keys.sort_by_key(|(_, hits)| hits.first().copied().unwrap_or((usize::MAX, usize::MAX)));
let mut drop: Vec<std::collections::BTreeSet<usize>> =
vec![std::collections::BTreeSet::new(); pages.len()];
for (norm, hits) in keys {
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 {
drop[*pi].insert(*ii);
}
if !removed.iter().any(|r| r == norm) {
removed.push(norm.clone());
}
}
}
for (pi, page) in pages.iter_mut().enumerate() {
if drop[pi].is_empty() {
continue;
}
let keep = std::mem::take(&mut page.items)
.into_iter()
.enumerate()
.filter(|(i, _)| !drop[pi].contains(i))
.map(|(_, it)| it)
.collect();
page.items = keep;
}
}
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, PdfOptions, 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, PdfOptions, 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, options, resp_tx) in rx.iter() {
let _ = resp_tx.send(parse_one(&pdfium, &path, &options));
}
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()))
}
#[derive(Debug, Clone, Default)]
pub struct PdfiumWorkerBackend {
options: PdfOptions,
}
impl PdfiumWorkerBackend {
pub fn new() -> Self {
Self::default()
}
pub fn with_options(mut self, options: PdfOptions) -> Self {
self.options = options;
self
}
}
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(), self.options.clone(), resp_tx))
.map_err(|_| DocError::corrupt("PDFium worker 已退出"))?;
resp_rx
.recv()
.map_err(|_| DocError::corrupt("PDFium worker 无响应"))?
}
}
#[derive(Debug, Clone, Default)]
pub struct PdfTextBackend {
options: PdfOptions,
}
impl PdfTextBackend {
pub fn new() -> Self {
Self::default()
}
pub fn with_options(mut self, options: PdfOptions) -> Self {
self.options = options;
self
}
}
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 total = per_page.len();
let limit = self.options.max_pages.unwrap_or(total).min(total);
let truncated = limit < total;
let pages: Vec<RawPage> = per_page
.into_iter()
.take(limit)
.map(|page_text| RawPage {
name: None,
items: page_text
.lines()
.map(|l| RawTextItem {
text: l.to_string(),
..Default::default()
})
.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.meta
.insert("pdf.truncated".into(), truncated.to_string());
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 lb = line
.iter()
.map(|e| e.1.bbox.unwrap().y1)
.fold(f32::MAX, f32::min);
let lh = line
.iter()
.map(|e| e.1.bbox.unwrap().y1 - e.1.bbox.unwrap().y0)
.fold(f32::MIN, f32::max);
let bh = bb.y1 - bb.y0;
if (lb - bb.y1).abs() <= 0.3 * (lh + bh) {
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),
..Default::default()
});
cur = Some((it.text.clone(), bb));
}
}
}
}
if let Some((t, cb)) = cur {
out.push(RawTextItem {
text: t,
bbox: Some(cb),
..Default::default()
});
}
}
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();
let mut sorted: Vec<&RawTextItem> = lines.iter().collect();
sorted.sort_by(|a, b| {
let (ab, bb) = (a.bbox.unwrap(), b.bbox.unwrap());
ab.y0.total_cmp(&bb.y0).then(ab.x0.total_cmp(&bb.x0))
});
for it in sorted {
let bb = it.bbox.unwrap();
let bh = bb.y1 - bb.y0;
let mut placed = false;
for row in rows.iter_mut() {
let lh = row
.iter()
.map(|r| r.bbox.unwrap().y1 - r.bbox.unwrap().y0)
.fold(f32::MIN, f32::max);
let lb = row
.iter()
.map(|r| r.bbox.unwrap().y1)
.fold(f32::MAX, f32::min);
if (lb - bb.y1).abs() <= 0.3 * (lh + bh) {
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 trect = {
let hits: Vec<Rect> = lines
.iter()
.zip(&used)
.filter(|(_, u)| **u)
.map(|(it, _)| it.bbox.unwrap())
.collect();
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),
}
};
let inside = |b: Rect| {
b.x0 >= trect.x0 - 1.0
&& b.x1 <= trect.x1 + 1.0
&& b.y0 >= trect.y0 - 1.0
&& b.y1 <= trect.y1 + 1.0
};
let cells_len: Vec<usize> = grid
.iter()
.flat_map(|r| r.iter())
.map(|it| it.text.chars().count())
.collect();
let in_table: usize = cells_len.iter().sum();
let long_ratio = if cells_len.is_empty() {
0.0
} else {
cells_len.iter().filter(|n| **n > 40).count() as f32 / cells_len.len() as f32
};
let stray: usize = lines
.iter()
.zip(&used)
.filter(|(it, u)| !**u && inside(it.bbox.unwrap()))
.map(|(it, _)| it.text.chars().count())
.sum();
if stray > in_table && long_ratio >= 0.2 {
return None;
}
let lopsided = grid
.iter()
.filter(|r| {
let lens: Vec<usize> = r.iter().map(|it| it.text.chars().count()).collect();
if lens.len() < 2 {
return false;
}
let max = *lens.iter().max().unwrap_or(&0);
if max < 40 {
return false;
}
let rest = lens
.iter()
.filter(|n| **n != max)
.max()
.copied()
.unwrap_or(0);
rest > 0 && max >= 3 * rest
})
.count();
if !grid.is_empty() && lopsided as f32 / grid.len() as f32 >= 0.5 {
return None;
}
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> {
if items.len() < 4 {
return items;
}
let (x_min, x_max) = items.iter().fold((f32::MAX, f32::MIN), |(lo, hi), it| {
let b = it.bbox.unwrap();
(lo.min(b.x0), hi.max(b.x1))
});
let span = x_max - x_min;
if span <= 1.0 || interleaved_ratio(&items, span) < 0.25 {
return items;
}
if let Some(out) = order_by_gutter(&items, x_min, x_max, span) {
return out;
}
order_by_leads(items, span)
}
fn order_by_gutter(
items: &[RawTextItem],
x_min: f32,
x_max: f32,
span: f32,
) -> Option<Vec<RawTextItem>> {
const MIN_GUTTER: usize = 20;
let n = span as usize + 1;
if n <= MIN_GUTTER * 3 {
return None;
}
let wide = span * 0.6;
let mut occ = vec![0usize; n];
for it in items {
let b = it.bbox.unwrap();
if (b.x1 - b.x0) > wide {
continue;
}
let lo = ((b.x0 - x_min).max(0.0) as usize).min(n - 1);
let hi = (((b.x1 - x_min) as usize).min(n - 1)).max(lo);
for v in occ.iter_mut().take(hi + 1).skip(lo) {
*v += 1;
}
}
let mut gutters: Vec<(usize, usize)> = Vec::new();
let mut run: Option<usize> = None;
for (i, v) in occ.iter().enumerate() {
if *v == 0 {
if run.is_none() {
run = Some(i);
}
} else if let Some(s) = run.take() {
if i - s >= MIN_GUTTER {
gutters.push((s, i));
}
}
}
if let Some(s) = run {
if n - s >= MIN_GUTTER {
gutters.push((s, n));
}
}
if gutters.is_empty() {
return None;
}
let center = |b: Rect| (b.x0 + b.x1) / 2.0;
let count_in = |lo: f32, hi: f32| {
items
.iter()
.filter(|it| {
let c = center(it.bbox.unwrap());
c >= lo && c < hi
})
.count()
};
let mut bounds: Vec<(f32, f32)> = Vec::new();
let mut prev = x_min;
for (g0, g1) in &gutters {
let seg = (prev, x_min + *g0 as f32);
if count_in(seg.0, seg.1) >= 2 {
bounds.push(seg);
prev = x_min + *g1 as f32;
}
}
let last = (prev, x_max);
if !bounds.is_empty() && count_in(last.0, last.1) >= 2 {
bounds.push(last);
}
if bounds.len() < 2 {
return None;
}
let mut out: Vec<RawTextItem> = Vec::with_capacity(items.len());
for (lo, hi) in &bounds {
let mut col: Vec<RawTextItem> = items
.iter()
.filter(|it| {
let c = center(it.bbox.unwrap());
c >= *lo && c < *hi
})
.cloned()
.collect();
col.sort_by(|a, b| {
let (ab, bb) = (a.bbox.unwrap(), b.bbox.unwrap());
ab.y0.total_cmp(&bb.y0).then(ab.x0.total_cmp(&bb.x0))
});
out.extend(col);
}
if out.len() != items.len() {
return None;
}
Some(out)
}
fn visual_line_leads(items: &[RawTextItem]) -> Vec<(f32, f32)> {
let mut order: Vec<&RawTextItem> = items.iter().collect();
order.sort_by(|a, b| {
let (ab, bb) = (a.bbox.unwrap(), b.bbox.unwrap());
ab.y0.total_cmp(&bb.y0).then(ab.x0.total_cmp(&bb.x0))
});
let mut bands: Vec<Vec<Rect>> = Vec::new();
for it in order {
let b = it.bbox.unwrap();
let bh = b.y1 - b.y0;
let mut placed = false;
for band in bands.iter_mut() {
let lb = band.iter().map(|r| r.y1).fold(f32::MAX, f32::min);
let lh = band.iter().map(|r| r.y1 - r.y0).fold(f32::MIN, f32::max);
if (lb - b.y1).abs() <= 0.3 * (lh + bh) {
band.push(b);
placed = true;
break;
}
}
if !placed {
bands.push(vec![b]);
}
}
let frag = |seg: &[Rect]| -> (f32, f32) {
let x0 = seg.iter().map(|r| r.x0).fold(f32::MAX, f32::min);
let x1 = seg.iter().map(|r| r.x1).fold(f32::MIN, f32::max);
(x0, x1 - x0)
};
let mut out: Vec<(f32, f32)> = Vec::new();
for mut band in bands {
band.sort_by(|a, b| a.x0.total_cmp(&b.x0));
let mut cur: Vec<Rect> = vec![band[0]];
for b in band.into_iter().skip(1) {
let prev = *cur.last().unwrap();
let h = ((prev.y1 - prev.y0).min(b.y1 - b.y0)).max(1.0);
if b.x0 - prev.x1 > (2.0 * h).max(12.0) {
out.push(frag(&cur));
cur = vec![b];
} else {
cur.push(b);
}
}
out.push(frag(&cur));
}
out
}
fn order_by_leads(items: Vec<RawTextItem>, span: f32) -> Vec<RawTextItem> {
let wide = span * 0.6;
let per_line = visual_line_leads(&items);
if per_line.len() < 4 {
return items;
}
let mut leads: Vec<f32> = per_line
.iter()
.filter(|(_, w)| *w <= wide)
.map(|(x, _)| *x)
.collect();
if leads.len() < 4 {
return items;
}
leads.sort_by(|a, b| a.total_cmp(b));
let lead_span = leads.last().copied().unwrap_or(0.0) - leads.first().copied().unwrap_or(0.0);
if lead_span <= 40.0 {
return items;
}
let thr = (lead_span * 0.25).max(40.0);
let mut groups: Vec<Vec<f32>> = Vec::new();
for a in leads {
match groups.last_mut() {
Some(g) if a - *g.last().unwrap() <= thr => g.push(a),
_ => groups.push(vec![a]),
}
}
let cols: Vec<f32> = groups
.iter()
.filter(|g| g.len() >= 3)
.map(|g| g[0])
.collect();
if cols.len() < 2 {
return items;
}
let total = items.len();
let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); cols.len()];
for (i, it) in items.iter().enumerate() {
let bx = it.bbox.unwrap().x0;
let mut idx = 0;
for (c, s) in cols.iter().enumerate() {
if bx >= *s - 0.5 {
idx = c;
}
}
buckets[idx].push(i);
}
if buckets[0].len() < 3 || buckets[1].len() < 3 {
return items;
}
let left_ys: Vec<f32> = buckets[0]
.iter()
.map(|&i| items[i].bbox.unwrap().y0)
.collect();
let paired = buckets[1]
.iter()
.filter(|&&i| {
let y = items[i].bbox.unwrap().y0;
left_ys.iter().any(|ly| (y - ly).abs() < 12.0)
})
.count();
if (paired as f32) / (buckets[1].len() as f32) < 0.3 {
return items;
}
let mut out = Vec::with_capacity(total);
for bucket in &mut buckets {
bucket.sort_by(|&a, &b| {
let (ay, ax) = (items[a].bbox.unwrap().y0, items[a].bbox.unwrap().x0);
let (by, bx) = (items[b].bbox.unwrap().y0, items[b].bbox.unwrap().x0);
ay.total_cmp(&by).then(ax.total_cmp(&bx))
});
out.extend(bucket.iter().map(|&i| items[i].clone()));
}
debug_assert_eq!(out.len(), total);
out
}
fn interleaved_ratio(items: &[RawTextItem], span: f32) -> f32 {
if items.len() < 2 {
return 0.0;
}
let half = (span * 0.5).max(1.0);
let mut hits = 0usize;
for w in items.windows(2) {
let a = w[0].bbox.unwrap();
let b = w[1].bbox.unwrap();
if (b.y0 - a.y0).abs() < 3.0 && (b.x0 - a.x0).abs() >= half {
hits += 1;
}
}
hits as f32 / (items.len() - 1) as f32
}
fn assemble(raw: &RawDoc, mode: LayoutMode) -> Document {
let pages: Vec<Page> = raw
.pages
.iter()
.map(|rp| {
let (lines, all_bbox) = if mode.reflows() {
let lines = merge_lines(rp.items.clone());
let all_bbox = lines.iter().all(|it| it.bbox.is_some());
(lines, all_bbox)
} else {
(rp.items.clone(), false)
};
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(),
..Default::default()
}),
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,
dedup_key: image_dedup_key(im),
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();
let mut doc = Document {
pages,
..Default::default()
};
let mut pdf: serde_json::Map<String, serde_json::Value> = serde_json::Map::new();
for (src, dst) in [
("pdf.version", "pdf_version"),
("pdf.title", "title"),
("pdf.author", "author"),
("pdf.producer", "producer"),
("pdf.creator", "creator"),
("headers_footers", "headers_footers"),
] {
if let Some(v) = raw.meta.get(src) {
pdf.insert(dst.to_string(), serde_json::Value::String(v.clone()));
}
}
for (src, dst) in [
("pdf.complexity", "complexity"),
("pdf.page_errors", "page_errors"),
] {
if let Some(v) = raw.meta.get(src) {
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(v) {
pdf.insert(dst.to_string(), parsed);
}
}
}
pdf.insert(
"truncated".to_string(),
serde_json::Value::Bool(
raw.meta
.get("pdf.truncated")
.map(|s| s == "true")
.unwrap_or(false),
),
);
pdf.insert(
"layout_modes".to_string(),
serde_json::json!(vec![mode.as_str(); doc.pages.len()]),
);
if let Ok(s) = serde_json::to_string(&pdf) {
doc.meta.extra.insert("pdf".to_string(), s);
}
doc.meta
.extra
.insert("needs_ocr".to_string(), raw.needs_ocr.to_string());
doc
}
#[derive(Debug, Clone, Default)]
pub struct PdfPlugin {
options: PdfOptions,
}
impl PdfPlugin {
pub fn new() -> Self {
Self::default()
}
pub fn with_options(mut self, options: PdfOptions) -> Self {
self.options = options;
self
}
}
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, self.options.layout_mode)
}
}
pub fn register_defaults(reg: &mut Registry) {
register_with(reg, &PdfOptions::default());
}
pub fn register_with(reg: &mut Registry, options: &PdfOptions) {
reg.register_parse(
Arc::new(PdfTextBackend::new().with_options(options.clone())),
None,
)
.unwrap();
if find_library().is_some() {
reg.register_parse(
Arc::new(PdfiumWorkerBackend::new().with_options(options.clone())),
None,
)
.unwrap();
}
reg.register_plugin(Arc::new(PdfPlugin::new().with_options(options.clone())));
}
pub fn is_terminal(e: &DocError) -> bool {
matches!(e.code, ErrorCode::Encrypted | ErrorCode::Unsupported)
}
#[cfg(test)]
mod tests {
use super::*;
fn item(text: &str, x0: f32, y0: f32, x1: f32, y1: f32) -> RawTextItem {
RawTextItem {
text: text.to_string(),
bbox: Some(Rect { x0, y0, x1, y1 }),
..Default::default()
}
}
fn order_stream(items: Vec<RawTextItem>) -> Vec<RawTextItem> {
column_order(items)
}
fn raw_image(data: Vec<u8>, ext: &str, px: Option<(u32, u32)>) -> RawImage {
RawImage {
data,
bbox: None,
ext: ext.to_string(),
px,
ocr_text: None,
}
}
#[test]
fn image_dedup_key_hashes_bytes() {
let im = raw_image(b"abc".to_vec(), "png", Some((60, 40)));
assert_eq!(
image_dedup_key(&im).as_deref(),
Some("900150983cd24fb0d6963f7d28e17f72")
);
let same = raw_image(b"abc".to_vec(), "jpg", Some((10, 10)));
assert_eq!(image_dedup_key(&im), image_dedup_key(&same));
let other = raw_image(b"abd".to_vec(), "png", Some((60, 40)));
assert_ne!(image_dedup_key(&im), image_dedup_key(&other));
}
#[test]
fn image_dedup_key_falls_back_to_size() {
assert_eq!(
image_dedup_key(&raw_image(Vec::new(), "png", Some((120, 80)))).as_deref(),
Some("120x80@png")
);
assert_eq!(
image_dedup_key(&raw_image(Vec::new(), "", Some((120, 80)))).as_deref(),
Some("120x80@")
);
assert_eq!(image_dedup_key(&raw_image(Vec::new(), "png", None)), None);
}
#[test]
fn layout_mode_parses_roundtrip() {
for (s, m) in [
("auto", LayoutMode::Auto),
("REBUILD", LayoutMode::Rebuild),
(" raw ", LayoutMode::Raw),
] {
let parsed = LayoutMode::parse(s).expect("应可解析");
assert_eq!(parsed, m);
assert_eq!(parsed.as_str(), m.as_str());
}
assert!(LayoutMode::parse("fast").is_none());
assert!(LayoutMode::Auto.reflows() && LayoutMode::Rebuild.reflows());
assert!(!LayoutMode::Raw.reflows());
assert_eq!(LayoutMode::default(), LayoutMode::Auto);
}
#[test]
fn pdf_version_maps_to_dotted_string() {
assert_eq!(pdf_version_str(PdfDocumentVersion::Pdf1_7), "1.7");
assert_eq!(pdf_version_str(PdfDocumentVersion::Pdf2_0), "2.0");
assert_eq!(pdf_version_str(PdfDocumentVersion::Other(21)), "2.1");
assert_eq!(pdf_version_str(PdfDocumentVersion::Unset), "");
}
#[test]
fn strip_headers_footers_removes_multiple_repeated_texts() {
let page = || RawPage {
items: vec![
item("Conf Title", 10.0, 5.0, 300.0, 15.0),
item("正文第一行", 10.0, 100.0, 300.0, 112.0),
item("正文第二行", 10.0, 120.0, 300.0, 132.0),
item("123", 550.0, 5.0, 580.0, 15.0),
],
width: Some(595.0),
height: Some(842.0),
..Default::default()
};
let mut pages = vec![page(), page(), page()];
let removed = strip_headers_footers(&mut pages);
assert_eq!(removed.len(), 2, "两条重复页眉都应记录：{removed:?}");
for p in &pages {
let texts: Vec<&str> = p.items.iter().map(|i| i.text.as_str()).collect();
assert_eq!(texts, vec!["正文第一行", "正文第二行"], "只应保留正文行");
}
}
#[test]
fn merge_lines_joins_fragments_on_same_baseline() {
let out = merge_lines(vec![
item("北", 10.0, 100.0, 22.0, 112.0),
item("京", 22.0, 100.5, 34.0, 112.5),
item("市", 34.0, 100.0, 46.0, 112.0),
]);
assert_eq!(out.len(), 1);
assert_eq!(out[0].text, "北京市");
}
#[test]
fn merge_lines_keeps_tightly_spaced_lines_apart() {
let out = merge_lines(vec![
item("第一行", 10.0, 100.0, 60.0, 112.0),
item("第二行", 10.0, 113.0, 60.0, 125.0),
]);
assert_eq!(out.len(), 2, "行距 1pt 也应判为两行：{:?}", out);
}
#[test]
fn merge_lines_splits_on_large_horizontal_gap() {
let out = merge_lines(vec![
item("左栏", 10.0, 100.0, 40.0, 112.0),
item("右栏", 200.0, 100.0, 230.0, 112.0),
]);
assert_eq!(out.len(), 2);
assert_eq!(out[0].text, "左栏");
assert_eq!(out[1].text, "右栏");
}
#[test]
fn column_order_reorders_interleaved_two_columns() {
let mut items = Vec::new();
for i in 0..4 {
let y = 100.0 + i as f32 * 12.0;
items.push(item(&format!("左{i}"), 54.0, y, 293.0, y + 8.0));
items.push(item(&format!("右{i}"), 317.0, y, 556.0, y + 8.0));
}
let out = order_stream(items);
let texts: Vec<&str> = out.iter().map(|it| it.text.as_str()).collect();
assert_eq!(
texts,
vec!["左0", "左1", "左2", "左3", "右0", "右1", "右2", "右3"],
"双栏应读成栏序而非左右逐行交错"
);
}
#[test]
fn column_order_keeps_right_column_with_centered_caption() {
let mut items = Vec::new();
for i in 0..4 {
let y = 100.0 + i as f32 * 12.0;
items.push(item(&format!("左{i}"), 54.0, y, 293.0, y + 8.0));
items.push(item(&format!("右{i}"), 317.0, y, 556.0, y + 8.0));
}
items.push(item("居中图注", 162.0, 400.0, 446.0, 408.0));
let out = order_stream(items);
let texts: Vec<&str> = out.iter().map(|it| it.text.as_str()).collect();
assert_eq!(
texts,
vec![
"左0",
"左1",
"左2",
"左3",
"居中图注",
"右0",
"右1",
"右2",
"右3"
],
"图注按 x0 归首栏，右栏仍须整段在后：{:?}",
texts
);
}
#[test]
fn order_by_leads_threshold_uses_lead_span() {
let mut items = Vec::new();
for i in 0..5 {
let y = 100.0 + i as f32 * 12.0;
items.push(item(&format!("左{i}"), 54.0, y, 140.0, y + 8.0));
items.push(item(&format!("右{i}"), 319.0, y, 560.0, y + 8.0));
}
for i in 0..3 {
let y = 100.0 + i as f32 * 12.0;
items.push(item(&format!("中{i}"), 160.0, y, 210.0, y + 8.0));
}
let out = order_by_leads(items, 506.0);
let texts: Vec<&str> = out.iter().map(|it| it.text.as_str()).collect();
let last_left = texts.iter().rposition(|t| t.starts_with("左")).unwrap();
let first_mid = texts.iter().position(|t| t.starts_with("中")).unwrap();
assert!(
first_mid > last_left,
"行首 150 应自成栏、不能让 54 与 150 合并：{texts:?}"
);
}
#[test]
fn order_by_leads_rejects_centered_title_single_column() {
let mut items = Vec::new();
for i in 0..3 {
let y = 100.0 + i as f32 * 12.0;
items.push(item("TITLE", 200.0, y, 380.0, y + 10.0));
}
for i in 0..6 {
let y = 150.0 + i as f32 * 12.0;
items.push(item(&format!("正文{i}"), 54.0, y, 300.0, y + 8.0));
}
let before: Vec<String> = items.iter().map(|it| it.text.clone()).collect();
let out = order_by_leads(items, 506.0);
let after: Vec<String> = out.iter().map(|it| it.text.clone()).collect();
assert_eq!(after, before, "居中标题 + 正文的单栏页不应被分栏重排");
}
#[test]
fn column_order_keeps_sequential_flow_untouched() {
let mut items = Vec::new();
for i in 0..4 {
let y = 100.0 + i as f32 * 12.0;
items.push(item(&format!("左{i}"), 54.0, y, 293.0, y + 8.0));
}
for i in 0..4 {
let y = 100.0 + i as f32 * 12.0;
items.push(item(&format!("右{i}"), 317.0, y, 556.0, y + 8.0));
}
let before: Vec<String> = items.iter().map(|it| it.text.clone()).collect();
let out = order_stream(items);
let after: Vec<String> = out.iter().map(|it| it.text.clone()).collect();
assert_eq!(after, before, "已是栏序的页不应被重排");
}
#[test]
fn detect_table_groups_rows_with_offset_label_cells() {
let mut lines = Vec::new();
for i in 0..8 {
let y = 100.0 + i as f32 * 10.0;
lines.push(item("(A) ", 100.0, y + 0.5, 140.0, y + 9.5));
lines.push(item("16", 200.0, y, 230.0, y + 8.0));
lines.push(item("32", 300.0, y, 330.0, y + 8.0));
lines.push(item("4.91", 400.0, y, 430.0, y + 8.0));
}
let (rows, used) = detect_table(&lines).expect("应认定为表格");
assert_eq!(rows.len(), 8, "8 个视觉行应保持 8 行：{:?}", rows);
assert_eq!(rows[0].len(), 4);
assert_eq!(used.iter().filter(|u| **u).count(), 32);
}
#[test]
fn merge_lines_keeps_bboxless_items() {
let out = merge_lines(vec![
RawTextItem {
text: "无坐标".into(),
..Default::default()
},
item("有坐标", 10.0, 100.0, 40.0, 112.0),
]);
assert_eq!(out.len(), 2);
assert_eq!(out[0].text, "无坐标");
assert!(out[0].bbox.is_none());
assert_eq!(out[1].text, "有坐标");
}
fn raw_doc_with_two_lines() -> RawDoc {
RawDoc {
pages: vec![RawPage {
items: vec![
item("Hello", 10.0, 100.0, 50.0, 112.0),
item("World", 54.0, 100.0, 94.0, 112.0),
],
width: Some(595.0),
height: Some(842.0),
..Default::default()
}],
..Default::default()
}
}
#[test]
fn assemble_honours_layout_mode() {
let raw = raw_doc_with_two_lines();
let auto = assemble(&raw, LayoutMode::Auto);
assert_eq!(auto.pages[0].blocks.len(), 1, "同行片段应合并");
assert_eq!(auto.pages[0].blocks[0].content, "Hello World");
let raw_mode = assemble(&raw, LayoutMode::Raw);
assert_eq!(raw_mode.pages[0].blocks.len(), 2, "raw 不合并");
assert_eq!(raw_mode.pages[0].blocks[0].content, "Hello");
assert_eq!(raw_mode.pages[0].blocks[1].content, "World");
let pdf: serde_json::Value =
serde_json::from_str(&raw_mode.meta.extra["pdf"]).expect("pdf 元数据应为 JSON");
assert_eq!(pdf["layout_modes"], serde_json::json!(["raw"]));
assert_eq!(pdf["truncated"], serde_json::json!(false));
assert_eq!(raw_mode.meta.extra["needs_ocr"], "false");
}
#[test]
fn detect_table_accepts_short_cells() {
let cols = [
("Region", 54.0, 100.0),
("Share", 200.0, 60.0),
("Note", 320.0, 70.0),
];
let mut lines = Vec::new();
for r in 0..3 {
let y = 100.0 + r as f32 * 12.0;
for (t, x, w) in cols {
lines.push(item(&format!("{t}{r}"), x, y, x + w, y + 8.0));
}
}
let hit = detect_table(&lines).expect("短单元格应认表");
assert_eq!(hit.0.len(), 3, "应识别出 3 行");
}
#[test]
fn detect_table_rejects_prose_rows_with_figure_labels() {
let prose = "Like most interpreters, SpiderMonkey has a foreign function interface that \
allows it to call C builtins and host system functions at runtime";
let stray =
"the interpreter API implementation functions to re-materialize the call stack \
on demand, which costs extra work per exit; ";
let mut lines = Vec::new();
for r in 0..3 {
let y = 100.0 + r as f32 * 20.0;
lines.push(item(prose, 54.0, y, 292.0, y + 8.0));
lines.push(item("?>9@AJ.0A:</C./8-2", 300.0, y, 380.0, y + 6.0));
lines.push(item("7<>;./89-@/#3'4", 400.0, y, 470.0, y + 6.0));
}
for y in [110.0_f32, 130.0] {
lines.push(item(&stray.repeat(4), 54.0, y, 292.0, y + 8.0));
}
assert!(
detect_table(&lines).is_none(),
"正文长句 + 图内标签不应被认成表格"
);
}
#[test]
fn detect_table_rejects_lopsided_row() {
let prose = "inner loop header, but the outer loop header is where the trace begins";
let mut lines = Vec::new();
for (r, (label_a, label_b)) in [("Bytecodes", "Overhead"), ("Record", "Enter")]
.iter()
.enumerate()
{
let y = 100.0 + r as f32 * 20.0;
lines.push(item(prose, 54.0, y, 292.0, y + 8.0));
lines.push(item(label_a, 300.0, y, 350.0, y + 6.0));
lines.push(item(label_b, 400.0, y, 450.0, y + 6.0));
}
assert!(
detect_table(&lines).is_none(),
"正文长句 + 短标签列不应被认成表格"
);
}
#[test]
fn pdf_options_builder() {
let d = PdfOptions::default();
assert_eq!(d.layout_mode, LayoutMode::Auto);
assert_eq!(d.max_pages, None);
let o = PdfOptions::new()
.with_layout_mode(LayoutMode::Raw)
.with_max_pages(5);
assert_eq!(o.layout_mode, LayoutMode::Raw);
assert_eq!(o.max_pages, Some(5));
}
#[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(), PdfOptions::default(), 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}"),
}
}
}
}
}