1use super::text_region::TextRegion;
7use crate::processors::BoundingBox;
8use image::RgbImage;
9use regex::Regex;
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::path::Path;
13use std::sync::Arc;
14use std::sync::LazyLock;
15
16static TITLE_NUMBERING_REGEX: LazyLock<Regex> = LazyLock::new(|| {
19 Regex::new(
20 r"(?x)
21 ^\s*
22 (
23 # Arabic numerals: 1, 1.2, 1.2.3, etc.
24 [1-9][0-9]*(?:\.[1-9][0-9]*)*[\.、]?
25 |
26 # Parenthesized Arabic numerals: (1), (1.2), etc.
27 [((][1-9][0-9]*(?:\.[1-9][0-9]*)*[))]
28 |
29 # Chinese numerals with punctuation: 一、 二、
30 [一二三四五六七八九十百千万亿零壹贰叁肆伍陆柒捌玖拾][、.]?
31 |
32 # Parenthesized Chinese numerals: (一)
33 [((][一二三四五六七八九十百千万亿零壹贰叁肆伍陆柒捌玖拾]+[))]
34 |
35 # Roman numerals with delimiter (period or followed by space)
36 (?:I|II|III|IV|V|VI|VII|VIII|IX|X)(?:\.|\b)
37 )
38 (\s+)
39 (.*)
40 $
41 ",
42 )
43 .unwrap_or_else(|e| panic!("Invalid title numbering regex: {e}"))
44});
45
46fn semantic_title_level_and_format(cleaned: &str) -> Option<(usize, String)> {
63 let trimmed = cleaned.trim();
64
65 let keyword = trimmed.trim_end_matches(':').to_ascii_uppercase();
67 if matches!(
68 keyword.as_str(),
69 "ABSTRACT" | "INTRODUCTION" | "REFERENCES" | "REFERENCE"
70 ) {
71 return Some((2, trimmed.to_string()));
72 }
73
74 if let Some(captures) = TITLE_NUMBERING_REGEX.captures(cleaned) {
75 let numbering = captures.get(1).map(|m| m.as_str().trim()).unwrap_or("");
76 let title_content = captures.get(3).map(|m| m.as_str()).unwrap_or("");
77
78 let dot_count = numbering.matches('.').count();
79 let level = (dot_count + 2).clamp(2, 6);
80
81 let formatted = if title_content.is_empty() {
82 numbering.trim_end_matches('.').to_string()
83 } else {
84 format!(
85 "{} {}",
86 numbering.trim_end_matches('.'),
87 title_content.trim_start()
88 )
89 };
90 return Some((level, formatted));
91 }
92
93 None
94}
95
96fn semantic_title_level(text: &str) -> Option<usize> {
97 let cleaned = text.replace("-\n", "").replace('\n', " ");
98 semantic_title_level_and_format(&cleaned).map(|(level, _)| level)
99}
100
101fn format_title_with_level(title: &str, clustered_level: Option<usize>) -> (usize, String) {
102 let cleaned = title.replace("-\n", "").replace('\n', " ");
104 if let Some((level, formatted)) = semantic_title_level_and_format(&cleaned) {
105 return (level, formatted);
106 }
107
108 let level = clustered_level.unwrap_or(2).clamp(2, 6);
110 (level, cleaned)
111}
112
113fn infer_paragraph_title_levels(elements: &[LayoutElement]) -> HashMap<usize, usize> {
119 let title_indices: Vec<usize> = elements
120 .iter()
121 .enumerate()
122 .filter(|(_, e)| e.element_type == LayoutElementType::ParagraphTitle)
123 .map(|(idx, _)| idx)
124 .collect();
125 if title_indices.is_empty() {
126 return HashMap::new();
127 }
128
129 let height_samples: Vec<(usize, f32)> = title_indices
130 .iter()
131 .filter_map(|&idx| {
132 let e = &elements[idx];
133 let height = (e.bbox.y_max() - e.bbox.y_min()).max(1.0);
134 let line_h = height / e.num_lines.unwrap_or(1).max(1) as f32;
135 let v = line_h.max(1.0);
136 if v.is_finite() { Some((idx, v)) } else { None }
137 })
138 .collect();
139
140 let indent_samples: Vec<(usize, f32)> = title_indices
141 .iter()
142 .filter_map(|&idx| {
143 let x = elements[idx].bbox.x_min();
144 if x.is_finite() { Some((idx, x)) } else { None }
145 })
146 .collect();
147 let semantic_levels: HashMap<usize, usize> = title_indices
148 .iter()
149 .filter_map(|&idx| {
150 elements[idx]
151 .text
152 .as_deref()
153 .and_then(semantic_title_level)
154 .map(|level| (idx, level))
155 })
156 .collect();
157
158 let font_levels = infer_levels_by_kmeans_feature(&height_samples, true);
159 let relative_levels = infer_levels_by_kmeans_feature(&indent_samples, false);
161
162 let mut voted = HashMap::new();
163 for idx in title_indices {
164 let semantic_level = semantic_levels.get(&idx).copied();
165 let font_level = font_levels.get(&idx).copied();
166 let relative_level = relative_levels.get(&idx).copied();
167
168 let mut score = [0u8; 7];
169 if let Some(level) = semantic_level {
170 score[level.clamp(1, 6)] += 2;
171 }
172 if let Some(level) = font_level {
173 score[level.clamp(1, 6)] += 1;
174 }
175 if let Some(level) = relative_level {
176 score[level.clamp(1, 6)] += 1;
177 }
178
179 let mut best_level = semantic_level.unwrap_or(2);
180 let mut best_score = 0u8;
181 for (level, &s) in score.iter().enumerate().skip(1) {
182 if s > best_score {
183 best_score = s;
184 best_level = level;
185 } else if s == best_score && s > 0 {
186 let is_semantic = semantic_level == Some(level);
187 let best_is_semantic = semantic_level == Some(best_level);
188 if (is_semantic && !best_is_semantic)
189 || (is_semantic == best_is_semantic && level < best_level)
190 {
191 best_level = level;
192 }
193 }
194 }
195
196 if best_score == 0 {
197 best_level = semantic_level
198 .or(font_level)
199 .or(relative_level)
200 .unwrap_or(2);
201 }
202
203 voted.insert(idx, best_level.clamp(1, 6));
204 }
205
206 voted
207}
208
209fn infer_levels_by_kmeans_feature(
214 samples: &[(usize, f32)],
215 descending: bool,
216) -> HashMap<usize, usize> {
217 let clean_samples: Vec<(usize, f32)> = samples
218 .iter()
219 .copied()
220 .filter(|(_, v)| v.is_finite())
221 .collect();
222 if clean_samples.len() < 2 {
223 return HashMap::new();
224 }
225
226 let mut values: Vec<f32> = clean_samples.iter().map(|(_, v)| *v).collect();
227 values.sort_by(|a, b| a.total_cmp(b));
228 let unique_count = values
229 .windows(2)
230 .filter(|w| (w[1] - w[0]).abs() > 1e-3)
231 .count()
232 + 1;
233 let k = unique_count.clamp(1, 4).min(clean_samples.len());
234 if k <= 1 {
235 return HashMap::new();
236 }
237
238 let mut centroids = (0..k)
239 .map(|i| {
240 let pos = ((i as f32 + 0.5) / k as f32 * values.len() as f32).floor() as usize;
241 values[pos.min(values.len() - 1)]
242 })
243 .collect::<Vec<_>>();
244
245 for _ in 0..16 {
246 let mut sums = vec![0.0f32; k];
247 let mut counts = vec![0usize; k];
248 for (_, value) in &clean_samples {
249 let mut best_idx = 0usize;
250 let mut best_dist = f32::INFINITY;
251 for (idx, c) in centroids.iter().enumerate() {
252 let dist = (value - c).abs();
253 if dist < best_dist {
254 best_dist = dist;
255 best_idx = idx;
256 }
257 }
258 sums[best_idx] += *value;
259 counts[best_idx] += 1;
260 }
261 for idx in 0..k {
262 if counts[idx] > 0 {
263 centroids[idx] = sums[idx] / counts[idx] as f32;
264 }
265 }
266 }
267
268 let mut centroid_order: Vec<(usize, f32)> = centroids.iter().copied().enumerate().collect();
269 if descending {
270 centroid_order.sort_by(|a, b| b.1.total_cmp(&a.1));
271 } else {
272 centroid_order.sort_by(|a, b| a.1.total_cmp(&b.1));
273 }
274 let rank_by_cluster: HashMap<usize, usize> = centroid_order
275 .into_iter()
276 .enumerate()
277 .map(|(rank, (cluster_idx, _))| (cluster_idx, rank))
278 .collect();
279
280 let mut result = HashMap::new();
281 for (element_idx, value) in &clean_samples {
282 let mut best_idx = 0usize;
283 let mut best_dist = f32::INFINITY;
284 for (idx, c) in centroids.iter().enumerate() {
285 let dist = (value - c).abs();
286 if dist < best_dist {
287 best_dist = dist;
288 best_idx = idx;
289 }
290 }
291 let rank = rank_by_cluster.get(&best_idx).copied().unwrap_or(0);
292 let level = (rank + 2).clamp(2, 6);
293 result.insert(*element_idx, level);
294 }
295
296 result
297}
298
299#[derive(Debug, Clone, Serialize, Deserialize)]
311pub struct RegionBlock {
312 pub bbox: BoundingBox,
314 pub confidence: f32,
316 pub order_index: Option<u32>,
318 pub element_indices: Vec<usize>,
320}
321
322#[derive(Debug, Clone, Serialize, Deserialize)]
331pub struct PageContinuationFlags {
332 pub paragraph_start: bool,
334 pub paragraph_end: bool,
336}
337
338impl PageContinuationFlags {
339 pub fn new(paragraph_start: bool, paragraph_end: bool) -> Self {
340 Self {
341 paragraph_start,
342 paragraph_end,
343 }
344 }
345
346 pub fn as_tuple(&self) -> (bool, bool) {
348 (self.paragraph_start, self.paragraph_end)
349 }
350}
351
352#[derive(Debug, Clone, Serialize, Deserialize)]
374pub struct StructureResult {
375 pub input_path: Arc<str>,
377 pub index: usize,
379 pub layout_elements: Vec<LayoutElement>,
381 pub tables: Vec<TableResult>,
383 pub formulas: Vec<FormulaResult>,
385 pub text_regions: Option<Vec<TextRegion>>,
387 pub orientation_angle: Option<f32>,
389 pub region_blocks: Option<Vec<RegionBlock>>,
392 #[serde(skip)]
397 pub rectified_img: Option<Arc<RgbImage>>,
398 pub page_continuation_flags: Option<PageContinuationFlags>,
402}
403
404impl StructureResult {
405 pub fn new(input_path: impl Into<Arc<str>>, index: usize) -> Self {
407 Self {
408 input_path: input_path.into(),
409 index,
410 layout_elements: Vec::new(),
411 tables: Vec::new(),
412 formulas: Vec::new(),
413 text_regions: None,
414 orientation_angle: None,
415 region_blocks: None,
416 rectified_img: None,
417 page_continuation_flags: None,
418 }
419 }
420
421 pub fn with_layout_elements(mut self, elements: Vec<LayoutElement>) -> Self {
423 self.layout_elements = elements;
424 self
425 }
426
427 pub fn with_tables(mut self, tables: Vec<TableResult>) -> Self {
429 self.tables = tables;
430 self
431 }
432
433 pub fn with_formulas(mut self, formulas: Vec<FormulaResult>) -> Self {
435 self.formulas = formulas;
436 self
437 }
438
439 pub fn with_text_regions(mut self, regions: Vec<TextRegion>) -> Self {
441 self.text_regions = Some(regions);
442 self
443 }
444
445 pub fn with_region_blocks(mut self, blocks: Vec<RegionBlock>) -> Self {
450 self.region_blocks = Some(blocks);
451 self
452 }
453
454 pub fn with_page_continuation_flags(mut self, flags: PageContinuationFlags) -> Self {
456 self.page_continuation_flags = Some(flags);
457 self
458 }
459
460 pub fn to_markdown(&self) -> String {
472 let table_bboxes: Vec<&BoundingBox> = self
474 .layout_elements
475 .iter()
476 .filter(|e| e.element_type == LayoutElementType::Table)
477 .map(|e| &e.bbox)
478 .collect();
479
480 let original_image_width = self
482 .rectified_img
483 .as_ref()
484 .map(|img| img.width() as f32)
485 .or_else(|| {
486 self.layout_elements
488 .iter()
489 .map(|e| e.bbox.x_max())
490 .fold(None, |acc, x| Some(acc.map_or(x, |max: f32| max.max(x))))
491 })
492 .unwrap_or(1.0);
493
494 let mut md = String::new();
495 let elements = &self.layout_elements;
496 let paragraph_title_levels = infer_paragraph_title_levels(elements);
497 let mut prev_text_element: Option<&LayoutElement> = None;
500
501 for (idx, element) in elements.iter().enumerate() {
502 if matches!(
504 element.element_type,
505 LayoutElementType::Number
506 | LayoutElementType::Footnote
507 | LayoutElementType::Header
508 | LayoutElementType::HeaderImage
509 | LayoutElementType::Footer
510 | LayoutElementType::FooterImage
511 | LayoutElementType::AsideText
512 ) {
513 continue;
514 }
515
516 if element.element_type == LayoutElementType::Text {
520 let overlaps_table = table_bboxes.iter().any(|table_bbox| {
521 element.bbox.ioa(table_bbox) > 0.3 });
523
524 if overlaps_table && element.confidence < 0.7 {
527 continue;
528 }
529 }
530
531 let seg_start_flag = get_seg_flag(element, prev_text_element);
535
536 let is_continuation = element.element_type == LayoutElementType::Text
537 && prev_text_element.is_some()
538 && !seg_start_flag;
539
540 if !is_continuation {
542 }
544
545 match element.element_type {
546 LayoutElementType::DocTitle => {
548 if !md.is_empty() {
549 md.push_str("\n\n");
550 }
551 if let Some(text) = &element.text {
552 let cleaned = clean_ocr_text(text);
553 let keyword = cleaned.trim().trim_end_matches(':').to_ascii_uppercase();
555 if matches!(
556 keyword.as_str(),
557 "ABSTRACT" | "INTRODUCTION" | "REFERENCES" | "REFERENCE"
558 ) {
559 md.push_str("## ");
560 } else {
561 md.push_str("# ");
562 }
563 md.push_str(&cleaned);
564 }
565 }
566 LayoutElementType::ParagraphTitle => {
568 if !md.is_empty() {
569 md.push_str("\n\n");
570 }
571 if let Some(text) = &element.text {
572 let cleaned = clean_ocr_text(text);
573 let clustered = paragraph_title_levels.get(&idx).copied();
574 let (level, formatted_title) = format_title_with_level(&cleaned, clustered);
575 for _ in 0..level {
576 md.push('#');
577 }
578 md.push(' ');
579 md.push_str(&formatted_title);
580 } else {
581 md.push_str("## ");
582 }
583 }
584 LayoutElementType::Table => {
587 if !md.is_empty() {
588 md.push_str("\n\n");
589 }
590 if let Some(table) =
591 self.tables.iter().find(|t| t.bbox.iou(&element.bbox) > 0.5)
592 {
593 if let Some(html) = &table.html_structure {
594 let simplified = simplify_table_html(html);
596 let table_with_border =
597 simplified.replacen("<table>", "<table border=\"1\">", 1);
598 let cleaned = clean_ocr_text(&table_with_border);
600 md.push_str("<div style=\"text-align: center;\">");
601 md.push_str(&cleaned);
602 md.push_str("</div>");
603 } else {
604 md.push_str("[Table]");
605 }
606 } else {
607 md.push_str("[Table]");
608 }
609 }
610 LayoutElementType::FormulaNumber => {
614 continue;
615 }
616 LayoutElementType::Formula => {
618 let raw_content = element.text.as_deref().map(|s| s.trim()).unwrap_or("");
620 if raw_content.is_empty() {
621 continue;
622 }
623 let latex_content =
625 if raw_content.starts_with("$$") && raw_content.ends_with("$$") {
626 &raw_content[2..raw_content.len() - 2]
627 } else if raw_content.starts_with('$') && raw_content.ends_with('$') {
628 &raw_content[1..raw_content.len() - 1]
629 } else {
630 raw_content
631 };
632
633 let is_inline = {
640 let has_prev_text = (0..idx)
641 .rev()
642 .find(|&i| {
643 let t = elements[i].element_type;
644 !t.is_formula() && t != LayoutElementType::FormulaNumber
645 })
646 .is_some_and(|i| {
647 let prev = &elements[i];
648 (prev.element_type == LayoutElementType::Text
649 || prev.element_type == LayoutElementType::ReferenceContent)
650 && is_same_line(&element.bbox, &prev.bbox)
651 });
652
653 let has_next_text = ((idx + 1)..elements.len())
654 .find(|&i| {
655 let t = elements[i].element_type;
656 !t.is_formula() && t != LayoutElementType::FormulaNumber
657 })
658 .is_some_and(|i| {
659 let next = &elements[i];
660 (next.element_type == LayoutElementType::Text
661 || next.element_type == LayoutElementType::ReferenceContent)
662 && is_same_line(&element.bbox, &next.bbox)
663 });
664
665 has_prev_text && has_next_text
668 };
669
670 if is_inline {
671 md.push('$');
673 md.push_str(latex_content);
674 md.push_str("$ ");
675 } else {
676 if !md.is_empty() {
678 md.push_str("\n\n");
679 }
680 md.push_str("$$");
681 md.push_str(latex_content);
682 md.push_str("$$");
683 }
684 }
685 LayoutElementType::Image | LayoutElementType::Chart => {
687 if !md.is_empty() {
688 md.push_str("\n\n");
689 }
690 md.push_str("<div style=\"text-align: center;\"><img src=\"");
692 let img_name = format!(
694 "imgs/img_in_{}_box_{:.0}_{:.0}_{:.0}_{:.0}.jpg",
695 if element.element_type == LayoutElementType::Chart {
696 "chart"
697 } else {
698 "image"
699 },
700 element.bbox.x_min(),
701 element.bbox.y_min(),
702 element.bbox.x_max(),
703 element.bbox.y_max()
704 );
705 md.push_str(&img_name);
706 md.push_str("\" alt=\"Image\" width=\"");
707 let image_width = element.bbox.x_max() - element.bbox.x_min();
709 let width_pct = (image_width / original_image_width * 100.0) as u32;
710 let width_pct = width_pct.clamp(1, 100);
711 md.push_str(&format!("{}%", width_pct));
712 md.push_str("\" /></div>");
713 }
714 LayoutElementType::Seal => {
716 if !md.is_empty() {
717 md.push_str("\n\n");
718 }
719 md.push_str("![Seal]");
720 if let Some(text) = &element.text {
721 md.push_str("\n> ");
722 md.push_str(text);
723 }
724 }
725 _ if element.element_type.is_caption() => {
727 if let Some(text) = &element.text {
728 if !md.is_empty() {
729 md.push_str("\n\n");
730 }
731 let cleaned = clean_ocr_text(text);
732 md.push_str("<div style=\"text-align: center;\">");
733 md.push_str(&cleaned);
734 md.push_str(" </div>");
735 }
736 }
737 LayoutElementType::Abstract => {
739 if let Some(text) = &element.text {
740 if !md.is_empty() {
741 md.push_str("\n\n");
742 }
743 let formatted = format_first_line(text, " ", &["abstract", "摘要"], "## ");
744 md.push_str(&formatted);
745 }
746 }
747 LayoutElementType::Reference => {
749 if let Some(text) = &element.text {
750 if !md.is_empty() {
751 md.push_str("\n\n");
752 }
753 let formatted =
754 format_first_line(text, "\n", &["references", "参考文献"], "## ");
755 md.push_str(&formatted);
756 }
757 }
758 LayoutElementType::Content => {
760 if let Some(text) = &element.text {
761 if !md.is_empty() {
762 md.push_str("\n\n");
763 }
764 let formatted = format_content_block(text);
765 md.push_str(&formatted);
766 }
767 }
768 LayoutElementType::Footnote => {
770 if let Some(text) = &element.text {
771 if !md.is_empty() {
772 md.push_str("\n\n");
773 }
774 let formatted = format_vision_footnote_block(text);
775 md.push_str(&formatted);
776 }
777 }
778 LayoutElementType::List => {
780 if let Some(text) = &element.text {
781 if !md.is_empty() {
782 md.push_str("\n\n");
783 }
784 let cleaned = format_text_block(text);
785 for line in cleaned.lines() {
787 let line = line.trim();
788 if !line.is_empty() {
789 md.push_str("- ");
790 md.push_str(line);
791 md.push('\n');
792 }
793 }
794 }
795 }
796 LayoutElementType::Algorithm => {
798 if let Some(text) = &element.text {
799 if !md.is_empty() {
800 md.push_str("\n\n");
801 }
802 md.push_str(text.trim_matches('\n'));
803 }
804 }
805 _ if element.element_type.is_header() || element.element_type.is_footer() => {
807 continue;
810 }
811 _ => {
813 if let Some(text) = &element.text {
814 let cleaned = clean_ocr_text(text);
815 if has_bullet_markers(&cleaned) {
816 if !md.is_empty() {
817 md.push_str("\n\n");
818 }
819 format_as_bullet_list(&cleaned, &mut md);
820 } else if is_continuation {
821 let formatted = format_text_block(text);
822 md.push_str(&formatted);
823 } else {
824 if !md.is_empty() {
825 md.push_str("\n\n");
826 }
827 let formatted = format_text_block(text);
828 md.push_str(&formatted);
829 }
830 }
831 }
832 }
833
834 if element.element_type == LayoutElementType::Text
835 || element.element_type == LayoutElementType::ReferenceContent
836 {
837 prev_text_element = Some(element);
838 }
839 }
840 md.trim().to_string()
841 }
842
843 pub fn calculate_continuation_flags(&self) -> PageContinuationFlags {
852 let elements = &self.layout_elements;
853
854 if elements.is_empty() {
855 return PageContinuationFlags::new(true, true);
856 }
857
858 let page_width = self
860 .rectified_img
861 .as_ref()
862 .map(|img| img.width() as f32)
863 .or_else(|| {
864 elements
865 .iter()
866 .map(|e| e.bbox.x_max())
867 .fold(None, |acc, x| Some(acc.map_or(x, |max: f32| max.max(x))))
868 });
869
870 let text_elements: Vec<_> = elements
872 .iter()
873 .filter(|e| {
874 matches!(
875 e.element_type,
876 LayoutElementType::Text
877 | LayoutElementType::DocTitle
878 | LayoutElementType::ParagraphTitle
879 | LayoutElementType::Abstract
880 | LayoutElementType::Reference
881 )
882 })
883 .collect();
884
885 if text_elements.is_empty() {
886 return PageContinuationFlags::new(true, true);
887 }
888
889 let first = &text_elements[0];
891 let paragraph_start = is_new_paragraph_start(first, page_width);
892
893 let last = &text_elements[text_elements.len() - 1];
895 let paragraph_end = is_paragraph_complete(last, page_width);
896
897 PageContinuationFlags::new(paragraph_start, paragraph_end)
898 }
899
900 pub fn to_html(&self) -> String {
904 let mut html = String::from(
905 "<!DOCTYPE html>\n<html>\n<head>\n<meta charset=\"UTF-8\">\n</head>\n<body>\n",
906 );
907
908 for element in &self.layout_elements {
909 match element.element_type {
910 LayoutElementType::DocTitle => {
912 html.push_str("<h1>");
913 if let Some(text) = &element.text {
914 html.push_str(&Self::escape_html(text));
915 }
916 html.push_str("</h1>\n");
917 }
918 LayoutElementType::ParagraphTitle => {
920 html.push_str("<h2>");
921 if let Some(text) = &element.text {
922 html.push_str(&Self::escape_html(text));
923 }
924 html.push_str("</h2>\n");
925 }
926 LayoutElementType::Table => {
928 if let Some(table) =
929 self.tables.iter().find(|t| t.bbox.iou(&element.bbox) > 0.5)
930 {
931 if let Some(table_html) = &table.html_structure {
932 let simplified = simplify_table_html(table_html);
934 let styled = simplified.replacen(
935 "<table>",
936 "<table border=\"1\" style=\"border-collapse: collapse;\">",
937 1,
938 );
939 html.push_str(&styled);
940 html.push('\n');
941 } else {
942 html.push_str("<p>[Table]</p>\n");
943 }
944 } else {
945 html.push_str("<p>[Table]</p>\n");
946 }
947 }
948 LayoutElementType::Formula | LayoutElementType::FormulaNumber => {
950 html.push_str("<p class=\"formula\">$$");
951 if let Some(latex) = &element.text {
952 html.push_str(&Self::escape_html(latex));
953 }
954 html.push_str("$$</p>\n");
955 }
956 LayoutElementType::Image | LayoutElementType::Chart => {
958 html.push_str("<figure>\n<img alt=\"Figure\" />\n");
959 if let Some(caption) = &element.text {
960 html.push_str("<figcaption>");
961 html.push_str(&Self::escape_html(caption));
962 html.push_str("</figcaption>\n");
963 }
964 html.push_str("</figure>\n");
965 }
966 LayoutElementType::Seal => {
968 html.push_str("<figure class=\"seal\">\n<img alt=\"Seal\" />\n");
969 if let Some(text) = &element.text {
970 html.push_str("<figcaption>");
971 html.push_str(&Self::escape_html(text));
972 html.push_str("</figcaption>\n");
973 }
974 html.push_str("</figure>\n");
975 }
976 _ if element.element_type.is_caption() => {
978 if let Some(text) = &element.text {
979 html.push_str("<figcaption>");
980 html.push_str(&Self::escape_html(text));
981 html.push_str("</figcaption>\n");
982 }
983 }
984 LayoutElementType::Abstract => {
986 html.push_str("<section class=\"abstract\">\n<h3>Abstract</h3>\n<p>");
987 if let Some(text) = &element.text {
988 html.push_str(&Self::escape_html(text));
989 }
990 html.push_str("</p>\n</section>\n");
991 }
992 LayoutElementType::Reference | LayoutElementType::ReferenceContent => {
994 html.push_str("<section class=\"references\">\n<p>");
995 if let Some(text) = &element.text {
996 html.push_str(&Self::escape_html(text));
997 }
998 html.push_str("</p>\n</section>\n");
999 }
1000 LayoutElementType::List => {
1002 html.push_str("<ul>\n");
1003 if let Some(text) = &element.text {
1004 for line in text.lines() {
1005 html.push_str("<li>");
1006 html.push_str(&Self::escape_html(line));
1007 html.push_str("</li>\n");
1008 }
1009 }
1010 html.push_str("</ul>\n");
1011 }
1012 _ if element.element_type.is_header() => {
1014 html.push_str("<header>");
1015 if let Some(text) = &element.text {
1016 html.push_str(&Self::escape_html(text));
1017 }
1018 html.push_str("</header>\n");
1019 }
1020 _ if element.element_type.is_footer() => {
1022 html.push_str("<footer>");
1023 if let Some(text) = &element.text {
1024 html.push_str(&Self::escape_html(text));
1025 }
1026 html.push_str("</footer>\n");
1027 }
1028 _ => {
1030 if let Some(text) = &element.text {
1031 html.push_str("<p>");
1032 html.push_str(&Self::escape_html(text));
1033 html.push_str("</p>\n");
1034 }
1035 }
1036 }
1037 }
1038 html.push_str("</body>\n</html>");
1039 html
1040 }
1041
1042 fn escape_html(text: &str) -> String {
1044 text.replace('&', "&")
1045 .replace('<', "<")
1046 .replace('>', ">")
1047 .replace('"', """)
1048 .replace('\'', "'")
1049 }
1050
1051 pub fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
1053 serde_json::to_value(self)
1054 }
1055
1056 pub fn save_results(
1073 &self,
1074 output_dir: impl AsRef<Path>,
1075 to_json: bool,
1076 to_html: bool,
1077 ) -> std::io::Result<()> {
1078 let output_dir = output_dir.as_ref();
1079 if !output_dir.exists() {
1080 std::fs::create_dir_all(output_dir)?;
1081 }
1082
1083 let input_path = Path::new(self.input_path.as_ref());
1084 let stem = if let Some(path_str) = input_path.to_str() {
1086 if let Some(hash_idx) = path_str.rfind('#') {
1087 let base = &path_str[..hash_idx];
1089 let page_num = &path_str[hash_idx + 1..];
1090 let base_stem = Path::new(base)
1091 .file_stem()
1092 .and_then(|s| s.to_str())
1093 .unwrap_or("result");
1094 format!("{}_{}", base_stem, page_num)
1095 } else {
1096 input_path
1097 .file_stem()
1098 .and_then(|s| s.to_str())
1099 .unwrap_or("result")
1100 .to_string()
1101 }
1102 } else {
1103 "result".to_string()
1104 };
1105
1106 if to_json {
1108 let json_path = output_dir.join(format!("{}.json", stem));
1109 let json_file = std::fs::File::create(json_path)?;
1110 serde_json::to_writer_pretty(json_file, self)?;
1111 }
1112
1113 if to_html {
1115 let html_path = output_dir.join(format!("{}.html", stem));
1116 std::fs::write(html_path, self.to_html())?;
1117 }
1118
1119 Ok(())
1120 }
1121}
1122
1123fn get_seg_flag(current: &LayoutElement, prev: Option<&LayoutElement>) -> bool {
1137 const COORD_THRESHOLD: f32 = 10.0;
1138
1139 let seg_start = current.seg_start_x.unwrap_or(current.bbox.x_min());
1140 let mut context_left = current.bbox.x_min();
1141 let mut context_right = current.bbox.x_max();
1142
1143 if let Some(prev) = prev {
1144 let prev_seg_end = prev.seg_end_x.unwrap_or(prev.bbox.x_max());
1145 let prev_num_lines = prev.num_lines.unwrap_or(1);
1146
1147 let overlap_blocks = context_left < prev.bbox.x_max() && context_right > prev.bbox.x_min();
1149
1150 let edge_distance;
1151 if overlap_blocks {
1152 context_left = context_left.min(prev.bbox.x_min());
1153 context_right = context_right.max(prev.bbox.x_max());
1154 edge_distance = 0.0;
1155 } else {
1156 edge_distance = (current.bbox.x_min() - prev.bbox.x_max()).abs();
1157 }
1158
1159 let prev_end_space_small = (context_right - prev_seg_end).abs() < COORD_THRESHOLD;
1160 let current_start_space_small = seg_start - context_left < COORD_THRESHOLD;
1161 let prev_lines_more_than_one = prev_num_lines > 1;
1162 let blocks_close = edge_distance
1163 < (prev.bbox.x_max() - prev.bbox.x_min())
1164 .max(current.bbox.x_max() - current.bbox.x_min());
1165
1166 if prev_end_space_small
1167 && current_start_space_small
1168 && prev_lines_more_than_one
1169 && blocks_close
1170 {
1171 return false; }
1173
1174 true } else {
1176 if seg_start - context_left < COORD_THRESHOLD {
1178 return false; }
1180 true
1181 }
1182}
1183
1184fn is_new_paragraph_start(element: &LayoutElement, page_width: Option<f32>) -> bool {
1189 let left = element.bbox.x_min();
1190 let threshold = page_width.map_or(50.0, |w| w * 0.05); left <= threshold
1192}
1193
1194fn is_paragraph_complete(element: &LayoutElement, page_width: Option<f32>) -> bool {
1199 let right = element.bbox.x_max();
1200
1201 if let Some(width) = page_width {
1203 let right_margin = width * 0.1;
1204 return right <= (width - right_margin);
1205 }
1206
1207 true
1209}
1210
1211pub fn concatenate_markdown_pages(results: &[StructureResult]) -> String {
1224 if results.is_empty() {
1225 return String::new();
1226 }
1227
1228 if results.len() == 1 {
1229 return results[0].to_markdown();
1230 }
1231
1232 let mut markdown = String::new();
1233 let mut prev_page_end_flag = true; for result in results.iter() {
1236 let flags = result
1237 .page_continuation_flags
1238 .as_ref()
1239 .cloned()
1240 .unwrap_or_else(|| result.calculate_continuation_flags());
1241
1242 let page_markdown = result.to_markdown();
1243
1244 if page_markdown.trim().is_empty() {
1246 prev_page_end_flag = flags.paragraph_end;
1247 continue;
1248 }
1249
1250 let page_first_continues = !flags.paragraph_start;
1251 let _page_last_continues = !flags.paragraph_end;
1252
1253 if page_first_continues && !prev_page_end_flag {
1255 let last_char = markdown.chars().last();
1258 let first_char = page_markdown.chars().next();
1259
1260 let last_is_chinese = last_char.is_some_and(is_chinese_char);
1261 let first_is_chinese = first_char.is_some_and(is_chinese_char);
1262
1263 if !last_is_chinese && !first_is_chinese {
1264 markdown.push(' ');
1266 markdown.push_str(page_markdown.trim_start());
1267 } else {
1268 markdown.push_str(page_markdown.trim_start());
1270 }
1271 } else {
1272 if !markdown.is_empty() {
1274 markdown.push_str("\n\n");
1275 }
1276 markdown.push_str(&page_markdown);
1277 }
1278
1279 prev_page_end_flag = flags.paragraph_end;
1280 }
1281
1282 markdown.trim().to_string()
1283}
1284
1285fn clean_ocr_text(text: &str) -> String {
1295 text.replace("-\n", "").replace('\n', " ")
1297}
1298
1299fn format_first_line(
1311 text: &str,
1312 spliter: &str,
1313 templates: &[&str],
1314 heading_prefix: &str,
1315) -> String {
1316 let parts: Vec<&str> = text.split(spliter).collect();
1317 let mut result_parts: Vec<String> = Vec::with_capacity(parts.len());
1318 let mut found_first = false;
1319
1320 for part in &parts {
1321 if !found_first {
1322 let trimmed = part.trim();
1323 if trimmed.is_empty() {
1324 result_parts.push(part.to_string());
1325 continue;
1326 }
1327 found_first = true;
1328 if templates.iter().any(|t| trimmed.eq_ignore_ascii_case(t)) {
1330 result_parts.push(format!("{}{}\n", heading_prefix, trimmed));
1332 } else {
1333 result_parts.push(part.to_string());
1334 }
1335 } else {
1336 result_parts.push(part.to_string());
1337 }
1338 }
1339
1340 result_parts.join(spliter)
1341}
1342
1343fn format_text_block(text: &str) -> String {
1349 let dehyphenated = text.replace("-\n", "");
1351 let step1 = dehyphenated.replace("\n\n", "\n");
1353 step1.replace('\n', "\n\n")
1355}
1356
1357fn format_content_block(text: &str) -> String {
1362 let step1 = text.replace("-\n", " \n");
1364 step1.replace('\n', " \n")
1366}
1367
1368fn format_vision_footnote_block(text: &str) -> String {
1372 let dehyphenated = text.replace("-\n", "");
1373 let step1 = dehyphenated.replace("\n\n", "\n");
1374 step1.replace('\n', "\n\n")
1375}
1376
1377const BULLET_MARKERS: &[char] = &['•', '●', '◦', '▪', '◆'];
1379
1380fn has_bullet_markers(text: &str) -> bool {
1382 BULLET_MARKERS.iter().any(|&m| text.contains(m))
1383}
1384
1385fn format_as_bullet_list(text: &str, md: &mut String) {
1390 for item in text.split(|c: char| BULLET_MARKERS.contains(&c)) {
1391 let item = item.trim();
1392 if !item.is_empty() {
1393 md.push_str("- ");
1394 md.push_str(item);
1395 md.push('\n');
1396 }
1397 }
1398}
1399
1400fn is_chinese_char(c: char) -> bool {
1404 match c {
1405 '\u{4E00}'..='\u{9FFF}' | '\u{3400}'..='\u{4DBF}' | '\u{20000}'..='\u{2A6DF}' | '\u{2A700}'..='\u{2B73F}' | '\u{2B740}'..='\u{2B81F}' | '\u{2B820}'..='\u{2CEAF}' | '\u{2CEB0}'..='\u{2EBEF}' => true,
1413 _ => false,
1414 }
1415}
1416
1417fn is_lowercase(c: char) -> bool {
1419 c.is_ascii_lowercase()
1420}
1421
1422fn is_uppercase(c: char) -> bool {
1424 c.is_ascii_uppercase()
1425}
1426
1427fn is_digit(c: char) -> bool {
1429 c.is_ascii_digit()
1430}
1431
1432fn dehyphenate(text: &str) -> String {
1439 let mut result = String::with_capacity(text.len());
1440 let chars: Vec<char> = text.chars().collect();
1441 let len = chars.len();
1442 let mut i = 0;
1443
1444 let is_url_context = |pos: usize| -> bool {
1446 let start = pos.saturating_sub(10);
1447 let end = (pos + 5).min(len);
1448 let window: String = chars[start..end].iter().collect();
1449 window.contains("http") || window.contains("www") || window.contains("://")
1450 };
1451
1452 while i < len {
1453 if chars[i] == '-' {
1454 if is_url_context(i) {
1455 result.push('-');
1456 i += 1;
1457 continue;
1458 }
1459
1460 let is_artifact = if i + 1 < len && chars[i + 1] == '\n' {
1463 if i + 2 < len {
1465 is_lowercase(chars[i + 2])
1466 } else {
1467 false
1468 }
1469 } else {
1470 false
1471 };
1472
1473 if is_artifact {
1474 i += 1; } else {
1477 result.push('-');
1478 }
1479 } else {
1480 result.push(chars[i]);
1481 }
1482 i += 1;
1483 }
1484
1485 result
1486}
1487
1488fn fix_merged_words(text: &str) -> String {
1494 let mut result = String::with_capacity(text.len());
1495 let chars: Vec<char> = text.chars().collect();
1496 let mut i = 0;
1497
1498 while i < chars.len() {
1499 let current = chars[i];
1500
1501 if i > 0 {
1502 let prev = chars[i - 1];
1503
1504 if is_lowercase(prev) && is_lowercase(current) {
1507 if i > 1 && chars[i - 2] == '\'' {
1510 result.push(' ');
1511 }
1512 } else if is_lowercase(prev) && is_uppercase(current) {
1515 if i + 1 < chars.len() && is_lowercase(chars[i + 1]) {
1518 result.push(' ');
1519 }
1520 }
1521 else if ((is_digit(prev) || prev == '%') && is_uppercase(current))
1525 || (is_letter(prev)
1526 && is_digit(current)
1527 && i + 1 < chars.len()
1528 && is_letter(chars[i + 1]))
1529 {
1530 result.push(' ');
1531 }
1532 }
1533
1534 result.push(current);
1535 i += 1;
1536 }
1537
1538 result
1539}
1540
1541fn is_letter(c: char) -> bool {
1543 is_lowercase(c) || is_uppercase(c)
1544}
1545
1546fn simplify_table_html(html: &str) -> String {
1551 html.replace("<html>", "")
1552 .replace("</html>", "")
1553 .replace("<body>", "")
1554 .replace("</body>", "")
1555}
1556
1557pub fn postprocess_text(text: &str) -> String {
1564 let text = dehyphenate(text);
1565 let text = fix_merged_words(&text);
1566
1567 let mut result = String::new();
1569 let mut in_space = false;
1570
1571 for c in text.chars() {
1572 if c.is_whitespace() {
1573 if !in_space && !result.is_empty() {
1574 result.push(' ');
1575 in_space = true;
1576 }
1577 } else {
1578 if c == '.' && !result.is_empty() {
1580 let last = result.chars().last().unwrap();
1581 if is_letter(last) || is_digit(last) {
1582 result.push('.');
1583 in_space = true;
1584 continue;
1585 }
1586 }
1587 if in_space && matches!(c, '.' | ',' | '!' | '?' | ';' | ':' | ')' | ']' | '}') {
1589 result.pop(); result.push(c);
1591 continue;
1592 }
1593 result.push(c);
1594 in_space = false;
1595 }
1596 }
1597
1598 result
1599}
1600
1601fn deduplicate_sections(markdown: &str) -> String {
1607 let mut result = String::new();
1608 let mut seen_sections: std::collections::HashSet<String> = std::collections::HashSet::new();
1609
1610 for line in markdown.lines() {
1611 let trimmed = line.trim();
1612
1613 let is_section_header =
1615 trimmed.starts_with("**") && trimmed.ends_with("**") && trimmed.len() > 4;
1616
1617 let section_name = if is_section_header {
1618 trimmed[2..trimmed.len() - 2].to_string()
1619 } else {
1620 String::new()
1621 };
1622
1623 if is_section_header {
1624 if seen_sections.contains(§ion_name) {
1625 continue;
1627 }
1628 seen_sections.insert(section_name);
1629 }
1630
1631 if !result.is_empty() {
1632 result.push('\n');
1633 }
1634 result.push_str(line);
1635 }
1636
1637 result
1638}
1639
1640fn is_same_line(bbox1: &BoundingBox, bbox2: &BoundingBox) -> bool {
1645 let y1_min = bbox1.y_min();
1646 let y1_max = bbox1.y_max();
1647 let y2_min = bbox2.y_min();
1648 let y2_max = bbox2.y_max();
1649
1650 let overlap_start = y1_min.max(y2_min);
1652 let overlap_end = y1_max.min(y2_max);
1653 let overlap = (overlap_end - overlap_start).max(0.0);
1654
1655 let height1 = y1_max - y1_min;
1657 let height2 = y2_max - y2_min;
1658 let min_height = height1.min(height2);
1659
1660 min_height > 0.0 && overlap / min_height > 0.5
1662}
1663
1664fn filter_empty_formulas(markdown: &str) -> String {
1668 let mut result = String::new();
1669 let lines: Vec<&str> = markdown.lines().collect();
1670 let mut i = 0;
1671
1672 while i < lines.len() {
1673 let line = lines[i];
1674
1675 if line.trim() == "$$" {
1677 if i + 1 < lines.len() && lines[i + 1].trim() == "$$" {
1679 i += 2;
1681 if i < lines.len() && lines[i].trim().is_empty() {
1683 i += 1;
1684 }
1685 continue;
1686 }
1687 let mut j = i + 1;
1689 let has_content = if j < lines.len() {
1690 let mut found = false;
1691 while j < lines.len() {
1692 if lines[j].trim() == "$$" {
1693 break;
1694 }
1695 if !lines[j].trim().is_empty() {
1696 found = true;
1697 break;
1698 }
1699 j += 1;
1700 }
1701 found
1702 } else {
1703 false
1704 };
1705
1706 if !has_content {
1707 while i < lines.len() && lines[i].trim() != "$$" {
1709 i += 1;
1710 }
1711 if i < lines.len() {
1712 i += 1; }
1714 continue;
1715 }
1716 }
1717
1718 if !result.is_empty() {
1719 result.push('\n');
1720 }
1721 result.push_str(line);
1722 i += 1;
1723 }
1724
1725 result
1726}
1727
1728pub fn postprocess_markdown(markdown: &str) -> String {
1732 let markdown = filter_empty_formulas(markdown);
1733 let markdown = deduplicate_sections(&markdown);
1734
1735 let mut result = String::new();
1737 let mut in_code_block = false;
1738 let mut in_formula = false;
1739
1740 for line in markdown.lines() {
1741 let trimmed = line.trim();
1742
1743 if trimmed.starts_with("```") {
1745 in_code_block = !in_code_block;
1746 result.push_str(line);
1747 result.push('\n');
1748 continue;
1749 }
1750
1751 if trimmed == "$$" {
1753 in_formula = !in_formula;
1754 result.push_str(line);
1755 result.push('\n');
1756 continue;
1757 }
1758
1759 if in_code_block {
1761 result.push_str(line);
1762 result.push('\n');
1763 continue;
1764 }
1765
1766 if in_formula {
1769 let contains_dollar = line.contains('$');
1772 let is_plain_text = line.split_whitespace().count() > 3 && !line.contains('\\');
1773
1774 if contains_dollar && is_plain_text {
1775 result.push_str(&line.replace('$', "\\$"));
1776 } else if contains_dollar {
1777 result.push_str(&line.replace('$', "\\$"));
1781 } else {
1782 result.push_str(line);
1783 }
1784 result.push('\n');
1785 continue;
1786 }
1787
1788 if trimmed.starts_with('#')
1790 || trimmed.starts_with('*')
1791 || trimmed.starts_with('>')
1792 || trimmed.starts_with('|')
1793 || trimmed.starts_with('-')
1794 || trimmed.starts_with('+')
1795 {
1796 result.push_str(line);
1797 } else {
1798 result.push_str(&postprocess_text(line));
1799 }
1800 result.push('\n');
1801 }
1802
1803 result
1804}
1805
1806pub trait StructureResultExt {
1808 fn to_concatenated_markdown(results: &[Self]) -> String
1810 where
1811 Self: Sized;
1812
1813 fn save_multi_page_results(
1815 results: &[Self],
1816 output_dir: impl AsRef<std::path::Path>,
1817 base_name: &str,
1818 to_json: bool,
1819 to_markdown: bool,
1820 to_html: bool,
1821 ) -> std::io::Result<()>
1822 where
1823 Self: Sized;
1824}
1825
1826impl StructureResultExt for StructureResult {
1827 fn to_concatenated_markdown(results: &[Self]) -> String {
1828 concatenate_markdown_pages(results)
1829 }
1830
1831 fn save_multi_page_results(
1832 results: &[Self],
1833 output_dir: impl AsRef<std::path::Path>,
1834 base_name: &str,
1835 to_json: bool,
1836 to_markdown: bool,
1837 to_html: bool,
1838 ) -> std::io::Result<()>
1839 where
1840 Self: Sized,
1841 {
1842 let output_dir = output_dir.as_ref();
1843 if !output_dir.exists() {
1844 std::fs::create_dir_all(output_dir)?;
1845 }
1846
1847 for (idx, result) in results.iter().enumerate() {
1849 let page_dir = output_dir.join(format!("page_{:03}", idx));
1850 std::fs::create_dir_all(&page_dir)?;
1851 result.save_results(&page_dir, to_json, to_html)?;
1852 }
1853
1854 if to_markdown {
1856 let concat_md_path = output_dir.join(format!("{}.md", base_name));
1857 std::fs::write(concat_md_path, Self::to_concatenated_markdown(results))?;
1858 }
1859
1860 if to_json {
1862 let concat_json_path = output_dir.join(format!("{}.json", base_name));
1863 let json_file = std::fs::File::create(concat_json_path)?;
1864 serde_json::to_writer_pretty(json_file, &results)?;
1865 }
1866
1867 Ok(())
1868 }
1869}
1870
1871#[derive(Debug, Clone, Serialize, Deserialize)]
1873pub struct LayoutElement {
1874 pub bbox: BoundingBox,
1876 pub element_type: LayoutElementType,
1878 pub confidence: f32,
1880 pub label: Option<String>,
1882 pub text: Option<String>,
1884 pub order_index: Option<u32>,
1891 #[serde(skip_serializing_if = "Option::is_none")]
1895 pub seg_start_x: Option<f32>,
1896 #[serde(skip_serializing_if = "Option::is_none")]
1900 pub seg_end_x: Option<f32>,
1901 #[serde(skip_serializing_if = "Option::is_none")]
1904 pub num_lines: Option<u32>,
1905}
1906
1907impl LayoutElement {
1908 pub fn new(bbox: BoundingBox, element_type: LayoutElementType, confidence: f32) -> Self {
1910 Self {
1911 bbox,
1912 element_type,
1913 confidence,
1914 label: None,
1915 text: None,
1916 order_index: None,
1917 seg_start_x: None,
1918 seg_end_x: None,
1919 num_lines: None,
1920 }
1921 }
1922
1923 pub fn with_label(mut self, label: impl Into<String>) -> Self {
1925 self.label = Some(label.into());
1926 self
1927 }
1928
1929 pub fn with_text(mut self, text: impl Into<String>) -> Self {
1931 self.text = Some(text.into());
1932 self
1933 }
1934}
1935
1936#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
1983pub enum LayoutElementType {
1984 DocTitle,
1986 ParagraphTitle,
1988 Text,
1990 Content,
1992 Abstract,
1994
1995 Image,
1997 Table,
1999 Chart,
2001 Formula,
2003
2004 FigureTitle,
2006 TableTitle,
2008 ChartTitle,
2010 FigureTableChartTitle,
2012
2013 Header,
2015 HeaderImage,
2017 Footer,
2019 FooterImage,
2021 Footnote,
2023
2024 Seal,
2026 Number,
2028 Reference,
2030 ReferenceContent,
2032 Algorithm,
2034 FormulaNumber,
2036 AsideText,
2038 List,
2040
2041 Region,
2044
2045 Other,
2047}
2048
2049impl LayoutElementType {
2050 pub fn as_str(&self) -> &'static str {
2054 match self {
2055 LayoutElementType::DocTitle => "doc_title",
2057 LayoutElementType::ParagraphTitle => "paragraph_title",
2058 LayoutElementType::Text => "text",
2059 LayoutElementType::Content => "content",
2060 LayoutElementType::Abstract => "abstract",
2061
2062 LayoutElementType::Image => "image",
2064 LayoutElementType::Table => "table",
2065 LayoutElementType::Chart => "chart",
2066 LayoutElementType::Formula => "formula",
2067
2068 LayoutElementType::FigureTitle => "figure_title",
2070 LayoutElementType::TableTitle => "table_title",
2071 LayoutElementType::ChartTitle => "chart_title",
2072 LayoutElementType::FigureTableChartTitle => "figure_table_chart_title",
2073
2074 LayoutElementType::Header => "header",
2076 LayoutElementType::HeaderImage => "header_image",
2077 LayoutElementType::Footer => "footer",
2078 LayoutElementType::FooterImage => "footer_image",
2079 LayoutElementType::Footnote => "footnote",
2080
2081 LayoutElementType::Seal => "seal",
2083 LayoutElementType::Number => "number",
2084 LayoutElementType::Reference => "reference",
2085 LayoutElementType::ReferenceContent => "reference_content",
2086 LayoutElementType::Algorithm => "algorithm",
2087 LayoutElementType::FormulaNumber => "formula_number",
2088 LayoutElementType::AsideText => "aside_text",
2089 LayoutElementType::List => "list",
2090
2091 LayoutElementType::Region => "region",
2093
2094 LayoutElementType::Other => "other",
2096 }
2097 }
2098
2099 pub fn from_label(label: &str) -> Self {
2104 match label.to_lowercase().as_str() {
2105 "doc_title" => LayoutElementType::DocTitle,
2107 "paragraph_title" | "title" => LayoutElementType::ParagraphTitle,
2108 "text" | "paragraph" => LayoutElementType::Text,
2109 "content" => LayoutElementType::Content,
2110 "abstract" => LayoutElementType::Abstract,
2111
2112 "image" | "figure" => LayoutElementType::Image,
2114 "table" => LayoutElementType::Table,
2115 "chart" | "flowchart" => LayoutElementType::Chart,
2116 "formula" | "equation" | "display_formula" | "inline_formula" => {
2117 LayoutElementType::Formula
2118 }
2119
2120 "figure_title" => LayoutElementType::FigureTitle,
2122 "table_title" => LayoutElementType::TableTitle,
2123 "chart_title" => LayoutElementType::ChartTitle,
2124 "figure_table_chart_title" | "caption" => LayoutElementType::FigureTableChartTitle,
2125
2126 "header" => LayoutElementType::Header,
2128 "header_image" => LayoutElementType::HeaderImage,
2129 "footer" => LayoutElementType::Footer,
2130 "footer_image" => LayoutElementType::FooterImage,
2131 "footnote" | "vision_footnote" => LayoutElementType::Footnote,
2132
2133 "seal" => LayoutElementType::Seal,
2135 "number" => LayoutElementType::Number,
2136 "reference" => LayoutElementType::Reference,
2137 "reference_content" => LayoutElementType::ReferenceContent,
2138 "algorithm" => LayoutElementType::Algorithm,
2139 "formula_number" => LayoutElementType::FormulaNumber,
2140 "aside_text" => LayoutElementType::AsideText,
2141 "list" => LayoutElementType::List,
2142 "vertical_text" => LayoutElementType::Text,
2143
2144 "region" => LayoutElementType::Region,
2146
2147 _ => LayoutElementType::Other,
2150 }
2151 }
2152
2153 pub fn semantic_category(&self) -> &'static str {
2172 match self {
2173 LayoutElementType::DocTitle | LayoutElementType::ParagraphTitle => "title",
2175
2176 LayoutElementType::Text | LayoutElementType::Content | LayoutElementType::Abstract => {
2178 "text"
2179 }
2180
2181 LayoutElementType::Image | LayoutElementType::Chart => "visual",
2183
2184 LayoutElementType::Table => "table",
2186
2187 LayoutElementType::FigureTitle
2189 | LayoutElementType::TableTitle
2190 | LayoutElementType::ChartTitle
2191 | LayoutElementType::FigureTableChartTitle => "caption",
2192
2193 LayoutElementType::Header | LayoutElementType::HeaderImage => "header",
2195
2196 LayoutElementType::Footer
2198 | LayoutElementType::FooterImage
2199 | LayoutElementType::Footnote => "footer",
2200
2201 LayoutElementType::Formula | LayoutElementType::FormulaNumber => "formula",
2203
2204 LayoutElementType::Seal
2206 | LayoutElementType::Number
2207 | LayoutElementType::Reference
2208 | LayoutElementType::ReferenceContent
2209 | LayoutElementType::Algorithm
2210 | LayoutElementType::AsideText => "special",
2211
2212 LayoutElementType::List => "list",
2214
2215 LayoutElementType::Region => "region",
2217
2218 LayoutElementType::Other => "other",
2220 }
2221 }
2222
2223 pub fn is_title(&self) -> bool {
2225 matches!(
2226 self,
2227 LayoutElementType::DocTitle | LayoutElementType::ParagraphTitle
2228 )
2229 }
2230
2231 pub fn is_visual(&self) -> bool {
2233 matches!(self, LayoutElementType::Image | LayoutElementType::Chart)
2234 }
2235
2236 pub fn is_caption(&self) -> bool {
2238 matches!(
2239 self,
2240 LayoutElementType::FigureTitle
2241 | LayoutElementType::TableTitle
2242 | LayoutElementType::ChartTitle
2243 | LayoutElementType::FigureTableChartTitle
2244 )
2245 }
2246
2247 pub fn is_header(&self) -> bool {
2249 matches!(
2250 self,
2251 LayoutElementType::Header | LayoutElementType::HeaderImage
2252 )
2253 }
2254
2255 pub fn is_footer(&self) -> bool {
2257 matches!(
2258 self,
2259 LayoutElementType::Footer
2260 | LayoutElementType::FooterImage
2261 | LayoutElementType::Footnote
2262 )
2263 }
2264
2265 pub fn is_formula(&self) -> bool {
2267 matches!(
2268 self,
2269 LayoutElementType::Formula | LayoutElementType::FormulaNumber
2270 )
2271 }
2272
2273 pub fn should_ocr(&self) -> bool {
2275 matches!(
2276 self,
2277 LayoutElementType::Text
2278 | LayoutElementType::Content
2279 | LayoutElementType::Abstract
2280 | LayoutElementType::DocTitle
2281 | LayoutElementType::ParagraphTitle
2282 | LayoutElementType::FigureTitle
2283 | LayoutElementType::TableTitle
2284 | LayoutElementType::ChartTitle
2285 | LayoutElementType::FigureTableChartTitle
2286 | LayoutElementType::Header
2287 | LayoutElementType::HeaderImage
2288 | LayoutElementType::Footer
2289 | LayoutElementType::FooterImage
2290 | LayoutElementType::Footnote
2291 | LayoutElementType::Reference
2292 | LayoutElementType::ReferenceContent
2293 | LayoutElementType::Algorithm
2294 | LayoutElementType::AsideText
2295 | LayoutElementType::List
2296 | LayoutElementType::Number
2297 )
2298 }
2299}
2300
2301pub fn remove_overlapping_layout_elements(
2306 layout_elements: &mut Vec<LayoutElement>,
2307 overlap_threshold: f32,
2308) -> usize {
2309 use std::collections::HashSet;
2310
2311 if layout_elements.len() <= 1 {
2312 return 0;
2313 }
2314
2315 let bboxes: Vec<_> = layout_elements.iter().map(|e| e.bbox.clone()).collect();
2316 let labels: Vec<&str> = layout_elements
2317 .iter()
2318 .map(|e| e.element_type.as_str())
2319 .collect();
2320
2321 let remove_indices =
2322 crate::processors::get_overlap_removal_indices(&bboxes, &labels, overlap_threshold);
2323 if remove_indices.is_empty() {
2324 return 0;
2325 }
2326
2327 let remove_set: HashSet<usize> = remove_indices.into_iter().collect();
2328 let before = layout_elements.len();
2329
2330 let mut idx = 0;
2331 layout_elements.retain(|_| {
2332 let keep = !remove_set.contains(&idx);
2333 idx += 1;
2334 keep
2335 });
2336
2337 before.saturating_sub(layout_elements.len())
2338}
2339
2340pub fn apply_standardized_layout_label_fixes(layout_elements: &mut [LayoutElement]) {
2344 if layout_elements.is_empty() {
2345 return;
2346 }
2347
2348 let mut footnote_indices: Vec<usize> = Vec::new();
2349 let mut paragraph_title_indices: Vec<usize> = Vec::new();
2350 let mut bottom_text_y_max: f32 = 0.0;
2351 let mut max_block_area: f32 = 0.0;
2352 let mut doc_title_num: usize = 0;
2353
2354 for (idx, elem) in layout_elements.iter().enumerate() {
2355 let area =
2356 (elem.bbox.x_max() - elem.bbox.x_min()) * (elem.bbox.y_max() - elem.bbox.y_min());
2357 max_block_area = max_block_area.max(area);
2358
2359 match elem.element_type {
2360 LayoutElementType::Footnote => footnote_indices.push(idx),
2361 LayoutElementType::ParagraphTitle => paragraph_title_indices.push(idx),
2362 LayoutElementType::Text => {
2363 bottom_text_y_max = bottom_text_y_max.max(elem.bbox.y_max());
2364 }
2365 LayoutElementType::DocTitle => doc_title_num += 1,
2366 _ => {}
2367 }
2368 }
2369
2370 for idx in footnote_indices {
2371 if layout_elements[idx].bbox.y_max() < bottom_text_y_max {
2372 layout_elements[idx].element_type = LayoutElementType::Text;
2373 layout_elements[idx].label = Some("text".to_string());
2374 }
2375 }
2376
2377 let only_one_paragraph_title = paragraph_title_indices.len() == 1 && doc_title_num == 0;
2378 if only_one_paragraph_title {
2379 let idx = paragraph_title_indices[0];
2380 let area = (layout_elements[idx].bbox.x_max() - layout_elements[idx].bbox.x_min())
2381 * (layout_elements[idx].bbox.y_max() - layout_elements[idx].bbox.y_min());
2382
2383 let title_area_ratio_threshold = 0.3f32;
2384 if area > max_block_area * title_area_ratio_threshold {
2385 layout_elements[idx].element_type = LayoutElementType::DocTitle;
2386 layout_elements[idx].label = Some("doc_title".to_string());
2387 }
2388 }
2389}
2390
2391#[derive(Debug, Clone, Serialize, Deserialize)]
2393pub struct TableResult {
2394 pub bbox: BoundingBox,
2396 pub table_type: TableType,
2398 pub classification_confidence: Option<f32>,
2400 pub structure_confidence: Option<f32>,
2402 pub cells: Vec<TableCell>,
2404 pub html_structure: Option<String>,
2406 pub cell_texts: Option<Vec<Option<String>>>,
2408 #[serde(skip)]
2410 pub structure_tokens: Option<Vec<String>>,
2411 #[serde(skip)]
2415 pub detected_cell_bboxes: Option<Vec<BoundingBox>>,
2416 #[serde(skip)]
2420 pub is_e2e: bool,
2421}
2422
2423impl TableResult {
2424 pub fn new(bbox: BoundingBox, table_type: TableType) -> Self {
2426 Self {
2427 bbox,
2428 table_type,
2429 classification_confidence: None,
2430 structure_confidence: None,
2431 cells: Vec::new(),
2432 html_structure: None,
2433 cell_texts: None,
2434 structure_tokens: None,
2435 detected_cell_bboxes: None,
2436 is_e2e: false,
2437 }
2438 }
2439
2440 pub fn with_classification_confidence(mut self, confidence: f32) -> Self {
2442 self.classification_confidence = Some(confidence);
2443 self
2444 }
2445
2446 pub fn with_structure_confidence(mut self, confidence: f32) -> Self {
2448 self.structure_confidence = Some(confidence);
2449 self
2450 }
2451
2452 pub fn with_cells(mut self, cells: Vec<TableCell>) -> Self {
2454 self.cells = cells;
2455 self
2456 }
2457
2458 pub fn with_html_structure(mut self, html: impl Into<String>) -> Self {
2460 self.html_structure = Some(html.into());
2461 self
2462 }
2463
2464 pub fn with_cell_texts(mut self, texts: Vec<Option<String>>) -> Self {
2466 self.cell_texts = Some(texts);
2467 self
2468 }
2469
2470 pub fn with_structure_tokens(mut self, tokens: Vec<String>) -> Self {
2472 self.structure_tokens = Some(tokens);
2473 self
2474 }
2475
2476 pub fn with_detected_cell_bboxes(mut self, bboxes: Vec<BoundingBox>) -> Self {
2478 self.detected_cell_bboxes = Some(bboxes);
2479 self
2480 }
2481
2482 pub fn with_e2e(mut self, is_e2e: bool) -> Self {
2484 self.is_e2e = is_e2e;
2485 self
2486 }
2487
2488 pub fn confidence(&self) -> Option<f32> {
2500 match (self.classification_confidence, self.structure_confidence) {
2501 (Some(cls), Some(str)) => Some(cls.min(str)),
2502 (None, Some(str)) => Some(str),
2503 (Some(cls), None) => Some(cls),
2504 (None, None) => None,
2505 }
2506 }
2507
2508 pub fn has_structure(&self) -> bool {
2513 !self.cells.is_empty() || self.html_structure.is_some()
2514 }
2515}
2516
2517#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
2519pub enum TableType {
2520 Wired,
2522 Wireless,
2524 Unknown,
2526}
2527
2528#[derive(Debug, Clone, Serialize, Deserialize)]
2530pub struct TableCell {
2531 pub bbox: BoundingBox,
2533 pub row: Option<usize>,
2535 pub col: Option<usize>,
2537 pub row_span: Option<usize>,
2539 pub col_span: Option<usize>,
2541 pub confidence: f32,
2543 pub text: Option<String>,
2545}
2546
2547impl TableCell {
2548 pub fn new(bbox: BoundingBox, confidence: f32) -> Self {
2550 Self {
2551 bbox,
2552 row: None,
2553 col: None,
2554 row_span: None,
2555 col_span: None,
2556 confidence,
2557 text: None,
2558 }
2559 }
2560
2561 pub fn with_position(mut self, row: usize, col: usize) -> Self {
2563 self.row = Some(row);
2564 self.col = Some(col);
2565 self
2566 }
2567
2568 pub fn with_span(mut self, row_span: usize, col_span: usize) -> Self {
2570 self.row_span = Some(row_span);
2571 self.col_span = Some(col_span);
2572 self
2573 }
2574
2575 pub fn with_text(mut self, text: impl Into<String>) -> Self {
2577 self.text = Some(text.into());
2578 self
2579 }
2580}
2581
2582#[derive(Debug, Clone, Serialize, Deserialize)]
2584pub struct FormulaResult {
2585 pub bbox: BoundingBox,
2587 pub latex: String,
2589 pub confidence: f32,
2591}
2592
2593impl FormulaResult {
2594 pub fn new(bbox: BoundingBox, latex: impl Into<String>, confidence: f32) -> Self {
2596 Self {
2597 bbox,
2598 latex: latex.into(),
2599 confidence,
2600 }
2601 }
2602}
2603
2604#[cfg(test)]
2605mod tests {
2606 use super::*;
2607
2608 #[test]
2609 fn test_structure_result_creation() {
2610 let result = StructureResult::new("test.jpg", 0);
2611 assert_eq!(result.input_path.as_ref(), "test.jpg");
2612 assert_eq!(result.index, 0);
2613 assert!(result.layout_elements.is_empty());
2614 assert!(result.tables.is_empty());
2615 assert!(result.formulas.is_empty());
2616 assert!(result.text_regions.is_none());
2617 }
2618
2619 #[test]
2620 fn test_layout_element_type_as_str() {
2621 assert_eq!(LayoutElementType::Text.as_str(), "text");
2622 assert_eq!(LayoutElementType::Table.as_str(), "table");
2623 assert_eq!(LayoutElementType::Formula.as_str(), "formula");
2624 }
2625
2626 #[test]
2627 fn test_table_result_creation() {
2628 let bbox = BoundingBox::from_coords(0.0, 0.0, 100.0, 100.0);
2629 let table = TableResult::new(bbox, TableType::Wired);
2630 assert_eq!(table.table_type, TableType::Wired);
2631 assert!(table.cells.is_empty());
2632 assert!(table.html_structure.is_none());
2633 }
2634
2635 #[test]
2636 fn test_structure_result_export() {
2637 let bbox = BoundingBox::from_coords(0.0, 0.0, 100.0, 100.0);
2638 let mut result = StructureResult::new("test.jpg", 0);
2639
2640 let title = LayoutElement::new(bbox.clone(), LayoutElementType::DocTitle, 1.0)
2641 .with_text("Test Document");
2642
2643 let text =
2644 LayoutElement::new(bbox.clone(), LayoutElementType::Text, 1.0).with_text("Hello world");
2645
2646 result = result.with_layout_elements(vec![title, text]);
2647
2648 let md = result.to_markdown();
2649 assert!(md.contains("# Test Document"));
2650 assert!(md.contains("Hello world"));
2651
2652 let html = result.to_html();
2653 assert!(html.contains("<h1>Test Document</h1>"));
2654 assert!(html.contains("<p>Hello world</p>"));
2655 }
2656
2657 #[test]
2658 fn test_format_title_with_level_keywords() {
2659 let (level, text) = format_title_with_level("Abstract", None);
2660 assert_eq!(level, 2);
2661 assert_eq!(text, "Abstract");
2662
2663 let (level, text) = format_title_with_level("References:", None);
2664 assert_eq!(level, 2);
2665 assert_eq!(text, "References:");
2666 }
2667
2668 #[test]
2669 fn test_format_title_with_level_cluster_fallback() {
2670 let (level, text) = format_title_with_level("Unnumbered Heading", Some(4));
2671 assert_eq!(level, 4);
2672 assert_eq!(text, "Unnumbered Heading");
2673 }
2674
2675 #[test]
2676 fn test_to_markdown_skips_footnote() {
2677 let mut result = StructureResult::new("test.jpg", 0);
2678 let body = LayoutElement::new(
2679 BoundingBox::from_coords(0.0, 0.0, 100.0, 30.0),
2680 LayoutElementType::Text,
2681 1.0,
2682 )
2683 .with_text("Body");
2684 let footnote = LayoutElement::new(
2685 BoundingBox::from_coords(0.0, 40.0, 100.0, 60.0),
2686 LayoutElementType::Footnote,
2687 1.0,
2688 )
2689 .with_text("Footnote text");
2690 result = result.with_layout_elements(vec![body, footnote]);
2691
2692 let md = result.to_markdown();
2693 assert!(md.contains("Body"));
2694 assert!(!md.contains("Footnote text"));
2695 }
2696
2697 #[test]
2698 fn test_to_markdown_doc_title_joins_lines_with_space() {
2699 let mut result = StructureResult::new("test.jpg", 0);
2700 let title = LayoutElement::new(
2701 BoundingBox::from_coords(0.0, 0.0, 100.0, 20.0),
2702 LayoutElementType::DocTitle,
2703 1.0,
2704 )
2705 .with_text("Main\nTitle");
2706 result = result.with_layout_elements(vec![title]);
2707 let md = result.to_markdown();
2708 assert!(md.contains("# Main Title"));
2709 }
2710
2711 #[test]
2712 fn test_to_markdown_content_uses_soft_breaks() {
2713 let mut result = StructureResult::new("test.jpg", 0);
2714 let toc = LayoutElement::new(
2715 BoundingBox::from_coords(0.0, 0.0, 100.0, 40.0),
2716 LayoutElementType::Content,
2717 1.0,
2718 )
2719 .with_text("1 Intro\n2 Method");
2720 result = result.with_layout_elements(vec![toc]);
2721 let md = result.to_markdown();
2722 assert!(md.contains("1 Intro \n2 Method"));
2723 }
2724
2725 #[test]
2726 fn test_infer_paragraph_title_levels_by_height() {
2727 let titles = vec![
2728 LayoutElement::new(
2729 BoundingBox::from_coords(0.0, 0.0, 100.0, 40.0),
2730 LayoutElementType::ParagraphTitle,
2731 1.0,
2732 )
2733 .with_text("Large"),
2734 LayoutElement::new(
2735 BoundingBox::from_coords(0.0, 50.0, 100.0, 74.0),
2736 LayoutElementType::ParagraphTitle,
2737 1.0,
2738 )
2739 .with_text("Medium"),
2740 LayoutElement::new(
2741 BoundingBox::from_coords(0.0, 80.0, 100.0, 98.0),
2742 LayoutElementType::ParagraphTitle,
2743 1.0,
2744 )
2745 .with_text("Small"),
2746 ];
2747
2748 let levels = infer_paragraph_title_levels(&titles);
2749 let l0 = levels.get(&0).copied().unwrap_or(2);
2750 let l1 = levels.get(&1).copied().unwrap_or(2);
2751 let l2 = levels.get(&2).copied().unwrap_or(2);
2752 assert!(l0 <= l1 && l1 <= l2);
2753 }
2754
2755 #[test]
2756 fn test_infer_paragraph_title_levels_semantic_vote_wins_tie() {
2757 let titles = vec![
2758 LayoutElement::new(
2759 BoundingBox::from_coords(0.0, 0.0, 100.0, 40.0),
2760 LayoutElementType::ParagraphTitle,
2761 1.0,
2762 )
2763 .with_text("1.1 Detail"),
2764 LayoutElement::new(
2765 BoundingBox::from_coords(0.0, 50.0, 100.0, 70.0),
2766 LayoutElementType::ParagraphTitle,
2767 1.0,
2768 )
2769 .with_text("2 Intro"),
2770 ];
2771
2772 let levels = infer_paragraph_title_levels(&titles);
2773 assert_eq!(levels.get(&0).copied(), Some(3));
2774 assert_eq!(levels.get(&1).copied(), Some(2));
2775 }
2776
2777 #[test]
2778 fn test_infer_paragraph_title_levels_uses_relative_indent_signal() {
2779 let titles = vec![
2780 LayoutElement::new(
2781 BoundingBox::from_coords(0.0, 0.0, 100.0, 24.0),
2782 LayoutElementType::ParagraphTitle,
2783 1.0,
2784 )
2785 .with_text("Heading A"),
2786 LayoutElement::new(
2787 BoundingBox::from_coords(40.0, 40.0, 140.0, 64.0),
2788 LayoutElementType::ParagraphTitle,
2789 1.0,
2790 )
2791 .with_text("Heading B"),
2792 ];
2793
2794 let levels = infer_paragraph_title_levels(&titles);
2795 let left_level = levels.get(&0).copied().unwrap_or(2);
2796 let indented_level = levels.get(&1).copied().unwrap_or(2);
2797 assert!(left_level < indented_level);
2798 }
2799}