pub mod alignment;
pub mod captions;
pub mod cell_merge;
pub mod drawings;
pub mod features;
pub mod grid;
pub mod heading_lines;
pub mod judge;
pub mod lattice;
pub mod line_gnn;
pub mod lines;
#[doc(hidden)]
pub mod models;
pub mod pairs;
pub mod raster;
pub mod region;
pub mod separations;
pub mod splits;
pub mod stitch;
pub mod style;
pub mod table_split;
pub mod tables;
#[cfg(feature = "ort-tables")]
pub mod tables_tatr;
pub mod vector_proposals;
pub mod verifier;
use crate::blocks::{Block, BlockKind};
#[cfg(test)]
use crate::table_trace::TableOrigin;
pub use features::DocumentFacts;
use features::{block_vector, page_features};
use lines::{lines_of, LineContext};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Switches {
pub list_item: bool,
pub table_region: bool,
pub table_structure: bool,
pub table_judge: bool,
pub table_lattice: bool,
pub list_opens: bool,
pub heading_lines: bool,
pub captions: bool,
pub line_gnn: bool,
pub line_image: bool,
pub line_gnn_v2: bool,
pub demote_per_cell: bool,
pub table_stitch_header: bool,
pub table_stitch_rows: bool,
pub table_stitch_merge: bool,
pub table_is_v2: bool,
pub table_drop_empty_columns: bool,
pub table_cell_text_by_line: bool,
pub table_v2_release_rejoin: bool,
pub table_v2_release_column_order: bool,
pub table_split_side_by_side: bool,
pub table_split_stacked: bool,
pub table_runs_on_figures: bool,
pub markdown_header_block_figures: bool,
pub table_corridor_split: bool,
pub markdown_header_block_rowspan_crossing: bool,
pub markdown_header_block_word_table: bool,
pub markdown_header_block_mixed_body: bool,
pub markdown_headerless_html: bool,
pub markdown_header_block_stub_years: bool,
}
pub const SWITCHES: Switches = Switches {
list_item: true,
list_opens: true,
heading_lines: true,
captions: CAPTIONS_SHIPPED,
line_gnn: LINE_GNN_SHIPPED,
line_image: LINE_IMAGE_SHIPPED,
line_gnn_v2: LINE_GNN_V2_SHIPPED,
demote_per_cell: DEMOTE_PER_CELL_SHIPPED,
table_region: true,
table_structure: true,
table_judge: true,
table_lattice: true,
table_stitch_header: true,
table_stitch_rows: true,
table_stitch_merge: true,
table_is_v2: TABLE_IS_V2_SHIPPED,
table_drop_empty_columns: true,
table_cell_text_by_line: true,
table_v2_release_rejoin: true,
table_v2_release_column_order: TABLE_V2_RELEASE_COLUMN_ORDER_SHIPPED,
table_split_side_by_side: false,
table_split_stacked: true,
table_runs_on_figures: true,
markdown_header_block_figures: true,
table_corridor_split: TABLE_CORRIDOR_SPLIT_SHIPPED,
markdown_header_block_rowspan_crossing: true,
markdown_header_block_word_table: true,
markdown_header_block_mixed_body: true,
markdown_headerless_html: true,
markdown_header_block_stub_years: true,
};
const LIST_ITEM_THRESHOLD: f64 = models::list_item::THRESHOLD;
const LIST_OPENS_THRESHOLD: f64 = models::list_opens::THRESHOLD;
const HEADING_LINES_THRESHOLD: f64 = models::heading_lines::THRESHOLD;
const CAPTION_LINES_THRESHOLD: f64 = models::caption_lines::THRESHOLD;
const LINE_GNN_HEADING: f64 = 0.6;
const LINE_GNN_SHIPPED: bool = true;
const LINE_IMAGE_SHIPPED: bool = false;
const LINE_GNN_V2_SHIPPED: bool = true;
const LINE_GNN_V2_HEADING: f64 = 0.5;
const DEMOTE_PER_CELL_SHIPPED: bool = true;
const TABLE_IS_V2_SHIPPED: bool = true;
const TABLE_V2_RELEASE_COLUMN_ORDER_SHIPPED: bool = false;
const TABLE_CORRIDOR_SPLIT_SHIPPED: bool = true;
const LINE_GNN_KEEP_HEADING: f64 = 0.5;
fn heading_bar(kind: &str, switches: Switches) -> f64 {
if !switches.line_gnn {
HEADING_LINES_THRESHOLD
} else if kind == "heading" {
LINE_GNN_KEEP_HEADING
} else if switches.line_gnn_v2 {
LINE_GNN_V2_HEADING
} else {
LINE_GNN_HEADING
}
}
const CAPTIONS_SHIPPED: bool = true;
const LIST_ITEM_DEMOTE_BELOW: f64 = 0.15;
fn probability(margin: f64) -> f64 {
1.0 / (1.0 + (-margin).exp())
}
#[doc(hidden)]
pub fn heading_probabilities(blocks: &[Block], pages: &[(f64, f64)]) -> Vec<Option<f64>> {
let facts = document_facts(blocks);
let mut out = vec![None; blocks.len()];
for (page, members) in pages_of(blocks).iter().enumerate() {
let refs: Vec<&Block> = members.iter().map(|&i| &blocks[i]).collect();
let size = pages.get(page).copied().unwrap_or((0.0, 0.0));
let vectors = page_vectors(&refs, size, &facts);
let rows = heading_lines::page_rows(&vectors);
for ((&i, v), rows) in members.iter().zip(&vectors).zip(&rows) {
if matches!(v.kind, "paragraph" | "heading" | "table_cell") && !rows.is_empty() {
let sum: f64 = rows
.iter()
.map(|row| probability(models::heading_lines::predict_margin(row)))
.sum();
out[i] = Some(sum / rows.len() as f64);
}
}
}
out
}
#[derive(Debug, Clone, Default, PartialEq)]
struct Decision {
item: Option<bool>,
opens_at: Vec<usize>,
first_opens: bool,
heading_lines: Vec<bool>,
p_heading_lines: Vec<f64>,
line_classes: Vec<[f64; 4]>,
caption_lines: Vec<bool>,
p_caption_lines: Vec<f64>,
line_breaks: Vec<usize>,
link_after: Option<f64>,
p_item: Option<f64>,
p_opens: Vec<f64>,
p_first_opens: Option<f64>,
}
#[doc(hidden)]
pub fn explain(blocks: &[Block], pages: &[(f64, f64)]) -> Vec<(usize, Option<f64>, Vec<f64>)> {
decide(blocks, pages, SWITCHES)
.into_iter()
.enumerate()
.map(|(i, d)| (i, d.p_item, d.p_opens))
.collect()
}
#[doc(hidden)]
#[derive(Debug, Clone)]
pub struct BlockVectors {
pub page: usize,
pub kind: &'static str,
pub block: [f64; 30],
pub lines: Vec<[f64; 23]>,
pub line_boxes: Vec<([f64; 4], f64)>,
pub style: Vec<[f64; 5]>,
}
fn pages_of(blocks: &[Block]) -> Vec<Vec<usize>> {
let page_count = blocks.iter().map(|b| b.page + 1).max().unwrap_or(0);
let mut by_page: Vec<Vec<usize>> = vec![Vec::new(); page_count];
for (i, b) in blocks.iter().enumerate() {
by_page[b.page].push(i);
}
by_page
}
pub fn document_facts(blocks: &[Block]) -> DocumentFacts {
let faces = blocks.iter().map(|b| {
(
b.spans.iter().find_map(|s| s.font_name.clone()),
b.text.chars().count(),
)
});
DocumentFacts::of_texts_and_faces(
pages_of(blocks)
.into_iter()
.map(|members| members.into_iter().map(|i| blocks[i].text.as_str())),
faces,
)
}
fn page_vectors(page: &[&Block], size: (f64, f64), facts: &DocumentFacts) -> Vec<BlockVectors> {
let number = page.first().map_or(0, |b| b.page);
let texts: Vec<String> = page.iter().map(|b| b.text.clone()).collect();
let mut features = page_features(page, size.0, size.1);
facts.fill(&mut features, &texts);
let grouped = lines::grouped_lines(&features.words, page.len());
vectors_from_parts(page, number, &features, &grouped)
}
fn vectors_from_parts(
page: &[&Block],
number: usize,
features: &features::PageFeatures,
grouped: &[Vec<Vec<&features::PageWord>>],
) -> Vec<BlockVectors> {
let heads = lines::heads_on_grouped(grouped);
let mut styles = style::page_style_on_grouped(grouped, &features.words).into_iter();
page.iter()
.enumerate()
.map(|(at, own)| {
let block = &features.blocks[at];
let context = LineContext {
bbox: own.bbox,
kind: block.kind,
size: block.size,
body_size: features.body_size,
prev_ends_colon: (at > 0).then(|| page[at - 1].text.trim().ends_with(':')),
};
let piece: &[Vec<&features::PageWord>] =
grouped.get(at).map(Vec::as_slice).unwrap_or(&[]);
BlockVectors {
page: number,
kind: block.kind,
block: block_vector(block),
lines: lines::line_vectors_on_lines(piece, &context, &heads),
line_boxes: lines::line_boxes_on_lines(piece),
style: styles.next().unwrap_or_default(),
}
})
.collect()
}
#[doc(hidden)]
pub fn document_vectors(blocks: &[Block], pages: &[(f64, f64)]) -> Vec<BlockVectors> {
let facts = document_facts(blocks);
let mut out: Vec<Option<BlockVectors>> = vec![None; blocks.len()];
for (page, members) in pages_of(blocks).iter().enumerate() {
let refs: Vec<&Block> = members.iter().map(|&i| &blocks[i]).collect();
let size = pages.get(page).copied().unwrap_or((0.0, 0.0));
for (&i, v) in members.iter().zip(page_vectors(&refs, size, &facts)) {
out[i] = Some(v);
}
}
out.into_iter()
.map(|v| v.expect("every block is on a page"))
.collect()
}
fn decide_page(
page: &[&Block],
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
images: &[[f64; 4]],
) -> Vec<Decision> {
decide_page_traced(page, size, facts, switches, images).0
}
fn decide_page_traced(
page: &[&Block],
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
images: &[[f64; 4]],
) -> (
Vec<Decision>,
Vec<BlockVectors>,
Option<line_gnn::LineVerdicts>,
) {
let vectors = page_vectors(page, size, facts);
if !switches.list_item
&& !switches.list_opens
&& !switches.heading_lines
&& !switches.captions
&& !switches.line_gnn
{
return (vec![Decision::default(); page.len()], vectors, None);
}
let line_rows = switches
.heading_lines
.then(|| heading_lines::page_rows(&vectors));
let classes = switches
.list_item
.then(|| heading_lines::page_class_probabilities(&vectors));
let gnn = switches
.line_gnn
.then(|| line_gnn::verdicts(&vectors, size, images, gnn_of(switches)));
let heading_rows = if switches.line_gnn {
heading_lines::page_rows(&vectors)
} else {
Vec::new()
};
let caption_rows = switches
.captions
.then(|| captions::page_caption_rows(&vectors, images));
let decisions = decide_from_parts(
page,
&vectors,
switches,
line_rows.as_ref(),
classes.as_ref(),
gnn.as_ref(),
&heading_rows,
caption_rows.as_ref(),
);
(decisions, vectors, gnn)
}
#[allow(clippy::too_many_arguments)]
fn decide_from_parts(
page: &[&Block],
vectors: &[BlockVectors],
switches: Switches,
line_rows: Option<&Vec<Vec<[f64; heading_lines::WIDTH]>>>,
classes: Option<&Vec<Vec<[f64; 4]>>>,
gnn: Option<&line_gnn::LineVerdicts>,
heading_rows: &[Vec<[f64; heading_lines::WIDTH]>],
caption_rows: Option<&Vec<Vec<[f64; captions::CAPTION_WIDTH]>>>,
) -> Vec<Decision> {
let mut decisions = vec![Decision::default(); page.len()];
let mut first_line = Vec::with_capacity(vectors.len());
let mut running = 0usize;
for v in vectors {
first_line.push(running);
running += v.lines.len();
}
for (at, (decision, v)) in decisions.iter_mut().zip(vectors).enumerate() {
if let Some(rows) = caption_rows {
if v.kind == "paragraph" {
let p: Vec<f64> = rows[at]
.iter()
.map(|row| probability(models::caption_lines::predict_margin(row)))
.collect();
decision.caption_lines = p.iter().map(|p| *p >= CAPTION_LINES_THRESHOLD).collect();
decision.p_caption_lines = p;
}
}
if let Some(gnn) = gnn {
if !v.lines.is_empty() {
let last = first_line[at] + v.lines.len() - 1;
decision.link_after = gnn.link_next.get(last).map(|p| f64::from(*p));
}
if matches!(v.kind, "paragraph" | "heading" | "table_cell") && !v.lines.is_empty() {
let from = first_line[at];
let to = from + v.lines.len();
let heading = gnn_of(switches).class("heading").unwrap_or(1);
let p: Vec<f64> = gnn.classes[from..to]
.iter()
.zip(&heading_rows[at])
.map(|(c, row)| {
0.5 * (f64::from(c[heading])
+ probability(models::heading_lines::predict_margin(row)))
})
.collect();
let bar = heading_bar(v.kind, switches);
decision.heading_lines = p.iter().map(|p| *p >= bar).collect();
decision.p_heading_lines = p;
if v.kind != "table_cell" {
decision.line_breaks = (from + 1..to)
.filter(|&i| gnn.link_next[i - 1] < 0.5)
.map(|i| i - from)
.filter(|&k| {
let (above, _) = v.line_boxes[k - 1];
let (below, _) = v.line_boxes[k];
below[3] <= above[1]
})
.collect();
}
}
}
if let Some(rows) = line_rows.filter(|_| !switches.line_gnn) {
decision.line_classes = rows[at]
.iter()
.map(heading_lines::class_probabilities)
.collect();
if matches!(v.kind, "paragraph" | "heading" | "table_cell") {
let p: Vec<f64> = rows[at]
.iter()
.map(|row| probability(models::heading_lines::predict_margin(row)))
.collect();
decision.heading_lines = p.iter().map(|p| *p >= HEADING_LINES_THRESHOLD).collect();
decision.p_heading_lines = p;
}
}
if switches.list_opens && matches!(v.kind, "paragraph" | "list_item") {
let p: Vec<f64> = v
.lines
.iter()
.map(|line| probability(models::list_opens::predict_margin(line)))
.collect();
decision.first_opens = p.first().is_some_and(|p| *p >= LIST_OPENS_THRESHOLD);
decision.p_first_opens = p.first().copied();
decision.opens_at = p
.iter()
.enumerate()
.skip(1)
.filter(|(_, p)| **p >= LIST_OPENS_THRESHOLD)
.map(|(k, _)| k)
.collect();
decision.p_opens = p.into_iter().skip(1).collect();
}
if switches.list_item && matches!(v.kind, "paragraph" | "list_item") {
let p = list_item_probability(v, classes.map(|c| c[at].as_slice()));
decision.item = (p >= LIST_ITEM_THRESHOLD).then_some(true);
decision.p_item = Some(p);
}
}
let members: Vec<usize> = (0..page.len()).collect();
demote_confirmed_runs(&members, page, vectors, &mut decisions);
decisions
}
#[allow(clippy::too_many_arguments)]
fn decide_page_on_shared(
page: &[&Block],
vectors: &[BlockVectors],
size: (f64, f64),
switches: Switches,
images: &[[f64; 4]],
rows: Option<&Vec<Vec<[f64; heading_lines::WIDTH]>>>,
classes: Option<&Vec<Vec<[f64; 4]>>>,
) -> Vec<Decision> {
if !switches.list_item
&& !switches.list_opens
&& !switches.heading_lines
&& !switches.captions
&& !switches.line_gnn
{
return vec![Decision::default(); page.len()];
}
let gnn = switches
.line_gnn
.then(|| line_gnn::verdicts(vectors, size, images, gnn_of(switches)));
let caption_rows = rows
.filter(|_| switches.captions)
.map(|rows| captions::page_caption_rows_on_rows(vectors, images, rows));
decide_from_parts(
page,
vectors,
switches,
rows.filter(|_| switches.heading_lines),
classes.filter(|_| switches.list_item),
gnn.as_ref(),
rows.filter(|_| switches.line_gnn)
.map(Vec::as_slice)
.unwrap_or(&[]),
caption_rows.as_ref(),
)
}
fn judge_and_decide_shared(
page: &[Block],
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
rules: &[crate::text::HorizontalRule],
images: &[[f64; 4]],
) -> (Option<tables::JudgedRegions>, Vec<Decision>) {
let refs: Vec<&Block> = page.iter().collect();
let texts: Vec<String> = page.iter().map(|b| b.text.clone()).collect();
let mut features = page_features(&refs, size.0, size.1);
facts.fill(&mut features, &texts);
let grouped = lines::grouped_lines(&features.words, page.len());
let number = page.first().map_or(0, |b| b.page);
let vectors = vectors_from_parts(&refs, number, &features, &grouped);
let rows = (switches.heading_lines
|| switches.list_item
|| switches.line_gnn
|| switches.captions
|| switches.table_judge)
.then(|| heading_lines::page_rows(&vectors));
let classes = (switches.list_item || switches.table_judge).then(|| {
heading_lines::class_probabilities_of_rows(
rows.as_ref().expect("rows are read wherever classes are"),
)
});
let judged = switches.table_judge.then(|| {
let signal = list_signal_from_parts(
&vectors,
classes
.as_ref()
.expect("classes are read wherever the signal is"),
);
let lines = judge::lines_of_grouped(&refs, &grouped, &signal);
tables::judge_regions_with_lines(page, size, rules, switches, &lines, &features.words)
});
let decisions = decide_page_on_shared(
&refs,
&vectors,
size,
switches,
images,
rows.as_ref(),
classes.as_ref(),
);
(judged, decisions)
}
#[cfg(test)]
fn judge_and_decide_legacy(
page: &[Block],
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
rules: &[crate::text::HorizontalRule],
images: &[[f64; 4]],
) -> (Option<tables::JudgedRegions>, Vec<Decision>) {
let judged = switches
.table_judge
.then(|| tables::judge_regions_traced(page, size, rules, facts, switches));
let decisions = {
let refs: Vec<&Block> = page.iter().collect();
decide_page(&refs, size, facts, switches, images)
};
(judged, decisions)
}
fn list_item_probability(v: &BlockVectors, classes: Option<&[[f64; 4]]>) -> f64 {
let mut x = [f64::NAN; 50];
x[..24].copy_from_slice(&v.block[..24]);
if let Some(first) = v.lines.first() {
x[24..47].copy_from_slice(first);
}
if let Some(p) = classes.filter(|c| !c.is_empty()) {
let list: Vec<f64> = p.iter().map(|c| c[2]).collect();
x[47] = list.iter().sum::<f64>() / list.len() as f64;
x[48] = list.iter().copied().fold(f64::NEG_INFINITY, f64::max);
x[49] = list[0];
}
probability(models::list_item::predict_margin(&x))
}
#[doc(hidden)]
pub fn line_class_probs(page: &[&Block], size: (f64, f64), facts: &DocumentFacts) -> Vec<[f32; 4]> {
heading_lines::page_class_probabilities(&page_vectors(page, size, facts))
.into_iter()
.flatten()
.map(|p| p.map(|v| v as f32))
.collect()
}
pub fn list_signal(page: &[&Block], size: (f64, f64), facts: &DocumentFacts) -> Vec<f32> {
let vectors = page_vectors(page, size, facts);
let classes = heading_lines::page_class_probabilities(&vectors);
list_signal_from_parts(&vectors, &classes)
}
fn list_signal_from_parts(vectors: &[BlockVectors], classes: &[Vec<[f64; 4]>]) -> Vec<f32> {
let mut out = Vec::new();
for (v, classes) in vectors.iter().zip(classes) {
if !matches!(v.kind, "paragraph" | "list_item") {
out.extend(std::iter::repeat_n(f32::NAN, v.lines.len()));
continue;
}
let item = list_item_probability(v, Some(classes));
let mut opened = 0.0f64;
for line in &v.lines {
opened = opened.max(probability(models::list_opens::predict_margin(line)));
out.push(item.max(opened) as f32);
}
}
out
}
fn hint_of(d: &Decision) -> tables::CellHint {
tables::CellHint {
heading: !d.heading_lines.is_empty() && d.heading_lines.iter().all(|h| *h),
p_heading: (!d.p_heading_lines.is_empty())
.then(|| d.p_heading_lines.iter().sum::<f64>() / d.p_heading_lines.len() as f64),
}
}
#[doc(hidden)]
pub fn cell_hints(
blocks: &[Block],
pages: &[(f64, f64)],
images: &[Vec<[f64; 4]>],
) -> Vec<tables::CellHint> {
let facts = document_facts(blocks);
let mut out = vec![tables::CellHint::default(); blocks.len()];
for (page, members) in pages_of(blocks).iter().enumerate() {
let refs: Vec<&Block> = members.iter().map(|&i| &blocks[i]).collect();
let size = pages.get(page).copied().unwrap_or((0.0, 0.0));
let pictures = images.get(page).map(Vec::as_slice).unwrap_or(&[]);
for (&i, d) in members
.iter()
.zip(decide_page(&refs, size, &facts, SWITCHES, pictures))
{
out[i] = hint_of(&d);
}
}
out
}
#[doc(hidden)]
#[derive(Debug, Clone, PartialEq)]
pub struct LineConfidence {
pub p_heading: Option<f64>,
pub class_probs: Vec<f32>,
pub p_link_next: f32,
}
#[doc(hidden)]
#[derive(Debug, Clone, PartialEq)]
pub struct PageConfidence {
pub lines: Vec<LineConfidence>,
pub doubt: f64,
pub heading_doubt: f64,
pub link_doubt: f64,
pub class_entropy: f64,
}
#[doc(hidden)]
pub fn page_confidence(
blocks: &[Block],
pages: &[(f64, f64)],
images: &[Vec<[f64; 4]>],
) -> Vec<PageConfidence> {
page_confidence_with_switches(blocks, pages, images, SWITCHES)
}
#[doc(hidden)]
pub fn page_confidence_with_switches(
blocks: &[Block],
pages: &[(f64, f64)],
images: &[Vec<[f64; 4]>],
switches: Switches,
) -> Vec<PageConfidence> {
let facts = document_facts(blocks);
let members = pages_of(blocks);
let count = members.len().max(pages.len());
let mut out = Vec::with_capacity(count);
for number in 0..count {
let on_page: Vec<&Block> = members
.get(number)
.map(|m| m.iter().map(|&i| &blocks[i]).collect())
.unwrap_or_default();
let size = pages.get(number).copied().unwrap_or((0.0, 0.0));
let pictures = images.get(number).map(Vec::as_slice).unwrap_or(&[]);
out.push(confidence_of_page(
&on_page, size, &facts, switches, pictures,
));
}
out
}
fn confidence_of_page(
page: &[&Block],
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
images: &[[f64; 4]],
) -> PageConfidence {
let (decisions, vectors, gnn) = decide_page_traced(page, size, facts, switches, images);
let mut lines = Vec::new();
let mut near = 0usize;
let mut decided = 0usize;
let mut entropy_sum = 0.0;
let mut boxes: Vec<[f64; 4]> = Vec::new();
for (decision, v) in decisions.iter().zip(&vectors) {
let bar = heading_bar(v.kind, switches);
for k in 0..v.lines.len() {
let p = decision.p_heading_lines.get(k).copied();
if let Some(p) = p {
decided += 1;
if (p - bar).abs() <= 0.15 {
near += 1;
}
}
let at = lines.len();
if let Some(g) = &gnn {
entropy_sum += line_gnn::normalised_entropy(&g.classes[at]);
}
lines.push(LineConfidence {
p_heading: p,
class_probs: gnn
.as_ref()
.map(|g| g.classes[at].clone())
.unwrap_or_default(),
p_link_next: gnn.as_ref().map(|g| g.link_next[at]).unwrap_or(f32::NAN),
});
boxes.push(v.line_boxes[k].0);
}
}
let heading_doubt = if decided == 0 {
0.0
} else {
near as f64 / decided as f64
};
let mut linked = 0usize;
let mut stacked = 0usize;
if let Some(g) = &gnn {
for i in 0..lines.len().saturating_sub(1) {
if boxes[i + 1][3] <= boxes[i][1] {
stacked += 1;
if (0.35..=0.65).contains(&f64::from(g.link_next[i])) {
linked += 1;
}
}
}
}
let (link_doubt, class_entropy, doubt) = match &gnn {
None => (f64::NAN, f64::NAN, heading_doubt),
Some(_) => {
let link = if stacked == 0 {
0.0
} else {
linked as f64 / stacked as f64
};
let entropy = if lines.is_empty() {
0.0
} else {
entropy_sum / lines.len() as f64
};
(link, entropy, heading_doubt.max(link).max(entropy))
}
};
PageConfidence {
lines,
doubt,
heading_doubt,
link_doubt,
class_entropy,
}
}
#[doc(hidden)]
pub fn final_tables_for_dump(
page: Vec<Block>,
size: (f64, f64),
facts: &DocumentFacts,
rules: &[crate::text::HorizontalRule],
rulings: &[crate::text::Ruling],
images: &[[f64; 4]],
) -> Vec<tables::FinalTableV2> {
final_tables_for_dump_with(page, size, facts, rules, rulings, images, SWITCHES)
}
#[doc(hidden)]
pub fn final_tables_for_dump_with(
page: Vec<Block>,
size: (f64, f64),
facts: &DocumentFacts,
rules: &[crate::text::HorizontalRule],
rulings: &[crate::text::Ruling],
images: &[[f64; 4]],
switches: Switches,
) -> Vec<tables::FinalTableV2> {
let (judged_before, decisions) =
judge_and_decide_shared(&page, size, facts, switches, rules, images);
let marked: Vec<bool> = page
.iter()
.map(|b| matches!(b.kind, BlockKind::ListItem { .. }) || leading_marker(&b.text).is_some())
.collect();
let (page, origin) = apply_traced(page, &decisions);
let (page, origin) = if switches.line_gnn && switches.list_item {
merge_items(
page,
origin,
&decisions,
&marked,
judged_before.as_ref().map(|j| j.regions.as_slice()),
)
} else {
(page, origin)
};
let judged: Option<Vec<Option<usize>>> = judged_before
.as_ref()
.map(|j| origin.iter().map(|&o| j.regions[o]).collect());
let trace = judged_before.as_ref().and_then(|j| j.trace.clone());
let meta = trace
.as_ref()
.map(tables::JudgeRegionMeta::of_trace)
.unwrap_or_default();
let builder_p: Vec<Option<f64>> = trace.as_ref().map_or_else(
|| vec![None; origin.len()],
|t| {
origin
.iter()
.map(|&o| t.builder_p.get(o).copied().flatten())
.collect()
},
);
let lattices = if switches.table_lattice {
lattice::lattices(rulings)
} else {
Vec::new()
};
let hints: Vec<tables::CellHint> = origin.iter().map(|&o| hint_of(&decisions[o])).collect();
let switches = Switches {
table_is_v2: true,
..switches
};
let drawn = !rules.is_empty() || !rulings.is_empty();
let (out, mine, _) = tables::refine_traced(
page,
size,
facts,
switches,
judged.as_deref(),
trace.as_ref(),
&lattices,
&hints,
drawn,
None,
true,
);
let finals = tables::final_tables_for_is_table_v2(
&out,
&mine,
judged.as_deref(),
&meta,
&builder_p,
&hints,
rules,
rulings,
size.1,
);
if switches.table_split_stacked {
tables::split_final_tables_stacked(out, finals).1
} else {
finals
}
}
fn decide(blocks: &[Block], pages: &[(f64, f64)], switches: Switches) -> Vec<Decision> {
let facts = document_facts(blocks);
let mut out = vec![Decision::default(); blocks.len()];
for (page, members) in pages_of(blocks).iter().enumerate() {
let refs: Vec<&Block> = members.iter().map(|&i| &blocks[i]).collect();
let size = pages.get(page).copied().unwrap_or((0.0, 0.0));
for (&i, d) in members
.iter()
.zip(decide_page(&refs, size, &facts, switches, &[]))
{
out[i] = d;
}
}
out
}
fn demote_confirmed_runs(
members: &[usize],
page: &[&Block],
vectors: &[BlockVectors],
decisions: &mut [Decision],
) {
let mut at = 0;
while at < members.len() {
if vectors[members[at]].kind != "list_item" {
at += 1;
continue;
}
let mut end = at;
while end < members.len() && vectors[members[end]].kind == "list_item" {
end += 1;
}
let run = &members[at..end];
let scored: Vec<f64> = run.iter().filter_map(|&i| decisions[i].p_item).collect();
let low = !scored.is_empty()
&& scored.iter().sum::<f64>() / (scored.len() as f64) < LIST_ITEM_DEMOTE_BELOW;
let counts = run.len() >= 2
&& run
.windows(2)
.all(|w| match (&page[w[0]].kind, &page[w[1]].kind) {
(
BlockKind::ListItem {
ordinal: Some(a), ..
},
BlockKind::ListItem {
ordinal: Some(b), ..
},
) => *b == a + 1,
_ => false,
});
let worded = run.iter().all(|&i| {
page[i]
.text
.split_whitespace()
.any(|w| w.chars().filter(|c| c.is_alphabetic()).count() >= 3)
});
let low = low && !(counts && worded);
if low {
for &i in run {
if decisions[i].p_item.is_some() {
decisions[i].item = Some(false);
}
}
}
at = end;
}
}
pub fn refine_page(
page: Vec<Block>,
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
) -> Vec<Block> {
refine_page_traced(page, size, facts, switches, &[], &[], &[], None, true).0
}
const LIST_MERGE_AT: f64 = 0.6;
fn under_the_item(item: &Block, below: &Block) -> bool {
let em = crate::blocks::dominant_font_size(item.spans.iter()).max(1.0);
let [left, _, right, _] = item.bbox;
let [b_left, _, b_right, _] = below.bbox;
let last = item.spans.iter().map(|s| s.y).fold(f64::INFINITY, f64::min);
let first = below
.spans
.iter()
.map(|s| s.y)
.fold(f64::NEG_INFINITY, f64::max);
let step = last - first;
let hanging = b_left >= left + MERGE_MIN_HANG_EM * em
|| (b_left > left + 0.3 * em && (b_left - text_left(item, em)).abs() <= 0.35 * em);
let flush = (b_left - left).abs() <= FLUSH_EM * em
&& right - left >= FLUSH_MIN_MEASURE_EM * em
&& b_right <= right + 0.5 * em;
step > 0.0
&& step <= MERGE_MAX_GAP_EM * em
&& (hanging || flush)
&& b_left < right
&& right >= b_right - 3.0 * em
}
fn text_left(item: &Block, em: f64) -> f64 {
let top = item
.spans
.iter()
.map(|s| s.y)
.fold(f64::NEG_INFINITY, f64::max);
let mut line: Vec<&crate::text::TextSpan> = item
.spans
.iter()
.filter(|s| (s.y - top).abs() <= 0.5 * em && !s.text.trim().is_empty())
.collect();
line.sort_by(|a, b| a.x.total_cmp(&b.x));
let Some(first) = line.first() else {
return item.bbox[0];
};
let trimmed = first.text.trim_start();
let glyph = trimmed.trim_end();
if !glyph.is_empty() && glyph.chars().count() <= 2 && !glyph.chars().any(char::is_alphanumeric)
{
return line.get(1).map_or(first.x, |s| s.x);
}
let lead = first.text.chars().count() - trimmed.chars().count();
let Some((len, _)) = leading_marker(trimmed) else {
return first.x;
};
let marker_chars = lead + trimmed[..len.min(trimmed.len())].chars().count();
if marker_chars < first.text.chars().count() {
first
.char_bounds
.get(marker_chars)
.map_or(first.x, |b| b[0])
} else {
line.get(1).map_or(first.x, |s| s.x)
}
}
const MERGE_MAX_GAP_EM: f64 = 2.2;
const MERGE_MIN_HANG_EM: f64 = 0.8;
const FLUSH_EM: f64 = 0.35;
const FLUSH_MIN_MEASURE_EM: f64 = 18.0;
fn merge_items(
out: Vec<Block>,
origin: Vec<usize>,
decisions: &[Decision],
marked: &[bool],
judged: Option<&[Option<usize>]>,
) -> (Vec<Block>, Vec<usize>) {
let mut blocks: Vec<Block> = Vec::with_capacity(out.len());
let mut from: Vec<usize> = Vec::with_capacity(out.len());
let mut reach: Vec<usize> = Vec::with_capacity(out.len());
for (block, o) in out.into_iter().zip(origin) {
let joins = match (blocks.last(), reach.last()) {
(Some(item), Some(&upto)) => {
matches!(item.kind, BlockKind::ListItem { .. })
&& matches!(
block.kind,
BlockKind::ListItem { .. } | BlockKind::Paragraph
)
&& o == upto + 1
&& !marked.get(o).copied().unwrap_or(true)
&& !(matches!(block.kind, BlockKind::Paragraph)
&& judged.is_some_and(|j| {
j.get(o)
.copied()
.flatten()
.is_some_and(|k| k != tables::REJECTED)
}))
&& item.page == block.page
&& under_the_item(item, &block)
&& decisions
.get(upto)
.and_then(|d| d.link_after)
.is_some_and(|p| p >= LIST_MERGE_AT)
}
_ => false,
};
if joins {
let item = blocks.last_mut().expect("a block to join");
item.text = crate::blocks::join_lines(&[item.text.clone(), block.text.clone()]);
item.spans.extend(block.spans);
crate::blocks::recompute_bbox(item);
*reach.last_mut().expect("a reach") = o;
} else {
blocks.push(block);
from.push(o);
reach.push(o);
}
}
(blocks, from)
}
#[allow(clippy::too_many_arguments)]
fn refine_page_traced(
page: Vec<Block>,
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
rules: &[crate::text::HorizontalRule],
rulings: &[crate::text::Ruling],
images: &[[f64; 4]],
tatr: Option<&tables::TatrContext>,
cell_path_words: bool,
) -> (Vec<Block>, Vec<usize>, Option<tables::PageTableSignals>) {
let (judged, decisions) = judge_and_decide_shared(&page, size, facts, switches, rules, images);
let marked: Vec<bool> = page
.iter()
.map(|b| matches!(b.kind, BlockKind::ListItem { .. }) || leading_marker(&b.text).is_some())
.collect();
let (page, origin) = apply_traced(page, &decisions);
let (page, origin) = if switches.line_gnn && switches.list_item {
merge_items(
page,
origin,
&decisions,
&marked,
judged.as_ref().map(|j| j.regions.as_slice()),
)
} else {
(page, origin)
};
if !switches.table_region
&& !switches.table_structure
&& judged.is_none()
&& !switches.table_lattice
{
return (page, origin, None);
}
let trace = judged.as_ref().and_then(|j| j.trace.clone());
let judged: Option<Vec<Option<usize>>> =
judged.map(|j| origin.iter().map(|&o| j.regions[o]).collect());
let lattices = if switches.table_lattice {
lattice::lattices(rulings)
} else {
Vec::new()
};
let hints: Vec<tables::CellHint> = origin.iter().map(|&o| hint_of(&decisions[o])).collect();
let drawn = !rules.is_empty() || !rulings.is_empty();
let (out, mine, mut signals) = tables::refine_traced(
page,
size,
facts,
switches,
judged.as_deref(),
trace.as_ref(),
&lattices,
&hints,
drawn,
tatr,
cell_path_words,
);
let (out, mine, verdicts) = if switches.table_is_v2 {
let meta = trace
.as_ref()
.map(tables::JudgeRegionMeta::of_trace)
.unwrap_or_default();
let builder_p: Vec<Option<f64>> = trace.as_ref().map_or_else(
|| vec![None; hints.len()],
|t| {
origin
.iter()
.map(|&o| t.builder_p.get(o).copied().flatten())
.collect()
},
);
let (out, mine, verdicts, map) = tables::apply_table_is_v2_traced(
out,
mine,
judged.as_deref(),
&meta,
&builder_p,
&hints,
rules,
rulings,
size.1,
switches,
);
if let Some(signals) = signals.as_mut() {
tables::remap_table_signal_ranges(signals, &map, &out);
}
(out, mine, Some(verdicts))
} else {
(out, mine, None)
};
let (out, mine, verdicts, mut signals) = if switches.table_split_stacked {
tables::split_stacked_after_v2(out, mine, verdicts, signals)
} else {
(out, mine, verdicts, signals)
};
if let (Some(signals), Some(verdicts)) = (signals.as_mut(), verdicts.as_ref()) {
tables::fill_is_table_v2_signals(signals, verdicts);
}
let origin = mine.into_iter().map(|i| origin[i]).collect();
(out, origin, signals)
}
#[doc(hidden)]
pub fn refine_with(blocks: Vec<Block>, pages: &[(f64, f64)], switches: Switches) -> Vec<Block> {
refine_with_rules(blocks, pages, &[], switches)
}
#[doc(hidden)]
pub fn refine_with_rules(
blocks: Vec<Block>,
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
switches: Switches,
) -> Vec<Block> {
refine_with_rulings(blocks, pages, rules, &[], switches)
}
#[doc(hidden)]
pub fn refine_with_rulings(
blocks: Vec<Block>,
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
rulings: &[Vec<crate::text::Ruling>],
switches: Switches,
) -> Vec<Block> {
refine_with_images(blocks, pages, rules, rulings, &[], switches)
}
#[doc(hidden)]
pub fn refine_with_images(
blocks: Vec<Block>,
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
rulings: &[Vec<crate::text::Ruling>],
images: &[Vec<[f64; 4]>],
switches: Switches,
) -> Vec<Block> {
refine_with_tatr(blocks, pages, rules, rulings, images, switches, None, true)
}
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn refine_with_tatr(
blocks: Vec<Block>,
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
rulings: &[Vec<crate::text::Ruling>],
images: &[Vec<[f64; 4]>],
switches: Switches,
tatr: Option<&tables::TatrContext>,
cell_path_words: bool,
) -> Vec<Block> {
let facts = document_facts(&blocks);
let members = pages_of(&blocks);
let mut slots: Vec<Option<Block>> = blocks.into_iter().map(Some).collect();
let mut placed: Vec<Vec<Block>> = (0..slots.len()).map(|_| Vec::new()).collect();
for (number, members) in members.iter().enumerate() {
let page: Vec<Block> = members
.iter()
.map(|&i| slots[i].take().expect("a block is on one page"))
.collect();
let size = pages.get(number).copied().unwrap_or((0.0, 0.0));
let drawn = rules.get(number).map(Vec::as_slice).unwrap_or(&[]);
let lines = rulings.get(number).map(Vec::as_slice).unwrap_or(&[]);
let pictures = images.get(number).map(Vec::as_slice).unwrap_or(&[]);
let (out, origin, _) = refine_page_traced(
page,
size,
&facts,
switches,
drawn,
lines,
pictures,
tatr,
cell_path_words,
);
for (block, from) in out.into_iter().zip(origin) {
placed[members[from]].push(block);
}
}
let mut out: Vec<Block> = placed.into_iter().flatten().collect();
for (i, block) in out.iter_mut().enumerate() {
block.index = i;
if !matches!(block.kind, BlockKind::TableCell { .. }) {
block.grid = None;
}
}
out
}
#[doc(hidden)]
pub fn table_signals_with(
blocks: &[Block],
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
rulings: &[Vec<crate::text::Ruling>],
images: &[Vec<[f64; 4]>],
switches: Switches,
) -> Vec<Option<tables::PageTableSignals>> {
let facts = document_facts(blocks);
let members = pages_of(blocks);
(0..pages.len())
.map(|number| {
let page: Vec<Block> = members
.get(number)
.map(|members| members.iter().map(|&i| blocks[i].clone()).collect())
.unwrap_or_default();
let size = pages.get(number).copied().unwrap_or((0.0, 0.0));
let drawn = rules.get(number).map(Vec::as_slice).unwrap_or(&[]);
let lines = rulings.get(number).map(Vec::as_slice).unwrap_or(&[]);
let pictures = images.get(number).map(Vec::as_slice).unwrap_or(&[]);
refine_page_traced(
page, size, &facts, switches, drawn, lines, pictures, None, true,
)
.2
})
.collect()
}
#[doc(hidden)]
pub fn refine_page_with_signals(
page: Vec<Block>,
size: (f64, f64),
facts: &DocumentFacts,
switches: Switches,
rules: &[crate::text::HorizontalRule],
rulings: &[crate::text::Ruling],
images: &[[f64; 4]],
) -> (Vec<Block>, Option<tables::PageTableSignals>) {
let (blocks, _, signals) = refine_page_traced(
page, size, facts, switches, rules, rulings, images, None, true,
);
(blocks, signals)
}
#[doc(hidden)]
pub fn refine_with_images_and_signals(
blocks: Vec<Block>,
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
rulings: &[Vec<crate::text::Ruling>],
images: &[Vec<[f64; 4]>],
switches: Switches,
) -> (Vec<Block>, Vec<Option<tables::PageTableSignals>>) {
refine_with_images_and_signals_with_cell_path_words(
blocks, pages, rules, rulings, images, switches, true,
)
}
#[doc(hidden)]
pub fn refine_with_images_and_signals_with_cell_path_words(
blocks: Vec<Block>,
pages: &[(f64, f64)],
rules: &[Vec<crate::text::HorizontalRule>],
rulings: &[Vec<crate::text::Ruling>],
images: &[Vec<[f64; 4]>],
switches: Switches,
cell_path_words: bool,
) -> (Vec<Block>, Vec<Option<tables::PageTableSignals>>) {
let facts = document_facts(&blocks);
let members = pages_of(&blocks);
let mut slots: Vec<Option<Block>> = blocks.into_iter().map(Some).collect();
let mut placed: Vec<Vec<Block>> = (0..slots.len()).map(|_| Vec::new()).collect();
let mut signals: Vec<Option<tables::PageTableSignals>> = Vec::with_capacity(pages.len());
for number in 0..members.len().max(pages.len()) {
let page: Vec<Block> = members
.get(number)
.map(|members| {
members
.iter()
.map(|&i| slots[i].take().expect("a block is on one page"))
.collect()
})
.unwrap_or_default();
let size = pages.get(number).copied().unwrap_or((0.0, 0.0));
let drawn = rules.get(number).map(Vec::as_slice).unwrap_or(&[]);
let lines = rulings.get(number).map(Vec::as_slice).unwrap_or(&[]);
let pictures = images.get(number).map(Vec::as_slice).unwrap_or(&[]);
let (out, origin, page_signals) = refine_page_traced(
page,
size,
&facts,
switches,
drawn,
lines,
pictures,
None,
cell_path_words,
);
if number < pages.len() {
signals.push(page_signals);
}
if let Some(members) = members.get(number) {
for (block, from) in out.into_iter().zip(origin) {
placed[members[from]].push(block);
}
}
}
let mut out: Vec<Block> = placed.into_iter().flatten().collect();
for (i, block) in out.iter_mut().enumerate() {
block.index = i;
if !matches!(block.kind, BlockKind::TableCell { .. }) {
block.grid = None;
}
}
(out, signals)
}
#[doc(hidden)]
pub fn page_blocks(blocks: &[Block], number: usize) -> Vec<Block> {
blocks
.iter()
.filter(|b| b.page == number)
.cloned()
.collect()
}
fn leading_marker(text: &str) -> Option<(usize, Option<u32>)> {
let trimmed = text.trim_start();
let lead = text.len() - trimmed.len();
let token = trimmed.split_whitespace().next()?;
let bullet = token.chars().count() <= 2
&& token.chars().next().is_some_and(|c| {
"\u{2022}\u{2023}\u{2043}\u{2219}\u{b7}\u{25aa}\u{25ab}\u{25a0}\u{25a1}\u{25cf}\u{25cb}\u{25e6}\u{25c6}\u{25c7}\u{2756}\u{2713}\u{2714}".contains(c)
});
let bullet = bullet || matches!(token, "-" | "\u{2013}" | "\u{2014}" | "*");
let glued = token
.chars()
.next()
.filter(|&c| crate::lists::is_a_pointer_or_a_picture(c));
if let Some(c) = glued {
let rest = &trimmed[c.len_utf8()..];
if crate::lists::opens_like_an_item(c, rest) {
let after = rest.len() - rest.trim_start().len();
return Some((lead + c.len_utf8() + after, None));
}
}
let digits: String = token
.trim_start_matches(['(', '['])
.trim_end_matches(['.', ')', ']'])
.to_string();
let ordered = digits != token && !digits.is_empty() && digits.chars().count() <= 3;
if !bullet && !ordered {
return None;
}
let rest = &trimmed[token.len()..];
if ordered
&& !token.starts_with(['(', '['])
&& digits.len() == 1
&& digits.chars().all(|c| c.is_ascii_uppercase())
&& crate::lists::opens_with_an_initial(rest)
{
return None;
}
let after = rest.len() - rest.trim_start().len();
let ordinal = if ordered {
digits.parse::<u32>().ok().or_else(|| {
let mut cs = digits.chars();
match (cs.next(), cs.next()) {
(Some(c), None) if c.is_ascii_alphabetic() => {
Some(u32::from(c.to_ascii_lowercase() as u8 - b'a' + 1))
}
_ => crate::lists::roman_value(&digits.to_ascii_lowercase()),
}
})
} else {
None
};
Some((lead + token.len() + after, ordinal))
}
fn text_of(lines: &[Vec<&features::PageWord>]) -> String {
let texts: Vec<String> = lines
.iter()
.map(|l| {
l.iter()
.map(|w| w.text.as_str())
.collect::<Vec<_>>()
.join(" ")
})
.collect();
crate::blocks::join_lines(&texts)
}
fn text_of_spans(spans: &[crate::text::TextSpan]) -> String {
let mut rows: Vec<Vec<crate::text::TextSpan>> = Vec::new();
for span in spans {
match rows.last_mut() {
Some(row)
if (row[0].y - span.y).abs() <= row[0].font_size.max(span.font_size) * 0.5 =>
{
row.push(span.clone())
}
_ => rows.push(vec![span.clone()]),
}
}
let texts: Vec<String> = rows
.into_iter()
.map(|spans| crate::text::TextBlock { spans }.text())
.collect();
crate::blocks::join_lines(&texts)
}
fn bbox_of(spans: &[crate::text::TextSpan]) -> Option<[f64; 4]> {
let mut out: Option<[f64; 4]> = None;
for s in spans {
let boxes: Vec<[f64; 4]> = if s.char_bounds.is_empty() {
vec![[s.x, s.y, s.x + s.width, s.y + s.height]]
} else {
s.char_bounds.clone()
};
for b in boxes {
out = Some(match out {
None => b,
Some(o) => [
o[0].min(b[0]),
o[1].min(b[1]),
o[2].max(b[2]),
o[3].max(b[3]),
],
});
}
}
out
}
fn by_heading_lines(block: Block, decision: &Decision) -> Vec<Block> {
let flags = &decision.heading_lines;
let p = &decision.p_heading_lines;
let mean = |from: usize, to: usize| -> Option<f64> {
let part = p.get(from..to)?;
(!part.is_empty()).then(|| part.iter().sum::<f64>() / part.len() as f64)
};
let body = |kind: BlockKind| match kind {
BlockKind::Heading { .. } => BlockKind::Paragraph,
other => other,
};
if matches!(
block.kind,
BlockKind::TableCell { header: true, .. } | BlockKind::TableCell { row: 1.., .. }
) {
return vec![block];
}
let breaks = &decision.line_breaks;
if flags.iter().all(|f| *f) && breaks.is_empty() {
let mut block = block;
if !matches!(block.kind, BlockKind::Heading { .. }) {
block.kind = BlockKind::Heading { level: 1 };
block.confidence = mean(0, p.len());
}
return vec![block];
}
if (!flags.iter().any(|f| *f) && breaks.is_empty())
|| matches!(block.kind, BlockKind::TableCell { .. })
{
let mut block = block;
if matches!(block.kind, BlockKind::Heading { .. }) {
block.kind = BlockKind::Paragraph;
block.confidence = mean(0, p.len());
}
return vec![block];
}
let cuts: Vec<usize> = (1..flags.len())
.filter(|&k| flags[k] != flags[k - 1] || breaks.contains(&k))
.collect();
let mut bounds = vec![0];
bounds.extend(&cuts);
bounds.push(flags.len());
let parts = if every_part_has_text(&block, &cuts) {
parts_of(&block, &cuts)
} else {
Vec::new()
};
if parts.len() != cuts.len() + 1 {
return vec![block];
}
let kind = block.kind;
parts
.into_iter()
.enumerate()
.map(|(r, mut piece)| {
piece.kind = if flags[bounds[r]] {
BlockKind::Heading { level: 1 }
} else {
body(kind)
};
piece.confidence = mean(bounds[r], bounds[r + 1]);
piece
})
.collect()
}
fn by_caption_lines(block: Block, decision: &Decision) -> Vec<Block> {
let flags = &decision.caption_lines;
let p = &decision.p_caption_lines;
let mean = |from: usize, to: usize| -> Option<f64> {
let part = p.get(from..to)?;
(!part.is_empty()).then(|| part.iter().sum::<f64>() / part.len() as f64)
};
if flags.iter().all(|f| *f) {
let mut block = block;
block.kind = BlockKind::Caption;
block.confidence = mean(0, p.len());
return vec![block];
}
let cuts: Vec<usize> = (1..flags.len())
.filter(|&k| flags[k] != flags[k - 1])
.collect();
let mut bounds = vec![0];
bounds.extend(&cuts);
bounds.push(flags.len());
let parts = if every_part_has_text(&block, &cuts) {
parts_of(&block, &cuts)
} else {
Vec::new()
};
if parts.len() != cuts.len() + 1 {
return vec![block];
}
parts
.into_iter()
.enumerate()
.map(|(r, mut piece)| {
if flags[bounds[r]] {
piece.kind = BlockKind::Caption;
piece.confidence = mean(bounds[r], bounds[r + 1]);
}
piece
})
.collect()
}
fn gnn_of(switches: Switches) -> &'static line_gnn::Net {
if switches.line_gnn_v2 {
&line_gnn::V2
} else {
&line_gnn::V1
}
}
fn as_item(mut block: Block, text: String) -> Block {
let (strip, ordinal) = leading_marker(&text).unwrap_or((0, None));
block.text = text[strip..].trim().to_string();
block.kind = BlockKind::ListItem { level: 0, ordinal };
block
}
const INDENT_MATCH_EM: f64 = 0.35;
fn size_of(block: &Block) -> f64 {
let mut sizes: Vec<f64> = block.spans.iter().map(|s| s.font_size).collect();
sizes.sort_by(|a, b| a.total_cmp(b));
sizes
.get(sizes.len() / 2)
.copied()
.filter(|s| *s > 0.0)
.unwrap_or(10.0)
}
fn promoted_level(item: &Block, others: &[&Block]) -> u8 {
let tolerance = INDENT_MATCH_EM * size_of(item);
let same_column = |b: &Block| {
let overlap = b.bbox[2].min(item.bbox[2]) - b.bbox[0].max(item.bbox[0]);
b.page == item.page && overlap > 0.0
};
let mut edges: Vec<f64> = others
.iter()
.filter(|b| matches!(b.kind, BlockKind::ListItem { .. }) && same_column(b))
.map(|b| b.bbox[0])
.chain(std::iter::once(item.bbox[0]))
.collect();
edges.sort_by(|a, b| a.total_cmp(b));
let mut steps: Vec<f64> = Vec::new();
for edge in edges {
if steps.last().is_none_or(|last| edge - last > tolerance) {
steps.push(edge);
}
}
let edge = item.bbox[0];
steps
.iter()
.enumerate()
.min_by(|(_, a), (_, b)| (edge - **a).abs().total_cmp(&(edge - **b).abs()))
.map_or(0, |(rank, _)| rank)
.min(crate::lists::MAX_LIST_LEVEL as usize) as u8
}
#[cfg(test)]
fn apply(blocks: Vec<Block>, decisions: &[Decision]) -> Vec<Block> {
apply_traced(blocks, decisions).0
}
fn apply_traced(blocks: Vec<Block>, decisions: &[Decision]) -> (Vec<Block>, Vec<usize>) {
let mut out: Vec<Block> = Vec::with_capacity(blocks.len());
let mut origin: Vec<usize> = Vec::with_capacity(blocks.len());
let mut promoted: Vec<usize> = Vec::new();
for (k, (block, decision)) in blocks.into_iter().zip(decisions).enumerate() {
let before = out.len();
let traced = |out: &Vec<Block>, origin: &mut Vec<usize>| origin.resize(out.len(), k);
debug_assert_eq!(origin.len(), before);
let heading_lines_act = decision.heading_lines.iter().any(|h| *h)
|| !decision.line_breaks.is_empty()
|| (!decision.heading_lines.is_empty()
&& matches!(block.kind, BlockKind::Heading { .. }));
if heading_lines_act && decision.opens_at.is_empty() {
out.extend(by_heading_lines(block, decision));
traced(&out, &mut origin);
continue;
}
if decision.caption_lines.iter().any(|c| *c)
&& decision.opens_at.is_empty()
&& matches!(block.kind, BlockKind::Paragraph)
{
out.extend(by_caption_lines(block, decision));
traced(&out, &mut origin);
continue;
}
if decision.opens_at.is_empty() || !every_part_has_text(&block, &decision.opens_at) {
match (decision.item, &block.kind) {
(Some(true), BlockKind::Paragraph) => {
let text = block.text.clone();
let mut item = as_item(block, text);
item.confidence = decision.p_item;
promoted.push(out.len());
out.push(item);
}
(Some(false), BlockKind::ListItem { .. }) => {
let mut block = block;
block.text = with_marker_back(&block);
block.kind = BlockKind::Paragraph;
block.confidence = decision.p_item;
out.push(block);
}
_ => out.push(block),
}
traced(&out, &mut origin);
continue;
}
let nested = match block.kind {
BlockKind::ListItem { level, .. } => Some(level),
_ => None,
};
for mut piece in split(block, decision) {
if let BlockKind::ListItem { ordinal, .. } = piece.kind {
match nested {
Some(level) => piece.kind = BlockKind::ListItem { level, ordinal },
None => promoted.push(out.len()),
}
}
out.push(piece);
}
traced(&out, &mut origin);
}
let levels: Vec<(usize, u8)> = promoted
.iter()
.map(|&i| {
let others: Vec<&Block> = out
.iter()
.enumerate()
.filter(|(j, _)| *j != i)
.map(|(_, b)| b)
.collect();
(i, promoted_level(&out[i], &others))
})
.collect();
for (i, level) in levels {
if let BlockKind::ListItem { ordinal, .. } = out[i].kind {
out[i].kind = BlockKind::ListItem { level, ordinal };
}
}
for (i, block) in out.iter_mut().enumerate() {
block.index = i;
}
(out, origin)
}
fn with_marker_back(block: &Block) -> String {
let words = words_of_block(block);
let lines = lines_of(0, &words);
let Some(marker) = lines
.first()
.and_then(|l| l.first())
.map(|w| w.text.as_str())
else {
return block.text.clone();
};
if block.text.starts_with(marker) {
return block.text.clone();
}
format!("{marker} {}", block.text)
}
fn words_of_block(block: &Block) -> Vec<features::PageWord> {
let mut words = Vec::new();
for (s, span) in block.spans.iter().enumerate() {
for w in features::words_of(span) {
words.push(features::PageWord {
text: w.text,
bbox: w.bbox,
block: 0,
span: s,
size: span.font_size,
bold: span.is_bold,
italic: span.is_italic,
font: span.font_name.clone(),
color: span.color.map(|[r, g, b, _]| [r, g, b]),
});
}
}
words
}
fn every_part_has_text(block: &Block, opens_at: &[usize]) -> bool {
let words = words_of_block(block);
let lines = lines_of(0, &words);
let mut parts: Vec<Vec<Vec<&features::PageWord>>> = Vec::new();
for (k, line) in lines.iter().enumerate() {
if parts.is_empty() || opens_at.contains(&k) {
parts.push(Vec::new());
}
if let Some(part) = parts.last_mut() {
part.push(line.clone());
}
}
parts.iter().all(|part| {
let text = text_of(part);
let strip = leading_marker(&text).map_or(0, |(strip, _)| strip);
!text[strip..].trim().is_empty()
})
}
fn parts_of(block: &Block, cuts: &[usize]) -> Vec<Block> {
let words = words_of_block(block);
let lines = lines_of(0, &words);
let mut line_of_span: Vec<Option<usize>> = vec![None; block.spans.len()];
for (k, line) in lines.iter().enumerate() {
for w in line {
if line_of_span[w.span].is_none_or(|l| k < l) {
line_of_span[w.span] = Some(k);
}
}
}
let mut cuts: Vec<usize> = cuts.to_vec();
cuts.sort_unstable();
cuts.dedup();
let part_of_line = |k: usize| cuts.iter().filter(|&&c| c <= k).count();
let mut spans: Vec<Vec<crate::text::TextSpan>> = vec![Vec::new(); cuts.len() + 1];
let mut last = 0usize;
for (s, span) in block.spans.iter().enumerate() {
let part = line_of_span[s].map_or(last, part_of_line);
last = part;
spans[part].push(span.clone());
}
let mut out = Vec::new();
for (p, part_spans) in spans.into_iter().enumerate() {
if part_spans.is_empty() {
continue;
}
let mut piece = block.clone();
piece.bbox = bbox_of(&part_spans).unwrap_or(block.bbox);
piece.text = text_of_spans(&part_spans);
piece.spans = part_spans;
if p > 0 {
piece.continues_previous = false;
}
out.push(piece);
}
out
}
fn split(block: Block, decision: &Decision) -> Vec<Block> {
let words = words_of_block(&block);
let lines = lines_of(0, &words);
let mut line_of_span: Vec<Option<usize>> = vec![None; block.spans.len()];
for (k, line) in lines.iter().enumerate() {
for w in line {
if line_of_span[w.span].is_none_or(|l| k < l) {
line_of_span[w.span] = Some(k);
}
}
}
let mut cuts = decision.opens_at.clone();
cuts.sort_unstable();
cuts.dedup();
let part_of_line = |k: usize| cuts.iter().filter(|&&c| c <= k).count();
let parts = cuts.len() + 1;
let mut spans: Vec<Vec<crate::text::TextSpan>> = vec![Vec::new(); parts];
let mut last = 0usize;
for (s, span) in block.spans.iter().enumerate() {
let part = line_of_span[s].map_or(last, part_of_line);
last = part;
spans[part].push(span.clone());
}
let mut out = Vec::with_capacity(parts);
for (p, part_spans) in spans.into_iter().enumerate() {
if part_spans.is_empty() {
continue;
}
let mut piece = block.clone();
piece.bbox = bbox_of(&part_spans).unwrap_or(block.bbox);
let text = text_of_spans(&part_spans);
piece.spans = part_spans;
piece.confidence = None;
if p == 0 {
piece.text = text;
piece = match (decision.item, &block.kind) {
_ if decision.first_opens => {
let t = piece.text.clone();
let mut item = as_item(piece, t);
item.confidence = decision.p_first_opens;
item
}
(Some(true), _) | (None, BlockKind::ListItem { .. }) => {
let t = piece.text.clone();
let mut item = as_item(piece, t);
item.confidence = decision.p_item;
item
}
_ => {
piece.kind = BlockKind::Paragraph;
piece
}
};
} else {
piece.continues_previous = false;
piece = as_item(piece, text);
piece.confidence = cuts
.get(p - 1)
.and_then(|cut| cut.checked_sub(1))
.and_then(|at| decision.p_opens.get(at).copied());
}
out.push(piece);
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::text::TextSpan;
fn line(text: &str, x: f64, y: f64) -> TextSpan {
let glyphs = text
.chars()
.enumerate()
.map(|(i, _)| {
let x0 = x + i as f64 * 5.0;
[x0, y, x0 + 5.0, y + 10.0]
})
.collect();
TextSpan {
text: text.to_string(),
x,
y,
width: text.chars().count() as f64 * 5.0,
height: 10.0,
font_size: 10.0,
char_bounds: glyphs,
..Default::default()
}
}
#[test]
fn a_list_item_cut_by_the_segmentation_is_put_back() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let item = BlockKind::ListItem {
level: 0,
ordinal: Some(5),
};
let out = vec![
block(
item,
"S.-H. Park, O. Simeone, and S. Shamai, Robust compression for cloud radio access",
72.0,
700.0,
),
block(
BlockKind::Paragraph,
"networks, IEEE Trans. Veh. Techn., vol. 62, pp. 692-703, 2013.",
84.0,
688.0,
),
block(
item,
"J. Hoydis, M. Kobayashi, and M. Debbah, Optimal channel training",
72.0,
676.0,
),
];
let decisions = |p: f64| {
vec![
Decision {
link_after: Some(p),
..Decision::default()
},
Decision {
link_after: Some(0.9),
..Decision::default()
},
Decision::default(),
]
};
let marked = [true, false, true];
let (joined, origin) =
merge_items(out.clone(), vec![0, 1, 2], &decisions(0.9), &marked, None);
assert_eq!(
joined.len(),
2,
"{:?}",
joined.iter().map(|b| &b.text).collect::<Vec<_>>()
);
assert!(
joined[0].text.ends_with("pp. 692-703, 2013."),
"{}",
joined[0].text
);
assert!(joined[1].text.starts_with("J. Hoydis"));
assert_eq!(origin, vec![0, 2]);
let (apart, _) = merge_items(out, vec![0, 1, 2], &decisions(0.3), &marked, None);
assert_eq!(apart.len(), 3);
}
#[test]
fn a_continuation_under_a_thin_bullets_text_joins_its_item() {
let mut item = Block::new(
BlockKind::ListItem {
level: 0,
ordinal: None,
},
0,
[0.0; 4],
"We were pleased that our market share gains continued".to_string(),
0,
);
item.spans = vec![
line(">", 72.0, 700.0),
line(
"We were pleased that our market share gains continued",
78.0,
700.0,
),
];
crate::blocks::recompute_bbox(&mut item);
let mut next = Block::new(
BlockKind::Paragraph,
0,
[0.0; 4],
"in the equity and fixed income markets.".to_string(),
0,
);
next.spans = vec![line("in the equity and fixed income markets.", 78.0, 688.0)];
crate::blocks::recompute_bbox(&mut next);
let decisions = vec![
Decision {
link_after: Some(0.99),
..Decision::default()
},
Decision::default(),
];
let (joined, _) = merge_items(
vec![item, next],
vec![0, 1],
&decisions,
&[true, false],
None,
);
assert_eq!(
joined.len(),
1,
"{:?}",
joined.iter().map(|b| &b.text).collect::<Vec<_>>()
);
assert!(joined[0].text.ends_with("fixed income markets."));
}
#[test]
fn a_flush_continuation_joins_its_item_where_the_link_head_says_so() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let item = block(
BlockKind::ListItem {
level: 0,
ordinal: Some(9),
},
"Summermeyer, D et al: Chimeric antigen receptor-modified T cells derived from defined CD44",
72.0,
700.0,
);
let next = block(
BlockKind::Paragraph,
"and CD8+ subsets confer superior antitumor reactivity in vivo.",
72.0,
688.0,
);
let decisions = vec![
Decision {
link_after: Some(0.99),
..Decision::default()
},
Decision::default(),
];
let (joined, _) = merge_items(
vec![item, next],
vec![0, 1],
&decisions,
&[true, false],
None,
);
assert_eq!(
joined.len(),
1,
"{:?}",
joined.iter().map(|b| &b.text).collect::<Vec<_>>()
);
assert!(joined[0].text.ends_with("reactivity in vivo."));
}
#[test]
fn a_flush_paragraph_after_a_short_item_stays_its_own() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let item = block(
BlockKind::ListItem {
level: 0,
ordinal: Some(2),
},
"2. Pears from the orchard and the garden",
72.0,
700.0,
);
let next = block(
BlockKind::Paragraph,
"and quinces from the walled garden were in.",
72.0,
688.0,
);
let decisions = vec![
Decision {
link_after: Some(0.99),
..Decision::default()
},
Decision::default(),
];
let (apart, _) = merge_items(
vec![item, next],
vec![0, 1],
&decisions,
&[true, false],
None,
);
assert_eq!(
apart.len(),
2,
"{:?}",
apart.iter().map(|b| &b.text).collect::<Vec<_>>()
);
}
#[test]
fn a_short_flush_line_under_a_short_item_stays_its_own() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let item = block(
BlockKind::ListItem {
level: 0,
ordinal: Some(2),
},
"2. Pears from the orchard",
72.0,
700.0,
);
let next = block(BlockKind::Paragraph, "and quinces, in.", 72.0, 688.0);
let decisions = vec![
Decision {
link_after: Some(0.99),
..Decision::default()
},
Decision::default(),
];
let (kept, _) = merge_items(
vec![item, next],
vec![0, 1],
&decisions,
&[true, false],
None,
);
assert_eq!(
kept.len(),
2,
"{:?}",
kept.iter().map(|b| &b.text).collect::<Vec<_>>()
);
}
#[test]
fn a_block_the_judge_chose_is_not_swallowed_as_a_list_continuation() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let out = || {
vec![
block(
BlockKind::ListItem {
level: 0,
ordinal: Some(5),
},
"S.-H. Park, O. Simeone, and S. Shamai, Robust compression for cloud radio access",
72.0,
700.0,
),
block(
BlockKind::Paragraph,
"networks, IEEE Trans. Veh. Techn., vol. 62, pp. 692-703, 2013.",
84.0,
688.0,
),
block(
BlockKind::Paragraph,
"and further notes on the same entry in the bibliography.",
84.0,
676.0,
),
]
};
let decisions = vec![
Decision {
link_after: Some(0.9),
..Decision::default()
},
Decision {
link_after: Some(0.9),
..Decision::default()
},
Decision::default(),
];
let marked = [true, false, false];
let (joined, _) = merge_items(out(), vec![0, 1, 2], &decisions, &marked, None);
assert_eq!(
joined.len(),
1,
"{:?}",
joined.iter().map(|b| &b.text).collect::<Vec<_>>()
);
let judged = [None, Some(0), Some(0)];
let (kept, origin) = merge_items(out(), vec![0, 1, 2], &decisions, &marked, Some(&judged));
assert_eq!(
kept.len(),
3,
"{:?}",
kept.iter().map(|b| &b.text).collect::<Vec<_>>()
);
assert_eq!(origin, vec![0, 1, 2]);
}
#[test]
fn a_block_the_judge_rejected_still_joins_its_item() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let out = vec![
block(
BlockKind::ListItem {
level: 0,
ordinal: Some(5),
},
"S.-H. Park, O. Simeone, and S. Shamai, Robust compression for cloud radio access",
72.0,
700.0,
),
block(
BlockKind::Paragraph,
"networks, IEEE Trans. Veh. Techn., vol. 62, pp. 692-703, 2013.",
84.0,
688.0,
),
block(
BlockKind::Paragraph,
"and further notes on the same entry in the bibliography.",
84.0,
676.0,
),
];
let decisions = vec![
Decision {
link_after: Some(0.9),
..Decision::default()
},
Decision {
link_after: Some(0.9),
..Decision::default()
},
Decision::default(),
];
let judged = [
None,
Some(super::tables::REJECTED),
Some(super::tables::REJECTED),
];
let (joined, _) = merge_items(
out,
vec![0, 1, 2],
&decisions,
&[true, false, false],
Some(&judged),
);
assert_eq!(
joined.len(),
1,
"{:?}",
joined.iter().map(|b| &b.text).collect::<Vec<_>>()
);
}
#[test]
fn a_list_item_the_judge_chose_still_joins_its_item() {
let block = |kind: BlockKind, text: &str, x: f64, y: f64| {
let mut b = Block::new(kind, 0, [0.0; 4], text.to_string(), 0);
b.spans = vec![line(text, x, y)];
crate::blocks::recompute_bbox(&mut b);
b
};
let out = vec![
block(
BlockKind::ListItem {
level: 0,
ordinal: Some(5),
},
"S.-H. Park, O. Simeone, and S. Shamai, Robust compression for cloud radio access",
72.0,
700.0,
),
block(
BlockKind::ListItem {
level: 0,
ordinal: None,
},
"networks, IEEE Trans. Veh. Techn., vol. 62, pp. 692-703, 2013.",
84.0,
688.0,
),
];
let decisions = vec![
Decision {
link_after: Some(0.9),
..Decision::default()
},
Decision::default(),
];
let judged = [None, Some(0)];
let (joined, _) = merge_items(out, vec![0, 1], &decisions, &[true, false], Some(&judged));
assert_eq!(
joined.len(),
1,
"{:?}",
joined.iter().map(|b| &b.text).collect::<Vec<_>>()
);
}
#[test]
fn each_piece_of_a_split_reports_its_own_probability() {
let decision = Decision {
item: None,
opens_at: vec![1, 2],
first_opens: true,
p_item: Some(0.11),
p_opens: vec![0.71, 0.72],
p_first_opens: Some(0.91),
heading_lines: Vec::new(),
p_heading_lines: Vec::new(),
line_classes: Vec::new(),
caption_lines: Vec::new(),
p_caption_lines: Vec::new(),
line_breaks: Vec::new(),
link_after: None,
};
let spans = vec![
line("first line", 72.0, 700.0),
line("second line", 72.0, 686.0),
line("third line", 72.0, 672.0),
];
let pieces = split(
block("first line second line third line", spans, 0),
&decision,
);
assert_eq!(pieces.len(), 3, "three lines, three parts: {pieces:?}");
let got: Vec<Option<f64>> = pieces.iter().map(|p| p.confidence).collect();
assert_eq!(
got,
vec![Some(0.91), Some(0.71), Some(0.72)],
"a piece reported a probability that belongs to another decision"
);
assert!(
got.iter().all(|c| *c != Some(0.11)),
"a piece inherited the block's p_item, which decided nothing about it"
);
}
#[test]
fn the_line_model_cuts_where_heading_lines_end_and_labels_each_part() {
let spans = || {
vec![
line("Operating review", 72.0, 700.0),
line("Sales rose in every region.", 72.0, 686.0),
line("Margins held.", 72.0, 672.0),
]
};
let decision = |flags: Vec<bool>, p: Vec<f64>| Decision {
heading_lines: flags,
p_heading_lines: p,
..Decision::default()
};
let text = "Operating review Sales rose in every region. Margins held.";
let parts = by_heading_lines(
block(text, spans(), 0),
&decision(vec![true, false, false], vec![0.9, 0.1, 0.2]),
);
assert_eq!(parts.len(), 2, "one cut, two parts: {parts:?}");
assert!(matches!(parts[0].kind, BlockKind::Heading { .. }));
assert_eq!(parts[0].text.trim(), "Operating review");
assert!(matches!(parts[1].kind, BlockKind::Paragraph));
assert_eq!(parts[0].confidence, Some(0.9));
assert!(parts[1].confidence.is_some_and(|c| (c - 0.15).abs() < 1e-9));
let mut heading = block(text, spans(), 0);
heading.kind = BlockKind::Heading { level: 2 };
let demoted = by_heading_lines(heading, &decision(vec![false; 3], vec![0.1; 3]));
assert_eq!(demoted.len(), 1);
assert!(matches!(demoted[0].kind, BlockKind::Paragraph));
let promoted = by_heading_lines(
block(text, spans(), 0),
&decision(vec![true; 3], vec![0.8; 3]),
);
assert_eq!(promoted.len(), 1);
assert!(matches!(promoted[0].kind, BlockKind::Heading { .. }));
}
fn block(text: &str, spans: Vec<TextSpan>, index: usize) -> Block {
Block {
kind: BlockKind::Paragraph,
page: 0,
bbox: bbox_of(&spans).unwrap_or([0.0; 4]),
text: text.to_string(),
spans,
index,
continues_previous: false,
confidence: None,
grid: None,
table_origin: TableOrigin::Other,
trace_learned: None,
}
}
#[test]
fn every_decision_is_read_from_one_snapshot() {
let page = |p: usize, base: usize| {
vec![
Block {
page: p,
..block(
"the following matters are to be reported on:",
vec![line(
"the following matters are to be reported on:",
72.0,
700.0,
)],
base,
)
},
Block {
page: p,
..block(
"\u{2022} that the accounts be laid before the members.",
vec![line(
"\u{2022} that the accounts be laid before the members.",
72.0,
680.0,
)],
base + 1,
)
},
Block {
page: p,
..block("1. 4,512", vec![line("1. 4,512", 480.0, 650.0)], base + 2)
},
]
};
let pages = [(612.0, 792.0), (612.0, 792.0)];
let forward: Vec<Block> = page(0, 0).into_iter().chain(page(1, 3)).collect();
let backward: Vec<Block> = page(1, 3).into_iter().chain(page(0, 0)).collect();
let a = decide(&forward, &pages, SWITCHES);
let b = decide(&backward, &pages, SWITCHES);
let by_index = |blocks: &[Block], d: Vec<Decision>| {
let mut v: Vec<(usize, Decision)> = blocks.iter().map(|b| b.index).zip(d).collect();
v.sort_by_key(|(i, _)| *i);
v
};
assert_eq!(by_index(&forward, a), by_index(&backward, b));
}
#[test]
fn a_split_block_takes_consecutive_places_and_everything_after_moves_down() {
let blocks = vec![
block("Before.", vec![line("Before.", 72.0, 700.0)], 0),
block(
"Apples are ripe. 2. Pears are not yet ripe.",
vec![
line("Apples are ripe.", 72.0, 680.0),
line("2. Pears are not", 72.0, 668.0),
line("yet ripe.", 72.0, 656.0),
],
1,
),
block("After.", vec![line("After.", 72.0, 630.0)], 2),
];
let decisions = vec![
Decision::default(),
Decision {
opens_at: vec![1],
..Decision::default()
},
Decision::default(),
];
let out = apply(blocks, &decisions);
let shape: Vec<(usize, &str)> = out.iter().map(|b| (b.index, b.text.as_str())).collect();
assert_eq!(
shape,
vec![
(0, "Before."),
(1, "Apples are ripe."),
(2, "Pears are not yet ripe."),
(3, "After."),
]
);
assert!(matches!(out[1].kind, BlockKind::Paragraph));
assert!(matches!(
out[2].kind,
BlockKind::ListItem {
ordinal: Some(2),
..
}
));
assert_eq!(out[1].spans.len(), 1);
assert_eq!(out[2].spans.len(), 2);
assert!(
out[2].bbox[3] <= out[1].bbox[1],
"the second part is below the first"
);
}
#[test]
fn a_demoted_item_keeps_its_body_text_and_gets_its_marker_back() {
let spans = vec![
line("1.", 72.0, 700.0),
line("appl", 90.0, 700.0),
line("es are ripe", 110.0, 700.0),
];
let item = Block {
kind: BlockKind::ListItem {
level: 0,
ordinal: Some(1),
},
..block("apples are ripe", spans, 0)
};
let out = apply(
vec![item],
&[Decision {
item: Some(false),
..Decision::default()
}],
);
assert!(matches!(out[0].kind, BlockKind::Paragraph));
assert_eq!(out[0].text, "1. apples are ripe");
}
#[test]
fn a_promoted_item_takes_the_indent_step_it_stands_on() {
let at = |x: f64, kind: BlockKind, index: usize| Block {
kind,
..block(
"apples",
vec![line("apples", x, 700.0 - index as f64 * 12.0)],
index,
)
};
let outer = at(
72.0,
BlockKind::ListItem {
level: 0,
ordinal: None,
},
0,
);
let paragraph = BlockKind::Paragraph;
let blocks = vec![outer, at(100.0, paragraph, 1), at(72.0, paragraph, 2)];
let decisions = vec![
Decision::default(),
Decision {
item: Some(true),
..Decision::default()
},
Decision {
item: Some(true),
..Decision::default()
},
];
let levels: Vec<u8> = apply(blocks, &decisions)
.iter()
.map(|b| match b.kind {
BlockKind::ListItem { level, .. } => level,
_ => u8::MAX,
})
.collect();
assert_eq!(levels, vec![0, 1, 0]);
}
#[test]
fn the_parts_of_a_split_item_keep_its_indent_step() {
let item = |x: f64, level: u8, index: usize, lines: &[&str]| Block {
kind: BlockKind::ListItem {
level,
ordinal: None,
},
..block(
&lines.join(" "),
lines
.iter()
.enumerate()
.map(|(k, t)| line(t, x, 700.0 - (index * 3 + k) as f64 * 12.0))
.collect(),
index,
)
};
let outer = item(72.0, 0, 0, &["apples"]);
let nested = item(100.0, 2, 1, &["pears", "\u{2022} plums"]);
let paragraph = Block {
kind: BlockKind::Paragraph,
..item(100.0, 0, 2, &["\u{2022} figs", "\u{2022} dates"])
};
let split_here = Decision {
opens_at: vec![1],
first_opens: true,
..Decision::default()
};
let decisions = vec![Decision::default(), split_here.clone(), split_here];
let levels: Vec<(String, u8)> = apply(vec![outer, nested, paragraph], &decisions)
.iter()
.map(|b| match b.kind {
BlockKind::ListItem { level, .. } => (b.text.clone(), level),
_ => (b.text.clone(), u8::MAX),
})
.collect();
let want: Vec<(String, u8)> = [
("apples", 0),
("pears", 2),
("plums", 2),
("figs", 1),
("dates", 1),
]
.iter()
.map(|(t, l)| (t.to_string(), *l))
.collect();
assert_eq!(levels, want);
}
#[test]
fn a_split_keeps_a_word_drawn_in_two_spans_whole() {
let mut tail = line("es are ripe.", 92.0, 680.0);
tail.x = 92.0;
let blocks = vec![block(
"apples are ripe. 2. Pears are not",
vec![
line("appl", 72.0, 680.0),
tail,
line("2. Pears are not", 72.0, 668.0),
],
0,
)];
let decisions = vec![Decision {
opens_at: vec![1],
..Decision::default()
}];
let texts: Vec<String> = apply(blocks, &decisions)
.into_iter()
.map(|b| b.text)
.collect();
assert_eq!(texts, vec!["apples are ripe.", "Pears are not"]);
}
#[test]
fn a_first_authors_initial_is_not_a_marker_to_take_off() {
assert_eq!(leading_marker("S. G. Ovchinnikov and V. N. Zabluda"), None);
assert_eq!(leading_marker("IV. Definitions"), Some((4, Some(4))));
assert_eq!(leading_marker("B. Scope of the works"), Some((3, Some(2))));
assert_eq!(leading_marker("II. A. Smith"), Some((4, Some(2))));
}
#[test]
fn a_marker_comes_off_the_front_and_its_number_stays() {
assert_eq!(leading_marker("3. Pears"), Some((3, Some(3))));
assert_eq!(leading_marker("(c) Plums"), Some((4, Some(3))));
assert_eq!(
leading_marker("\u{2022} Apples"),
Some(("\u{2022} ".len(), None))
);
assert_eq!(leading_marker("Apples and pears"), None);
for dash in ["-", "\u{2013}", "\u{2014}"] {
let text = format!("{dash} Pears");
assert_eq!(
leading_marker(&text),
Some((dash.len() + 1, None)),
"{text:?}"
);
}
assert_eq!(
leading_marker("-5.2 per cent"),
None,
"a negative figure is no dash"
);
assert_eq!(leading_marker("* Pears"), Some((2, None)));
assert_eq!(leading_marker(">Our market share"), Some((1, None)));
assert_eq!(leading_marker("\u{f0b7} Pears"), Some((4, None)));
assert_eq!(leading_marker("iv. Pears"), Some((4, Some(4))));
assert_eq!(leading_marker("(xii) Pears"), Some((6, Some(12))));
assert_eq!(
leading_marker("i. Pears"),
Some((3, Some(9))),
"a lone `i` is a letter"
);
}
#[test]
fn the_image_arm_ships_off() {
const { assert!(!SWITCHES.line_image) };
}
#[test]
fn the_stacked_figures_default_has_c_on_and_b_on() {
const { assert!(SWITCHES.table_runs_on_figures) };
const { assert!(SWITCHES.markdown_header_block_figures) };
}
#[test]
fn the_shipped_default_has_rows_on_and_header_on() {
const { assert!(SWITCHES.table_stitch_rows) };
const { assert!(SWITCHES.table_stitch_header) };
}
#[test]
fn the_shipped_default_drops_empty_columns() {
const { assert!(SWITCHES.table_drop_empty_columns) };
}
#[test]
fn the_shipped_default_reads_a_cell_line_by_line() {
const { assert!(SWITCHES.table_cell_text_by_line) };
}
#[test]
fn the_shipped_default_rejoins_a_refusal_in_row_order() {
const { assert!(SWITCHES.table_v2_release_rejoin) };
const { assert!(!SWITCHES.table_v2_release_column_order) };
}
#[test]
fn the_shipped_default_has_stacked_on_and_side_by_side_off() {
const { assert!(SWITCHES.table_split_stacked) };
const { assert!(!SWITCHES.table_split_side_by_side) };
}
#[test]
fn the_shipped_default_splits_a_joined_stripe_at_white() {
const { assert!(SWITCHES.table_corridor_split) };
}
#[test]
fn the_retrained_bar_counts_a_line_at_0_55() {
let v2 = Switches {
line_gnn: true,
line_gnn_v2: true,
..SWITCHES
};
let v1 = Switches {
line_gnn: true,
line_gnn_v2: false,
..SWITCHES
};
assert_eq!(heading_bar("paragraph", v2), 0.5);
assert_eq!(heading_bar("paragraph", v1), 0.6);
let p = 0.55;
assert!(
p >= heading_bar("paragraph", v2),
"0.55 is a heading at the retrained bar"
);
assert!(
p < heading_bar("paragraph", v1),
"0.55 is prose at the old bar"
);
let q = 0.40;
assert!(
(q - heading_bar("paragraph", v2)).abs() <= 0.15,
"0.40 counts in the doubt around 0.5"
);
assert!(
(q - heading_bar("paragraph", v1)).abs() > 0.15,
"0.40 is outside the doubt around 0.6"
);
assert_eq!(heading_bar("heading", v2), LINE_GNN_KEEP_HEADING);
assert_eq!(heading_bar("heading", v1), LINE_GNN_KEEP_HEADING);
let off = Switches {
line_gnn: false,
..SWITCHES
};
assert_eq!(heading_bar("paragraph", off), HEADING_LINES_THRESHOLD);
}
#[test]
fn the_pass_and_its_doubt_agree_on_every_line_of_a_real_page() {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/doclaynet/nyse_ssl_2003_p59.pdf"
);
let doc = crate::PdfDocument::open(std::fs::read(path).expect("fixture")).expect("open");
let sizes: Vec<(f64, f64)> = (0..doc.page_count())
.map(|p| {
doc.page_geometry(p)
.map(|g| g.effective_dimensions())
.unwrap_or((0.0, 0.0))
})
.collect();
let blocks = doc.extract_blocks_before_learning().expect("blocks");
let facts = document_facts(&blocks);
let mut lines_checked = 0usize;
let mut verdict_split = 0usize;
let mut doubt_split = 0usize;
for (page_no, members) in pages_of(&blocks).iter().enumerate() {
let page: Vec<&Block> = members.iter().map(|&i| &blocks[i]).collect();
let size = sizes.get(page_no).copied().unwrap_or((0.0, 0.0));
let (decisions, vectors, _) = decide_page_traced(&page, size, &facts, SWITCHES, &[]);
let conf = confidence_of_page(&page, size, &facts, SWITCHES, &[]);
let mut at = 0usize;
let mut near = 0usize;
let mut decided = 0usize;
for (decision, v) in decisions.iter().zip(&vectors) {
assert_eq!(
decision.heading_lines.len(),
decision.p_heading_lines.len(),
"a decided line carries both a verdict and its probability"
);
let bar = heading_bar(v.kind, SWITCHES);
for k in 0..v.lines.len() {
let p = decision.p_heading_lines.get(k).copied();
assert_eq!(
p, conf.lines[at].p_heading,
"line {at}: the doubt reads another probability than the pass decided on"
);
if let Some(p) = p {
decided += 1;
assert_eq!(
decision.heading_lines[k],
p >= bar,
"line {at} at {p}: the pass and its doubt disagree"
);
if (p - bar).abs() <= 0.15 {
near += 1;
}
if v.kind != "heading" {
if (0.5..0.6).contains(&p) {
verdict_split += 1;
}
if ((p - 0.5).abs() <= 0.15) != ((p - 0.6).abs() <= 0.15) {
doubt_split += 1;
}
}
}
at += 1;
}
}
assert_eq!(at, conf.lines.len(), "every line belongs to a block");
let expected = if decided == 0 {
0.0
} else {
near as f64 / decided as f64
};
assert!(
(conf.heading_doubt - expected).abs() < 1e-12,
"heading doubt {} is not the share around the pass's bar ({expected})",
conf.heading_doubt
);
lines_checked += at;
}
assert!(lines_checked > 0, "the page has no lines to agree on");
assert!(
verdict_split > 0,
"no line splits the old and the retrained verdict"
);
assert!(
doubt_split > 0,
"no line splits the doubt around 0.5 from the doubt around 0.6"
);
}
fn first_page_of(path: &str) -> (Vec<Block>, (f64, f64), DocumentFacts) {
let doc = crate::PdfDocument::open(std::fs::read(path).expect("fixture")).expect("open");
let size = doc
.page_geometry(0)
.map(|g| g.effective_dimensions())
.unwrap_or((0.0, 0.0));
let blocks = doc.extract_blocks_before_learning().expect("blocks");
let facts = document_facts(&blocks);
let page = page_blocks(&blocks, 0);
assert!(!page.is_empty(), "the fixture's first page has no blocks");
(page, size, facts)
}
#[test]
fn the_judge_and_the_decisions_measure_a_page_three_times() {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/doclaynet/nyse_ssl_2003_p59.pdf"
);
let (page, size, facts) = first_page_of(path);
features::reset_page_features_calls();
let _ = judge_and_decide_legacy(&page, size, &facts, SWITCHES, &[], &[]);
assert_eq!(
features::page_features_calls(),
3,
"the legacy composition no longer measures the page three times"
);
}
#[test]
fn the_shared_pass_measures_a_page_once() {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/doclaynet/nyse_ssl_2003_p59.pdf"
);
let (page, size, facts) = first_page_of(path);
features::reset_page_features_calls();
let _ = judge_and_decide_shared(&page, size, &facts, SWITCHES, &[], &[]);
assert_eq!(
features::page_features_calls(),
1,
"the shared pass measures the page more than once"
);
}
#[test]
fn refine_page_measures_twice() {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/doclaynet/nyse_ssl_2003_p59.pdf"
);
let (page, size, facts) = first_page_of(path);
features::reset_page_features_calls();
let _ = refine_page(page, size, &facts, SWITCHES);
assert_eq!(
features::page_features_calls(),
2,
"the page pass no longer measures twice"
);
}
#[test]
fn the_shared_pass_decides_what_the_separate_pass_decided() {
let pages = [
concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/doclaynet/nasdaq_wwd_2002_p7.pdf"
),
concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/doclaynet/nyse_ssl_2003_p59.pdf"
),
concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/learned/fw7.pdf"
),
];
let mut changed = 0usize;
for path in pages {
let (page, size, facts) = first_page_of(path);
let shared = judge_and_decide_shared(&page, size, &facts, SWITCHES, &[], &[]);
let legacy = judge_and_decide_legacy(&page, size, &facts, SWITCHES, &[], &[]);
assert_eq!(shared.0, legacy.0, "{path}: the judge differs");
assert_eq!(shared.1, legacy.1, "{path}: the decisions differ");
let decided = legacy.1.iter().any(|d| d != &Decision::default())
|| legacy
.0
.as_ref()
.is_some_and(|j| j.regions.iter().any(Option::is_some));
changed += usize::from(decided);
}
assert!(
changed > 0,
"the models change nothing on any page, so the identity tests nothing"
);
}
}