1#![warn(missing_docs)]
8
9pub mod api;
10pub mod models;
11pub mod output;
12pub mod pdf;
13pub mod pipeline;
14pub mod utils;
15
16pub mod hybrid;
18
19pub mod tagged;
20
21use crate::api::config::ProcessingConfig;
22use crate::models::content::ContentElement;
23use crate::models::document::PdfDocument;
24use crate::pdf::bookmark_extractor::extract_bookmarks;
25use crate::pdf::chunk_parser::extract_page_chunks;
26use crate::pdf::page_info;
27#[cfg(not(target_arch = "wasm32"))]
28use crate::pdf::raster_table_ocr::{
29 recover_dominant_image_text_chunks, recover_page_raster_table_cell_text,
30 recover_raster_table_borders,
31};
32use crate::pipeline::orchestrator::{run_pipeline, PipelineState};
33use crate::pipeline::stages::heading_detector::refine_heading_hierarchy;
34use crate::tagged::struct_tree::build_mcid_map;
35use std::time::Instant;
36
37#[cfg(not(target_arch = "wasm32"))]
49pub fn convert(
50 input_path: &std::path::Path,
51 config: &ProcessingConfig,
52) -> Result<PdfDocument, EdgePdfError> {
53 let timing_enabled = timing_enabled();
54 let total_start = Instant::now();
55
56 let phase_start = Instant::now();
57 let raw_doc = pdf::loader::load_pdf(input_path, config.password.as_deref())?;
58 log_phase_duration(timing_enabled, "load_pdf", phase_start);
59
60 let phase_start = Instant::now();
62 let page_info_list = page_info::extract_page_info(&raw_doc.document);
63 log_phase_duration(timing_enabled, "extract_page_info", phase_start);
64
65 let pages_map = raw_doc.document.get_pages();
67 let page_info_by_number: Vec<Option<&page_info::PageInfo>> =
70 if config.raster_table_ocr_enabled() {
71 let mut index = vec![None; pages_map.len().saturating_add(1)];
72 for info in &page_info_list {
73 if let Some(slot) = index.get_mut(info.page_number as usize) {
74 *slot = Some(info);
75 }
76 }
77 index
78 } else {
79 Vec::new()
80 };
81 let mut page_contents = Vec::with_capacity(pages_map.len());
82
83 let phase_start = Instant::now();
84 for (&page_num, &page_id) in &pages_map {
85 let page_chunks = extract_page_chunks(&raw_doc.document, page_num, page_id)?;
86 let mut recovered_text_chunks = Vec::new();
87 let mut recovered_tables = Vec::new();
88 if config.raster_table_ocr_enabled() {
89 if let Some(Some(page_info)) = page_info_by_number.get(page_num as usize) {
90 recovered_text_chunks = recover_dominant_image_text_chunks(
91 input_path,
92 &page_info.crop_box,
93 page_num,
94 &page_chunks.text_chunks,
95 &page_chunks.image_chunks,
96 );
97 recovered_tables = recover_raster_table_borders(
98 input_path,
99 &page_info.crop_box,
100 page_num,
101 &page_chunks.text_chunks,
102 &page_chunks.image_chunks,
103 );
104 }
105 }
106 let mut elements: Vec<ContentElement> = page_chunks
107 .text_chunks
108 .into_iter()
109 .map(ContentElement::TextChunk)
110 .collect();
111 elements.extend(
112 recovered_text_chunks
113 .into_iter()
114 .map(ContentElement::TextChunk),
115 );
116
117 elements.extend(
118 page_chunks
119 .image_chunks
120 .into_iter()
121 .map(ContentElement::Image),
122 );
123 elements.extend(
124 page_chunks
125 .line_chunks
126 .into_iter()
127 .map(ContentElement::Line),
128 );
129 elements.extend(
130 page_chunks
131 .line_art_chunks
132 .into_iter()
133 .map(ContentElement::LineArt),
134 );
135 elements.extend(
136 recovered_tables
137 .into_iter()
138 .map(ContentElement::TableBorder),
139 );
140
141 page_contents.push(elements);
142 }
143 log_phase_duration(timing_enabled, "extract_page_chunks", phase_start);
144
145 let phase_start = Instant::now();
147 let mcid_map = build_mcid_map(&raw_doc.document);
148 let mut pipeline_state = PipelineState::with_mcid_map(page_contents, config.clone(), mcid_map)
149 .with_page_info(page_info_list);
150 run_pipeline(&mut pipeline_state)?;
151 let bookmarks = extract_bookmarks(&raw_doc.document);
153 if !bookmarks.is_empty() {
154 refine_heading_hierarchy(&mut pipeline_state.pages, &bookmarks);
155 } else {
156 refine_heading_hierarchy(&mut pipeline_state.pages, &[]);
157 }
158 log_phase_duration(timing_enabled, "run_pipeline", phase_start);
159
160 let file_name = input_path
162 .file_name()
163 .and_then(|n| n.to_str())
164 .unwrap_or("unknown.pdf")
165 .to_string();
166
167 let mut doc = PdfDocument::new(file_name);
168 doc.source_path = Some(input_path.display().to_string());
169 doc.number_of_pages = pages_map.len() as u32;
170 doc.author = raw_doc.metadata.author;
171 doc.title = raw_doc.metadata.title;
172 doc.creation_date = raw_doc.metadata.creation_date;
173 doc.modification_date = raw_doc.metadata.modification_date;
174
175 let phase_start = Instant::now();
176 if config.raster_table_ocr_enabled() {
177 for (page_idx, page) in pipeline_state.pages.iter_mut().enumerate() {
178 if let Some(page_info) = pipeline_state.page_info.get(page_idx) {
179 recover_page_raster_table_cell_text(
180 input_path,
181 &page_info.crop_box,
182 page_info.page_number,
183 page,
184 );
185 }
186 }
187 }
188 log_phase_duration(
189 timing_enabled,
190 "recover_page_raster_table_cell_text",
191 phase_start,
192 );
193
194 let phase_start = Instant::now();
196 for page in pipeline_state.pages {
197 doc.kids.extend(page);
198 }
199 log_phase_duration(timing_enabled, "flatten_document", phase_start);
200 log_phase_duration(timing_enabled, "convert_total", total_start);
201
202 Ok(doc)
203}
204
205pub fn convert_bytes(
226 data: &[u8],
227 file_name: &str,
228 config: &ProcessingConfig,
229) -> Result<PdfDocument, EdgePdfError> {
230 let session = extract_session(data, file_name, config, None)?;
231 assemble(session, OcrAssembleMode::SyncEngine)
232}
233
234pub type ProgressFn<'a> = dyn FnMut(&str, u32, u32) + 'a;
236
237pub struct ExtractSession {
239 pub file_name: String,
241 pub config: ProcessingConfig,
243 pub author: Option<String>,
245 pub title: Option<String>,
247 pub creation_date: Option<String>,
249 pub modification_date: Option<String>,
251 pub number_of_pages: u32,
253 pub page_contents: Vec<Vec<ContentElement>>,
255 pub page_info_list: Vec<page_info::PageInfo>,
257 pub mcid_map: crate::tagged::struct_tree::McidMap,
259 pub bookmarks: Vec<pdf::bookmark_extractor::Bookmark>,
261 #[cfg(feature = "image")]
263 pub candidates: Vec<pdf::inmem_raster::RasterCandidate>,
264 #[cfg(feature = "image")]
266 pub ocr_words: std::collections::HashMap<u32, Vec<pdf::ocr::OcrWord>>,
267}
268
269pub enum OcrAssembleMode {
271 SyncEngine,
273 Provided,
275 #[cfg(feature = "image")]
277 Engine(std::sync::Arc<dyn pdf::ocr::OcrEngine>),
278}
279
280pub fn extract_session(
284 data: &[u8],
285 file_name: &str,
286 config: &ProcessingConfig,
287 mut on_progress: Option<&mut ProgressFn<'_>>,
288) -> Result<ExtractSession, EdgePdfError> {
289 let raw_doc = pdf::loader::load_pdf_from_bytes(data, config.password.as_deref())?;
290 let page_info_list = page_info::extract_page_info(&raw_doc.document);
291 let pages_map = raw_doc.document.get_pages();
292 let page_total = pages_map.len() as u32;
293 let mut page_contents = Vec::with_capacity(pages_map.len());
294 #[cfg(feature = "image")]
295 let mut candidates = Vec::new();
296 #[cfg(feature = "image")]
297 let mut next_id = 0u32;
298
299 let mut page_done = 0u32;
300 for (&page_num, &page_id) in &pages_map {
301 let page_chunks = extract_page_chunks(&raw_doc.document, page_num, page_id)?;
302
303 #[cfg(feature = "image")]
304 if config.raster_table_ocr_enabled() {
305 let page_cands = pdf::inmem_raster::collect_raster_candidates(
306 &raw_doc.document,
307 page_id,
308 &page_chunks.image_chunks,
309 &page_chunks.text_chunks,
310 &mut next_id,
311 );
312 candidates.extend(page_cands);
313 }
314
315 let mut elements: Vec<ContentElement> = page_chunks
316 .text_chunks
317 .into_iter()
318 .map(ContentElement::TextChunk)
319 .collect();
320 elements.extend(
321 page_chunks
322 .image_chunks
323 .into_iter()
324 .map(ContentElement::Image),
325 );
326 elements.extend(
327 page_chunks
328 .line_chunks
329 .into_iter()
330 .map(ContentElement::Line),
331 );
332 elements.extend(
333 page_chunks
334 .line_art_chunks
335 .into_iter()
336 .map(ContentElement::LineArt),
337 );
338 page_contents.push(elements);
340
341 page_done += 1;
342 if let Some(cb) = on_progress.as_mut() {
343 cb("planning", page_done, page_total);
344 }
345 }
346
347 let mcid_map = build_mcid_map(&raw_doc.document);
348 let bookmarks = extract_bookmarks(&raw_doc.document);
349
350 Ok(ExtractSession {
351 file_name: file_name.to_string(),
352 config: config.clone(),
353 author: raw_doc.metadata.author,
354 title: raw_doc.metadata.title,
355 creation_date: raw_doc.metadata.creation_date,
356 modification_date: raw_doc.metadata.modification_date,
357 number_of_pages: pages_map.len() as u32,
358 page_contents,
359 page_info_list,
360 mcid_map,
361 bookmarks,
362 #[cfg(feature = "image")]
363 candidates,
364 #[cfg(feature = "image")]
365 ocr_words: std::collections::HashMap::new(),
366 })
367}
368
369impl ExtractSession {
370 #[cfg(feature = "image")]
372 pub fn provide_ocr(&mut self, id: u32, words: Vec<pdf::ocr::OcrWord>) {
373 self.ocr_words.insert(id, words);
374 }
375
376 #[cfg(feature = "image")]
378 pub fn candidate_metas(&self) -> Vec<pdf::inmem_raster::RasterCandidateMeta> {
379 self.candidates.iter().map(|c| c.meta()).collect()
380 }
381
382 #[cfg(feature = "image")]
384 pub fn candidate_gray(&self, id: u32) -> Option<&[u8]> {
385 self.candidates
386 .iter()
387 .find(|c| c.id == id)
388 .map(|c| c.gray.as_slice())
389 }
390}
391
392pub fn assemble(
394 mut session: ExtractSession,
395 mode: OcrAssembleMode,
396) -> Result<PdfDocument, EdgePdfError> {
397 #[cfg(feature = "image")]
398 if session.config.raster_table_ocr_enabled() {
399 use image::GrayImage;
400 use pdf::ocr::{default_engine, OcrEngine};
401
402 for cand in &session.candidates {
403 let words = match &mode {
404 OcrAssembleMode::SyncEngine => {
405 if let Some(gray) =
406 GrayImage::from_raw(cand.width, cand.height, cand.gray.clone())
407 {
408 OcrEngine::recognize(&*default_engine(), &gray)
409 } else {
410 Vec::new()
411 }
412 }
413 OcrAssembleMode::Provided => {
414 session.ocr_words.get(&cand.id).cloned().unwrap_or_default()
415 }
416 OcrAssembleMode::Engine(engine) => {
417 if let Some(gray) =
418 GrayImage::from_raw(cand.width, cand.height, cand.gray.clone())
419 {
420 engine.recognize(&gray)
421 } else {
422 Vec::new()
423 }
424 }
425 };
426 if let Some(table) = pdf::inmem_raster::build_table_from_candidate(cand, &words) {
427 if pdf::inmem_raster::bordered_table_is_plausible(&table) {
428 let page_idx = cand.page.saturating_sub(1) as usize;
429 if let Some(page) = session.page_contents.get_mut(page_idx) {
430 page.push(ContentElement::TableBorder(table));
431 }
432 }
433 }
434 }
435 }
436
437 #[cfg(not(feature = "image"))]
438 let _ = mode;
439
440 let mut pipeline_state = PipelineState::with_mcid_map(
441 session.page_contents,
442 session.config.clone(),
443 session.mcid_map,
444 )
445 .with_page_info(session.page_info_list);
446 run_pipeline(&mut pipeline_state)?;
447 refine_heading_hierarchy(&mut pipeline_state.pages, &session.bookmarks);
448
449 let mut doc = PdfDocument::new(session.file_name);
450 doc.number_of_pages = session.number_of_pages;
451 doc.author = session.author;
452 doc.title = session.title;
453 doc.creation_date = session.creation_date;
454 doc.modification_date = session.modification_date;
455
456 for page in pipeline_state.pages {
457 doc.kids.extend(page);
458 }
459
460 Ok(doc)
461}
462
463#[derive(Debug, thiserror::Error)]
465pub enum EdgePdfError {
466 #[error("PDF loading error: {0}")]
468 LoadError(String),
469
470 #[error("Pipeline error at stage {stage}: {message}")]
472 PipelineError {
473 stage: u32,
475 message: String,
477 },
478
479 #[error("Output error: {0}")]
481 OutputError(String),
482
483 #[error("I/O error: {0}")]
485 IoError(#[from] std::io::Error),
486
487 #[error("Configuration error: {0}")]
489 ConfigError(String),
490
491 #[error("PDF parse error: {0}")]
493 LopdfError(String),
494}
495
496impl From<lopdf::Error> for EdgePdfError {
497 fn from(e: lopdf::Error) -> Self {
498 EdgePdfError::LopdfError(e.to_string())
499 }
500}
501
502fn timing_enabled() -> bool {
503 std::env::var("EDGEPARSE_TIMING")
504 .map(|value| {
505 matches!(
506 value.to_ascii_lowercase().as_str(),
507 "1" | "true" | "yes" | "on"
508 )
509 })
510 .unwrap_or(false)
511}
512
513fn log_phase_duration(enabled: bool, phase: &str, start: Instant) {
514 if enabled {
515 log::info!(
516 "Timing {}: {:.2} ms",
517 phase,
518 start.elapsed().as_secs_f64() * 1000.0
519 );
520 }
521}
522
523#[cfg(test)]
524mod tests {
525 use super::*;
526 use lopdf::{
527 content::{Content, Operation},
528 dictionary, Object, Stream,
529 };
530 use std::io::Write;
531
532 fn create_test_pdf_file(path: &std::path::Path) {
534 let mut doc = lopdf::Document::with_version("1.5");
535 let pages_id = doc.new_object_id();
536
537 let font_id = doc.add_object(dictionary! {
538 "Type" => "Font",
539 "Subtype" => "Type1",
540 "BaseFont" => "Helvetica",
541 });
542
543 let resources_id = doc.add_object(dictionary! {
544 "Font" => dictionary! {
545 "F1" => font_id,
546 },
547 });
548
549 let content = Content {
550 operations: vec![
551 Operation::new("BT", vec![]),
552 Operation::new("Tf", vec!["F1".into(), 12.into()]),
553 Operation::new("Td", vec![72.into(), 700.into()]),
554 Operation::new("Tj", vec![Object::string_literal("Hello EdgeParse!")]),
555 Operation::new("Td", vec![0.into(), Object::Real(-20.0)]),
556 Operation::new("Tj", vec![Object::string_literal("Second line of text.")]),
557 Operation::new("ET", vec![]),
558 ],
559 };
560
561 let encoded = content.encode().unwrap();
562 let content_id = doc.add_object(Stream::new(dictionary! {}, encoded));
563
564 let page_id = doc.add_object(dictionary! {
565 "Type" => "Page",
566 "Parent" => pages_id,
567 "Contents" => content_id,
568 "Resources" => resources_id,
569 "MediaBox" => vec![0.into(), 0.into(), 595.into(), 842.into()],
570 });
571
572 let pages = dictionary! {
573 "Type" => "Pages",
574 "Kids" => vec![page_id.into()],
575 "Count" => 1,
576 };
577 doc.objects.insert(pages_id, Object::Dictionary(pages));
578
579 let catalog_id = doc.add_object(dictionary! {
580 "Type" => "Catalog",
581 "Pages" => pages_id,
582 });
583 doc.trailer.set("Root", catalog_id);
584
585 let mut file = std::fs::File::create(path).unwrap();
586 doc.save_to(&mut file).unwrap();
587 file.flush().unwrap();
588 }
589
590 #[test]
591 fn test_convert_end_to_end() {
592 let dir = std::env::temp_dir().join("edgeparse_test");
593 std::fs::create_dir_all(&dir).unwrap();
594 let pdf_path = dir.join("test_convert.pdf");
595
596 create_test_pdf_file(&pdf_path);
597
598 let config = ProcessingConfig::default();
599 let result = convert(&pdf_path, &config);
600 assert!(result.is_ok(), "convert() failed: {:?}", result.err());
601
602 let doc = result.unwrap();
603 assert_eq!(doc.number_of_pages, 1);
604 assert!(
605 !doc.kids.is_empty(),
606 "Expected content elements in document"
607 );
608
609 let mut all_text = String::new();
611 for element in &doc.kids {
612 match element {
613 models::content::ContentElement::TextChunk(tc) => {
614 all_text.push_str(&tc.value);
615 all_text.push(' ');
616 }
617 models::content::ContentElement::TextLine(tl) => {
618 all_text.push_str(&tl.value());
619 all_text.push(' ');
620 }
621 models::content::ContentElement::TextBlock(tb) => {
622 all_text.push_str(&tb.value());
623 all_text.push(' ');
624 }
625 models::content::ContentElement::Paragraph(p) => {
626 all_text.push_str(&p.base.value());
627 all_text.push(' ');
628 }
629 models::content::ContentElement::Heading(h) => {
630 all_text.push_str(&h.base.base.value());
631 all_text.push(' ');
632 }
633 _ => {}
634 }
635 }
636
637 assert!(
638 all_text.contains("Hello"),
639 "Expected 'Hello' in extracted text, got: {}",
640 all_text
641 );
642 assert!(
643 all_text.contains("Second"),
644 "Expected 'Second' in extracted text, got: {}",
645 all_text
646 );
647
648 let _ = std::fs::remove_file(&pdf_path);
650 }
651
652 #[cfg(feature = "image")]
653 struct FakeOcrEngine;
654
655 #[cfg(feature = "image")]
656 impl pdf::ocr::OcrEngine for FakeOcrEngine {
657 fn recognize(&self, _gray: &image::GrayImage) -> Vec<pdf::ocr::OcrWord> {
658 vec![
659 pdf::ocr::OcrWord {
660 line_key: (1, 1, 1),
661 left: 10,
662 top: 10,
663 width: 40,
664 height: 12,
665 text: "A".into(),
666 confidence: 95.0,
667 },
668 pdf::ocr::OcrWord {
669 line_key: (1, 1, 1),
670 left: 80,
671 top: 10,
672 width: 40,
673 height: 12,
674 text: "B".into(),
675 confidence: 95.0,
676 },
677 pdf::ocr::OcrWord {
678 line_key: (1, 1, 2),
679 left: 10,
680 top: 40,
681 width: 40,
682 height: 12,
683 text: "C".into(),
684 confidence: 95.0,
685 },
686 pdf::ocr::OcrWord {
687 line_key: (1, 1, 2),
688 left: 80,
689 top: 40,
690 width: 40,
691 height: 12,
692 text: "D".into(),
693 confidence: 95.0,
694 },
695 ]
696 }
697 }
698
699 #[cfg(feature = "image")]
700 fn count_table_borders(doc: &PdfDocument) -> usize {
701 doc.kids
702 .iter()
703 .filter(|e| matches!(e, models::content::ContentElement::TableBorder(_)))
704 .count()
705 }
706
707 #[cfg(feature = "image")]
709 #[test]
710 fn two_phase_equals_sync_fake_engine() {
711 let dir = std::env::temp_dir().join("edgeparse_two_phase");
712 std::fs::create_dir_all(&dir).unwrap();
713 let pdf_path = dir.join("hello.pdf");
714 create_test_pdf_file(&pdf_path);
715 let data = std::fs::read(&pdf_path).unwrap();
716 let config = ProcessingConfig::default();
717 let engine: std::sync::Arc<dyn pdf::ocr::OcrEngine> = std::sync::Arc::new(FakeOcrEngine);
718
719 let sync_doc = {
720 let session = extract_session(&data, "hello.pdf", &config, None).unwrap();
721 assemble(session, OcrAssembleMode::Engine(engine.clone())).unwrap()
722 };
723
724 let two_phase_doc = {
725 let mut session = extract_session(&data, "hello.pdf", &config, None).unwrap();
726 let ids: Vec<_> = session.candidates.iter().map(|c| c.id).collect();
727 for id in ids {
728 let gray = session.candidate_gray(id).unwrap().to_vec();
729 let meta = session.candidates.iter().find(|c| c.id == id).unwrap();
730 let img = image::GrayImage::from_raw(meta.width, meta.height, gray).unwrap();
731 let words = engine.recognize(&img);
732 session.provide_ocr(id, words);
733 }
734 assemble(session, OcrAssembleMode::Provided).unwrap()
735 };
736
737 assert_eq!(sync_doc.number_of_pages, two_phase_doc.number_of_pages);
738 assert_eq!(
739 count_table_borders(&sync_doc),
740 count_table_borders(&two_phase_doc)
741 );
742 let sync_md = output::markdown::to_markdown(&sync_doc).unwrap();
743 let phase_md = output::markdown::to_markdown(&two_phase_doc).unwrap();
744 assert_eq!(sync_md, phase_md);
745 let _ = std::fs::remove_file(&pdf_path);
746 }
747
748 #[cfg(feature = "image")]
750 #[test]
751 fn candidate_prefilter_keeps_docs_110_and_122() {
752 let roots = [
753 std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../benchmark/pdfs"),
754 std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../odl-bench/pdfs"),
755 ];
756 for id in ["01030000000110", "01030000000122"] {
757 let pdf = roots
758 .iter()
759 .map(|r| r.join(format!("{id}.pdf")))
760 .find(|p| p.exists());
761 let Some(pdf) = pdf else {
762 eprintln!("skip {id}: PDF not found");
763 continue;
764 };
765 let data = std::fs::read(&pdf).unwrap();
766 let config = ProcessingConfig::default();
767 let session = extract_session(&data, &format!("{id}.pdf"), &config, None).unwrap();
768 assert!(
769 !session.candidates.is_empty(),
770 "{id}: expected ≥1 OCR candidate, got 0 (pre-filter dropped tables)"
771 );
772 let doc = assemble(session, OcrAssembleMode::SyncEngine).unwrap();
774 let _ = count_table_borders(&doc);
777 }
778 }
779}