diff --git a/napi/package.json b/napi/package.json index f415c90..59eb1b9 100644 --- a/napi/package.json +++ b/napi/package.json @@ -1,6 +1,6 @@ { "name": "@firecrawl/pdf-inspector", - "version": "1.5.0", + "version": "1.6.0", "description": "Fast PDF classification and text extraction. Detect text-based vs scanned PDFs, extract text by region with quality checks. Native Rust performance via napi-rs.", "main": "index.js", "types": "index.d.ts", diff --git a/napi/src/lib.rs b/napi/src/lib.rs index 0a9a2e6..6dbac26 100644 --- a/napi/src/lib.rs +++ b/napi/src/lib.rs @@ -317,6 +317,141 @@ pub fn extract_tables_in_regions( }) } +/// One cropped table region plus its raw structure-recovery output, for +/// `extractTablesWithStructure`. +/// +/// `structureTokens` and `cellBboxes` are typically produced by an external +/// table-structure recognition model (e.g. SLANet on PaddleOCR) running on +/// a rendered crop of the page. pdf-inspector uses the structure to lay out +/// the cells and pulls the cell text from the native PDF — no OCR involved. +#[napi(object)] +pub struct TsrTableInputJs { + /// 0-indexed page number where the crop was taken from. + pub page: u32, + /// Crop bbox on the page, `[x1, y1, x2, y2]` in PDF points with + /// top-left origin. + pub crop_pdf_pt_bbox: Vec, + /// DPI the crop image was rendered at (e.g. `200.0`). + pub render_dpi: f64, + /// Raw structure tokens emitted by the TSR model, in document order. + pub structure_tokens: Vec, + /// One bbox per cell (in document order). May be 4-element + /// `[x1,y1,x2,y2]` or 8-element 4-corner polygon, in crop image-pixel + /// space. + pub cell_bboxes: Vec>, +} + +/// Extract markdown tables using externally-supplied structure recovery. +/// +/// For each input, pairs structure tokens with cell bboxes (rowspan/colspan +/// aware), converts each cell bbox from crop image-pixels into page PDF +/// points, pulls the cell's text from the native PDF, and emits a markdown +/// pipe-table. +/// +/// Returns one markdown string per input, in input order. +#[napi] +pub fn extract_tables_with_structure( + buffer: Buffer, + inputs: Vec, +) -> Result> { + let bytes: Vec = buffer.to_vec(); + let parsed = parse_tsr_inputs(&inputs); + + catch_panic("extract_tables_with_structure", move || { + pdf_inspector::extract_tables_with_structure_mem(&bytes, &parsed) + .map_err(|e| to_napi_err(e, "extract_tables_with_structure")) + }) +} + +/// One resolved cell from `extractTablesWithStructureCells`. +#[napi(object)] +pub struct StructuredCellJs { + /// 0-indexed grid row. + pub row: u32, + /// 0-indexed grid column. + pub col: u32, + /// 1 for a normal cell. + pub rowspan: u32, + /// 1 for a normal cell. + pub colspan: u32, + /// `true` when the cell is a `` or sits inside ``. + pub is_header: bool, + /// Text extracted from the native PDF for this cell (may be empty). + pub text: String, + /// Axis-aligned bbox `[x1, y1, x2, y2]` in page PDF-points, top-left + /// origin. Useful for debug overlays or per-cell post-processing. + pub page_pt_bbox: Vec, +} + +/// Extract structured cells using externally-supplied structure recovery. +/// +/// Lower-level sibling of [`extractTablesWithStructure`]: instead of +/// rendering markdown, returns the resolved cells (row, col, rowspan, +/// colspan, isHeader, text, pagePtBbox) so callers can drive their own +/// rendering, debug overlays, or per-cell post-processing. +/// +/// Returns one `Array` per input, in input order. +#[napi] +pub fn extract_tables_with_structure_cells( + buffer: Buffer, + inputs: Vec, +) -> Result>> { + let bytes: Vec = buffer.to_vec(); + let parsed = parse_tsr_inputs(&inputs); + + catch_panic("extract_tables_with_structure_cells", move || { + let result = pdf_inspector::extract_tables_with_structure_cells_mem(&bytes, &parsed) + .map_err(|e| to_napi_err(e, "extract_tables_with_structure_cells"))?; + Ok(result + .into_iter() + .map(|cells| { + cells + .into_iter() + .map(|c| StructuredCellJs { + row: c.row as u32, + col: c.col as u32, + rowspan: c.rowspan as u32, + colspan: c.colspan as u32, + is_header: c.is_header, + text: c.text, + page_pt_bbox: c.page_pt_bbox.iter().map(|v| *v as f64).collect(), + }) + .collect() + }) + .collect()) + }) +} + +fn parse_tsr_inputs(inputs: &[TsrTableInputJs]) -> Vec { + inputs + .iter() + .map(|i| { + let crop = if i.crop_pdf_pt_bbox.len() == 4 { + [ + i.crop_pdf_pt_bbox[0] as f32, + i.crop_pdf_pt_bbox[1] as f32, + i.crop_pdf_pt_bbox[2] as f32, + i.crop_pdf_pt_bbox[3] as f32, + ] + } else { + [0.0, 0.0, 0.0, 0.0] + }; + let cell_bboxes: Vec> = i + .cell_bboxes + .iter() + .map(|bb| bb.iter().map(|v| *v as f32).collect()) + .collect(); + pdf_inspector::TsrTableInput { + page: i.page, + crop_pdf_pt_bbox: crop, + render_dpi: i.render_dpi as f32, + structure_tokens: i.structure_tokens.clone(), + cell_bboxes, + } + }) + .collect() +} + /// Per-page markdown extraction result. #[napi(object)] pub struct PageMarkdownResult { @@ -360,9 +495,8 @@ pub fn extract_pages_markdown( ) -> Result { let bytes: Vec = buffer.to_vec(); catch_panic("extract_pages_markdown", move || { - let result = - pdf_inspector::extract_pages_markdown_mem(&bytes, pages.as_deref()) - .map_err(|e| to_napi_err(e, "extract_pages_markdown"))?; + let result = pdf_inspector::extract_pages_markdown_mem(&bytes, pages.as_deref()) + .map_err(|e| to_napi_err(e, "extract_pages_markdown"))?; Ok(PagesExtractionResult { pages: result .pages diff --git a/src/lib.rs b/src/lib.rs index 5fcafeb..55b59c2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -784,6 +784,198 @@ pub fn extract_tables_in_regions_mem( Ok(results) } +// ========================================================================= +// Region-based table extraction with external structure recovery (TSR) +// ========================================================================= + +/// Input for [`extract_tables_with_structure_mem`]: one cropped table region +/// plus the raw structure-recovery output for it. +/// +/// The structure tokens and bboxes are typically produced by an external +/// table-structure recognition model (e.g. SLANet on PaddleOCR) running on +/// a rendered crop of the page. pdf-inspector uses the structure to lay out +/// the cells and pulls the cell text from the native PDF — no OCR involved. +#[derive(Debug, Clone)] +pub struct TsrTableInput { + /// 0-indexed page number where the crop was taken from. + pub page: u32, + /// Crop bbox on the page, `[x1, y1, x2, y2]` in PDF points with + /// **top-left origin** (matches the layout model's coordinate space). + pub crop_pdf_pt_bbox: [f32; 4], + /// DPI the crop image was rendered at (e.g. `200.0`). Used to convert + /// cell bboxes from image-pixels back to PDF points. + pub render_dpi: f32, + /// Raw structure tokens emitted by the TSR model, in document order. + /// See [`tables::structured::parse_structure`] for the accepted grammar. + pub structure_tokens: Vec, + /// One bbox per cell (in document order, parallel to the cell open-tags + /// in `structure_tokens`). May be 4-element `[x1,y1,x2,y2]` or + /// 8-element 4-corner polygon, in **crop image-pixel space**. + pub cell_bboxes: Vec>, +} + +/// Extract structured cells using externally-supplied structure recovery. +/// +/// For each input, this: +/// 1. Pairs each cell open-tag in `structure_tokens` with the next bbox in +/// `cell_bboxes` (document order), tracking row/col with rowspan/colspan +/// awareness. +/// 2. Converts each cell bbox from crop image-pixels into page PDF-points. +/// 3. Pulls the cell's text by overlap-testing PDF text items inside that +/// bbox — same primitives used by [`extract_text_in_regions_mem`]. +/// +/// Returns one `Vec` per input, in input order. Each cell +/// carries its (row, col, rowspan, colspan, is_header) metadata, the +/// extracted text, and its page-PDF-pt bbox so callers can do their own +/// rendering, debug overlays, or per-cell post-processing. +/// +/// Inputs whose page is out of range or whose tokens parse to zero cells +/// produce an empty `Vec`. +/// +/// See [`extract_tables_with_structure_mem`] if you just want the rendered +/// markdown. +pub fn extract_tables_with_structure_cells_mem( + buffer: &[u8], + inputs: &[TsrTableInput], +) -> Result>, PdfError> { + use tables::structured::{ + cell_px_to_page_pt, parse_structure, polygon_to_aabb, StructuredCell, + }; + + validate_pdf_bytes(buffer)?; + let (doc, _page_count) = load_document_from_mem(buffer)?; + let pages = doc.get_pages(); + + let needed_pages: HashSet = inputs.iter().map(|t| t.page + 1).collect(); + let font_cmaps = FontCMaps::from_doc_pages_fast(&doc, Some(&needed_pages)); + + let mut items_by_page: HashMap> = HashMap::new(); + let mut page_heights: HashMap = HashMap::new(); + let mut page_thresholds: HashMap = HashMap::new(); + let mut rotated_pages: HashSet = HashSet::new(); + + for (page_num, &page_id) in pages.iter() { + if !needed_pages.contains(page_num) { + continue; + } + let height = get_page_height(&doc, page_id).unwrap_or(792.0); + page_heights.insert(*page_num, height); + + let ((mut items, _rects, _lines), _has_gid, coords_rotated) = + extractor::content_stream::extract_page_text_items( + &doc, + page_id, + *page_num, + &font_cmaps, + false, + )?; + let threshold = text_utils::fix_letterspaced_items(&mut items); + if threshold > 0.10 { + page_thresholds.insert(*page_num, threshold); + } + if coords_rotated { + rotated_pages.insert(*page_num); + } + items_by_page.insert(*page_num, items); + } + + let mut results: Vec> = Vec::with_capacity(inputs.len()); + + for input in inputs { + let page_1idx = input.page + 1; + let Some(items) = items_by_page.get(&page_1idx) else { + // Out-of-range page or page with no extractable text — emit empty. + results.push(Vec::new()); + continue; + }; + let page_h = page_heights.get(&page_1idx).copied().unwrap_or(792.0); + let adaptive_threshold = page_thresholds.get(&page_1idx).copied().unwrap_or(0.10); + let coords = if rotated_pages.contains(&page_1idx) { + RegionCoordSpace::Rotated90Ccw + } else { + RegionCoordSpace::Standard + }; + + let crop_origin = [input.crop_pdf_pt_bbox[0], input.crop_pdf_pt_bbox[1]]; + + let slots = parse_structure(&input.structure_tokens); + if slots.is_empty() { + results.push(Vec::new()); + continue; + } + + let mut cells: Vec = Vec::with_capacity(slots.len()); + for slot in &slots { + let cell_text; + let page_pt_bbox; + + if let Some(coords_arr) = input.cell_bboxes.get(slot.bbox_idx) { + if let Some(aabb_px) = polygon_to_aabb(coords_arr) { + let aabb_pt = cell_px_to_page_pt(aabb_px, input.render_dpi, crop_origin); + let raw = collect_text_in_region_with_options( + items, + aabb_pt[0], + aabb_pt[1], + aabb_pt[2], + aabb_pt[3], + page_h, + coords, + adaptive_threshold, + ); + // Markdown cells must be one line — collapse line breaks + // produced by the line-grouping pass. + cell_text = raw.replace(['\n', '\r'], " "); + page_pt_bbox = aabb_pt; + } else { + cell_text = String::new(); + page_pt_bbox = [0.0, 0.0, 0.0, 0.0]; + } + } else { + cell_text = String::new(); + page_pt_bbox = [0.0, 0.0, 0.0, 0.0]; + } + + cells.push(StructuredCell { + row: slot.row, + col: slot.col, + rowspan: slot.rowspan, + colspan: slot.colspan, + is_header: slot.is_header, + text: cell_text, + page_pt_bbox, + }); + } + + results.push(cells); + } + + Ok(results) +} + +/// Extract markdown tables using externally-supplied structure recovery. +/// +/// Convenience wrapper around [`extract_tables_with_structure_cells_mem`] +/// that renders each cell list to markdown via +/// [`tables::cells_to_markdown`]. Returns one markdown string per input, +/// in input order. Inputs whose page is out of range or whose tokens parse +/// to zero cells produce an empty string. +pub fn extract_tables_with_structure_mem( + buffer: &[u8], + inputs: &[TsrTableInput], +) -> Result, PdfError> { + let cells_lists = extract_tables_with_structure_cells_mem(buffer, inputs)?; + Ok(cells_lists + .into_iter() + .map(|cells| { + if cells.is_empty() { + String::new() + } else { + tables::cells_to_markdown(&cells) + } + }) + .collect()) +} + /// Get page height in points from MediaBox. fn get_page_height(doc: &Document, page_id: lopdf::ObjectId) -> Option { let page_dict = doc.get_dictionary(page_id).ok()?; diff --git a/src/tables/mod.rs b/src/tables/mod.rs index c2e55eb..b69d2be 100644 --- a/src/tables/mod.rs +++ b/src/tables/mod.rs @@ -9,6 +9,7 @@ mod detect_struct; mod financial; mod format; mod grid; +pub mod structured; pub use detect_heuristic::detect_tables; pub(crate) use detect_heuristic::is_table_of_contents; @@ -17,6 +18,7 @@ pub(crate) use detect_rects::cluster_rects; pub use detect_rects::{detect_tables_from_rects, RectHintRegion}; pub use detect_struct::detect_tables_from_struct_tree; pub use format::table_to_markdown; +pub use structured::{cells_to_markdown, StructuredCell}; use crate::types::TextItem; diff --git a/src/tables/structured.rs b/src/tables/structured.rs new file mode 100644 index 0000000..54ae5e9 --- /dev/null +++ b/src/tables/structured.rs @@ -0,0 +1,738 @@ +//! Structure-recovery-aware (TSR) table assembly. +//! +//! Consumes the raw output of an external table-structure recognition model +//! (e.g. SLANet on PaddleOCR): a flat list of HTML structure tokens plus a +//! parallel list of per-cell bboxes. Pairs each cell open-tag with its bbox +//! in document order, tracks row/column position with rowspan/colspan +//! awareness, and emits a markdown pipe-table. +//! +//! No real HTML parser is needed — the token grammar is restricted (see +//! [`parse_structure`]), so a small state machine is enough. +//! +//! Cell text is supplied separately by the caller (typically by overlap- +//! testing PDF text items against each cell's page-PDF-pt bbox). + +use std::collections::HashSet; + +/// A single resolved cell, with both structural metadata and its bbox in +/// page PDF-points (top-left origin). +#[derive(Debug, Clone)] +pub struct StructuredCell { + /// 0-indexed grid row. + pub row: usize, + /// 0-indexed grid column. + pub col: usize, + /// 1 for a normal cell. + pub rowspan: usize, + /// 1 for a normal cell. + pub colspan: usize, + /// `true` when the cell is a `` or sits inside ``. + pub is_header: bool, + /// Cell text (filled in by the caller after overlap-testing PDF items). + pub text: String, + /// Axis-aligned bbox `[x1, y1, x2, y2]` in page PDF-points, top-left origin. + pub page_pt_bbox: [f32; 4], +} + +/// Intermediate parse result before the caller fills in text + page coords. +#[derive(Debug, Clone)] +pub(crate) struct CellSlot { + pub row: usize, + pub col: usize, + pub rowspan: usize, + pub colspan: usize, + pub is_header: bool, + /// Index into the parallel `cell_bboxes` array. + pub bbox_idx: usize, +} + +/// Parse a sequence of SLANet structure tokens into ordered cell slots. +/// +/// Token grammar (no real HTML parsing required): +/// - Section markers: ``, ``, ``, `` and +/// wrapper tokens (``, ``, ``, plus closing variants) +/// are tracked or skipped. +/// - Row markers: `` opens a new row, `` is informational. +/// - Empty cell, single token: `` or ``. +/// - Cell with attributes, multi-token sequence: ``, then later `` +/// (or ``). Cells get paired with the next bbox in document order. +/// +/// Cells inside `` and any `" => { + in_thead = true; + } + "" => { + in_thead = false; + } + "" => { + if started_first_row { + row += 1; + } + col = 0; + started_first_row = true; + } + "" | "" => { + let is_th = tok == ""; + while occupied.contains(&(row, col)) { + col += 1; + } + slots.push(CellSlot { + row, + col, + rowspan: 1, + colspan: 1, + is_header: in_thead || is_th, + bbox_idx, + }); + bbox_idx += 1; + col += 1; + } + " { + let is_th = tok == "". + i += 1; + while i < tokens.len() && tokens[i].trim() != ">" { + let attr = tokens[i].as_str(); + if let Some(v) = parse_int_attr(attr, "rowspan") { + rowspan = v.max(1); + } else if let Some(v) = parse_int_attr(attr, "colspan") { + colspan = v.max(1); + } + i += 1; + } + // i now points at the `>` token (or off the end if malformed). + while occupied.contains(&(row, col)) { + col += 1; + } + slots.push(CellSlot { + row, + col, + rowspan, + colspan, + is_header: in_thead || is_th, + bbox_idx, + }); + for r in row..row + rowspan { + for c in col..col + colspan { + occupied.insert((r, c)); + } + } + bbox_idx += 1; + col += colspan; + } + // Wrapper / informational tokens — no-op. + _ => {} + } + i += 1; + } + + slots +} + +/// Parse an attribute fragment like ` colspan="4"` or `rowspan='2'`. +/// +/// Tolerates leading whitespace and either single or double quotes. +fn parse_int_attr(s: &str, name: &str) -> Option { + let trimmed = s.trim(); + if !trimmed.starts_with(name) { + return None; + } + let rest = trimmed[name.len()..].trim_start(); + let rest = rest.strip_prefix('=')?.trim_start(); + let value = rest + .trim_start_matches(['"', '\'']) + .trim_end_matches(['"', '\'']); + value.parse().ok() +} + +/// Convert a SLANet polygon (4 or 8 elements) into an axis-aligned +/// `[x1, y1, x2, y2]` rect. +/// +/// 8-element form: `[x1,y1, x2,y1, x2,y2, x1,y2]` (4 corners). We ignore the +/// implicit corner order and just take min/max so rotated polygons collapse +/// to a sane bounding box. +/// +/// 4-element form: `[x1, y1, x2, y2]` (axis-aligned, older SLANet variants). +pub(crate) fn polygon_to_aabb(coords: &[f32]) -> Option<[f32; 4]> { + match coords.len() { + 4 => { + let x1 = coords[0].min(coords[2]); + let y1 = coords[1].min(coords[3]); + let x2 = coords[0].max(coords[2]); + let y2 = coords[1].max(coords[3]); + Some([x1, y1, x2, y2]) + } + 8 => { + let xs = [coords[0], coords[2], coords[4], coords[6]]; + let ys = [coords[1], coords[3], coords[5], coords[7]]; + let x1 = xs.iter().copied().fold(f32::INFINITY, f32::min); + let y1 = ys.iter().copied().fold(f32::INFINITY, f32::min); + let x2 = xs.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let y2 = ys.iter().copied().fold(f32::NEG_INFINITY, f32::max); + if x1.is_finite() && y1.is_finite() && x2.is_finite() && y2.is_finite() { + Some([x1, y1, x2, y2]) + } else { + None + } + } + _ => None, + } +} + +/// Convert a cell rect from crop image-pixel space to page PDF-points +/// (top-left origin), given the crop's PDF-point offset on the page and the +/// DPI the crop image was rendered at. +pub(crate) fn cell_px_to_page_pt( + cell_px: [f32; 4], + render_dpi: f32, + crop_origin_pt: [f32; 2], +) -> [f32; 4] { + let pt_per_px = if render_dpi > 0.0 { + 72.0 / render_dpi + } else { + 1.0 + }; + let [x_off, y_off] = crop_origin_pt; + [ + cell_px[0] * pt_per_px + x_off, + cell_px[1] * pt_per_px + y_off, + cell_px[2] * pt_per_px + x_off, + cell_px[3] * pt_per_px + y_off, + ] +} + +/// Sanitize cell text for inclusion in a markdown pipe-table cell: +/// collapse whitespace runs, drop newlines/tabs (cells must be one line), +/// and escape pipes that would otherwise break the table. +fn sanitize_cell(text: &str) -> String { + let mut s = String::with_capacity(text.len()); + let mut prev_space = false; + for c in text.chars() { + match c { + '|' => { + s.push_str("\\|"); + prev_space = false; + } + '\n' | '\r' | '\t' | ' ' => { + if !prev_space { + s.push(' '); + } + prev_space = true; + } + other => { + s.push(other); + prev_space = false; + } + } + } + s.trim().to_string() +} + +/// Render a list of explicitly-positioned cells as a markdown pipe-table. +/// +/// Grid dimensions are inferred from the cells' (row, col, rowspan, colspan) +/// extents. A cell with colspan/rowspan > 1 is rendered in its top-left +/// position; the absorbed grid positions are emitted as empty cells so the +/// markdown stays a valid rectangular grid that downstream readers can +/// column-count correctly. +/// +/// The separator row (`|---|...|`) is emitted after the **last** row that +/// contains a header cell (`is_header == true`). When no cells are flagged +/// as headers — e.g. the upstream TSR model didn't emit ``/` + // Expected: A at (0,0), B at (0,1), C at (1,1) — col 0 of row 1 + // is occupied by A's rowspan. + let tokens: Vec = vec![ + "
` cells are flagged as headers. +/// rowspan/colspan attributes are honoured and prior-row rowspans push +/// later-row cells to the right. +pub(crate) fn parse_structure(tokens: &[String]) -> Vec { + let mut slots: Vec = Vec::new(); + let mut occupied: HashSet<(usize, usize)> = HashSet::new(); + let mut row: usize = 0; + let mut col: usize = 0; + let mut bbox_idx: usize = 0; + let mut in_thead = false; + let mut started_first_row = false; + + let mut i = 0; + while i < tokens.len() { + let tok = tokens[i].trim(); + match tok { + "
` — +/// the separator falls back to "after row 0" so the output is still a +/// valid pipe-table. +pub fn cells_to_markdown(cells: &[StructuredCell]) -> String { + if cells.is_empty() { + return String::new(); + } + let num_rows = cells + .iter() + .map(|c| c.row + c.rowspan.max(1)) + .max() + .unwrap_or(0); + let num_cols = cells + .iter() + .map(|c| c.col + c.colspan.max(1)) + .max() + .unwrap_or(0); + if num_rows == 0 || num_cols == 0 { + return String::new(); + } + + // Separator goes after the last header row, falling back to row 0 when + // no header cells exist. Clamped into range so a malformed cell with + // row >= num_rows can't push it past the table. + let separator_after_row = cells + .iter() + .filter(|c| c.is_header) + .map(|c| c.row) + .max() + .unwrap_or(0) + .min(num_rows.saturating_sub(1)); + + let mut grid: Vec> = vec![vec![String::new(); num_cols]; num_rows]; + for cell in cells { + if cell.row < num_rows && cell.col < num_cols { + grid[cell.row][cell.col] = sanitize_cell(&cell.text); + } + } + + let mut output = String::new(); + for (row_idx, row) in grid.iter().enumerate() { + output.push('|'); + for cell in row { + output.push_str(cell); + output.push('|'); + } + output.push('\n'); + if row_idx == separator_after_row { + output.push('|'); + for _ in 0..num_cols { + output.push_str("---|"); + } + output.push('\n'); + } + } + output +} + +#[cfg(test)] +mod tests { + use super::*; + + fn t(s: &str) -> String { + s.to_string() + } + + /// Tokens for the synthetic 3×3 grid example (one colspan-4 row + two + /// data rows of 4 cells each = 9 cells total, 3 rows × 4 cols). + fn synthetic_3x3_tokens() -> Vec { + vec![ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + "", + "", + ] + .into_iter() + .map(t) + .collect() + } + + /// Bboxes for the synthetic 3×3 grid (8-element polygon form), all + /// within a 400×120 px crop. + fn synthetic_3x3_bboxes() -> Vec> { + vec![ + vec![3.0, 2.0, 395.0, 2.0, 396.0, 59.0, 3.0, 59.0], + vec![26.0, 62.0, 140.0, 62.0, 141.0, 120.0, 26.0, 120.0], + vec![149.0, 64.0, 248.0, 64.0, 248.0, 119.0, 149.0, 119.0], + vec![257.0, 64.0, 350.0, 64.0, 350.0, 119.0, 257.0, 119.0], + vec![359.0, 64.0, 395.0, 64.0, 395.0, 119.0, 359.0, 119.0], + vec![26.0, 122.0, 140.0, 122.0, 140.0, 178.0, 26.0, 178.0], + vec![149.0, 124.0, 248.0, 124.0, 248.0, 179.0, 149.0, 179.0], + vec![257.0, 124.0, 350.0, 124.0, 350.0, 179.0, 257.0, 179.0], + vec![359.0, 124.0, 395.0, 124.0, 395.0, 179.0, 359.0, 179.0], + ] + } + + #[test] + fn parse_structure_synthetic_3x3() { + let tokens = synthetic_3x3_tokens(); + let slots = parse_structure(&tokens); + + assert_eq!(slots.len(), 9, "should parse 9 cells"); + + // Cell 0: row 0 col 0, colspan 4 + assert_eq!(slots[0].row, 0); + assert_eq!(slots[0].col, 0); + assert_eq!(slots[0].colspan, 4); + assert_eq!(slots[0].rowspan, 1); + + // Cells 1..5: row 1, cols 0..3 + for (i, slot) in slots.iter().enumerate().skip(1).take(4) { + assert_eq!(slot.row, 1, "cell {i}: row should be 1"); + assert_eq!(slot.col, i - 1, "cell {i}: col should be {}", i - 1); + assert_eq!(slot.colspan, 1); + assert_eq!(slot.rowspan, 1); + } + + // Cells 5..9: row 2, cols 0..3 + for (i, slot) in slots.iter().enumerate().skip(5).take(4) { + assert_eq!(slot.row, 2, "cell {i}: row should be 2"); + assert_eq!(slot.col, i - 5); + assert_eq!(slot.colspan, 1); + } + } + + #[test] + fn polygon_to_aabb_8elt() { + // Synthetic cell bbox 0 + let coords = vec![3.0, 2.0, 395.0, 2.0, 396.0, 59.0, 3.0, 59.0]; + let aabb = polygon_to_aabb(&coords).unwrap(); + assert_eq!(aabb, [3.0, 2.0, 396.0, 59.0]); + } + + #[test] + fn polygon_to_aabb_4elt() { + let coords = vec![5.0, 10.0, 50.0, 60.0]; + let aabb = polygon_to_aabb(&coords).unwrap(); + assert_eq!(aabb, [5.0, 10.0, 50.0, 60.0]); + } + + #[test] + fn polygon_to_aabb_4elt_unordered() { + // Caller may pass corners in any order; min/max should normalise. + let coords = vec![50.0, 60.0, 5.0, 10.0]; + let aabb = polygon_to_aabb(&coords).unwrap(); + assert_eq!(aabb, [5.0, 10.0, 50.0, 60.0]); + } + + #[test] + fn polygon_to_aabb_invalid_len() { + assert!(polygon_to_aabb(&[1.0, 2.0, 3.0]).is_none()); + assert!(polygon_to_aabb(&[1.0; 6]).is_none()); + assert!(polygon_to_aabb(&[]).is_none()); + } + + #[test] + fn synthetic_3x3_aabbs_inside_crop() { + // All 9 bboxes should produce valid (x1= 0.0 && aabb[2] <= 500.0, "bbox {i}: within crop"); + assert!(aabb[1] >= 0.0 && aabb[3] <= 200.0, "bbox {i}: within crop"); + } + } + + #[test] + fn parse_int_attr_basic() { + assert_eq!(parse_int_attr(" colspan=\"4\"", "colspan"), Some(4)); + assert_eq!(parse_int_attr(" rowspan=\"2\"", "rowspan"), Some(2)); + assert_eq!(parse_int_attr("colspan='3'", "colspan"), Some(3)); + assert_eq!(parse_int_attr(" colspan=\"4\"", "rowspan"), None); + assert_eq!(parse_int_attr(" class=\"foo\"", "colspan"), None); + } + + #[test] + fn parse_structure_rowspan_pushes_next_row_right() { + //
AB
C
", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(t) + .collect(); + + let slots = parse_structure(&tokens); + assert_eq!(slots.len(), 3); + assert_eq!((slots[0].row, slots[0].col), (0, 0)); + assert_eq!(slots[0].rowspan, 2); + assert_eq!((slots[1].row, slots[1].col), (0, 1)); + // C should be at (1, 1) because (1, 0) is occupied by A's rowspan. + assert_eq!((slots[2].row, slots[2].col), (1, 1)); + } + + #[test] + fn parse_structure_thead_marks_headers() { + // H1H2 + // D1D2 + let tokens: Vec = vec![ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(t) + .collect(); + + let slots = parse_structure(&tokens); + assert_eq!(slots.len(), 4); + assert!(slots[0].is_header && slots[1].is_header); + assert!(!slots[2].is_header && !slots[3].is_header); + } + + #[test] + fn parse_structure_th_outside_thead_still_header() { + // A row-header style: leading in tbody. + let tokens: Vec = vec![ + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(t) + .collect(); + + let slots = parse_structure(&tokens); + assert_eq!(slots.len(), 2); + assert!(slots[0].is_header); + assert!(!slots[1].is_header); + } + + #[test] + fn parse_structure_th_with_attrs() { + let tokens: Vec = vec![ + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(t) + .collect(); + let slots = parse_structure(&tokens); + assert_eq!(slots.len(), 1); + assert_eq!(slots[0].colspan, 2); + assert!(slots[0].is_header); + } + + #[test] + fn cells_to_markdown_synthetic_3x3() { + // Build the cells the parser would produce for the synthetic grid, + // and provide some sample text so we can sanity-check output. + let cells = vec![ + StructuredCell { + row: 0, + col: 0, + rowspan: 1, + colspan: 4, + is_header: false, + text: "Title".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + StructuredCell { + row: 1, + col: 0, + rowspan: 1, + colspan: 1, + is_header: false, + text: "a".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + StructuredCell { + row: 1, + col: 1, + rowspan: 1, + colspan: 1, + is_header: false, + text: "b".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + StructuredCell { + row: 1, + col: 2, + rowspan: 1, + colspan: 1, + is_header: false, + text: "c".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + StructuredCell { + row: 1, + col: 3, + rowspan: 1, + colspan: 1, + is_header: false, + text: "d".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + ]; + + let md = cells_to_markdown(&cells); + // Header row contains the spanning cell text in col 0 and pads to 4 cols. + // Absorbed-by-colspan positions render as empty cells (no padding). + assert!(md.starts_with("|Title||||\n"), "got: {md}"); + assert!(md.contains("|---|---|---|---|\n")); + assert!(md.contains("|a|b|c|d|\n")); + } + + #[test] + fn cells_to_markdown_escapes_pipes() { + let cells = vec![ + StructuredCell { + row: 0, + col: 0, + rowspan: 1, + colspan: 1, + is_header: false, + text: "a|b".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + StructuredCell { + row: 0, + col: 1, + rowspan: 1, + colspan: 1, + is_header: false, + text: "x".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }, + ]; + let md = cells_to_markdown(&cells); + assert!(md.contains("|a\\|b|x|")); + } + + #[test] + fn cells_to_markdown_collapses_whitespace_and_newlines() { + let cells = vec![StructuredCell { + row: 0, + col: 0, + rowspan: 1, + colspan: 1, + is_header: false, + text: "foo \n bar\tbaz".into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + }]; + let md = cells_to_markdown(&cells); + assert!(md.contains("|foo bar baz|")); + } + + fn cell(row: usize, col: usize, is_header: bool, text: &str) -> StructuredCell { + StructuredCell { + row, + col, + rowspan: 1, + colspan: 1, + is_header, + text: text.into(), + page_pt_bbox: [0.0, 0.0, 0.0, 0.0], + } + } + + #[test] + fn cells_to_markdown_separator_after_last_header_row() { + // Two-row header (a multi-row thead), then two body rows. Separator + // should land after row 1 (the LAST header row), not after row 0. + let cells = vec![ + cell(0, 0, true, "H0a"), + cell(0, 1, true, "H0b"), + cell(1, 0, true, "H1a"), + cell(1, 1, true, "H1b"), + cell(2, 0, false, "d0a"), + cell(2, 1, false, "d0b"), + cell(3, 0, false, "d1a"), + cell(3, 1, false, "d1b"), + ]; + let md = cells_to_markdown(&cells); + let expected = "|H0a|H0b|\n|H1a|H1b|\n|---|---|\n|d0a|d0b|\n|d1a|d1b|\n"; + assert_eq!(md, expected, "got: {md}"); + } + + #[test] + fn cells_to_markdown_separator_when_row_0_not_header() { + // Row 0 is not flagged as a header but row 1 is. Separator should + // follow row 1 (the header), demonstrating that we don't blindly + // emit after row 0. + let cells = vec![ + cell(0, 0, false, "x0a"), + cell(0, 1, false, "x0b"), + cell(1, 0, true, "Hdr1"), + cell(1, 1, true, "Hdr2"), + cell(2, 0, false, "data1"), + cell(2, 1, false, "data2"), + ]; + let md = cells_to_markdown(&cells); + // Confirm the separator is NOT after row 0. + assert!(!md.starts_with("|x0a|x0b|\n|---|"), "got: {md}"); + // Confirm it IS after row 1. + assert!( + md.contains("|Hdr1|Hdr2|\n|---|---|\n|data1|data2|"), + "got: {md}" + ); + } + + #[test] + fn cells_to_markdown_no_headers_falls_back_to_row_0() { + // No header cells at all — fallback: separator after row 0 so the + // output is still a valid markdown pipe-table. + let cells = vec![ + cell(0, 0, false, "a"), + cell(0, 1, false, "b"), + cell(1, 0, false, "c"), + cell(1, 1, false, "d"), + ]; + let md = cells_to_markdown(&cells); + assert_eq!(md, "|a|b|\n|---|---|\n|c|d|\n"); + } +} diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index 60389ef..c4835e9 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -1474,6 +1474,306 @@ fn test_bits_pilani_page8_table_detection() { assert!(!region.needs_ocr, "Page 8 table should still be detected"); } +// ========================================================================= +// extract_tables_with_structure_mem tests (TSR-aware path) +// ========================================================================= + +/// Build an 8-element 4-corner polygon `[x1,y1, x2,y1, x2,y2, x1,y2]` from +/// an axis-aligned rect — matches the format SLANet emits for cell bboxes. +fn poly(x1: f32, y1: f32, x2: f32, y2: f32) -> Vec { + vec![x1, y1, x2, y1, x2, y2, x1, y2] +} + +#[test] +fn test_extract_tables_with_structure_real_pdf_bits_pilani() { + use pdf_inspector::{extract_tables_with_structure_mem, TsrTableInput}; + // Hand-crafted TSR fixture targeting page 4 (0-indexed=3) of + // bits_pilani_feedback.pdf, which contains a clean tabular layout. + // + // We construct a 2×2 table: + // row 0 (header): "Department" "Core Courses" + // row 1 (data): "BIO" "8.23" + // + // The PDF page is US Letter (792pt tall). We render at 72 dpi so + // image-px maps 1:1 to PDF-pt — that lets us write cell bboxes in + // the same units as our hand-measured page-pt coordinates. + let buf = std::fs::read("tests/fixtures/bits_pilani_feedback.pdf").unwrap(); + + // The PDF page is A4 in points (≈595.44 × 841.68). The table sits in + // the upper part of the page; we crop a window large enough to enclose + // both rows we care about. + // + // Crop bounds in PDF points (top-left origin): + // x: 80..280, y: 170..240 + let crop = [80.0_f32, 170.0, 280.0, 240.0]; + let dpi = 72.0_f32; + + // Cell bboxes in CROP image-pixel space (= crop-relative PDF-pt at + // 72 dpi). The y ranges are tightened against neighbouring rows + // ("First Degree" above the header at native y=666.7, "Feedback Score" + // between the header and data rows at native y=640.9, "CE" below the + // BIO row at native y=591.1) so each cell only overlaps its target + // text item. + let cell_bboxes = vec![ + // Header row: y crop-relative (7, 18) → page-pt y (177, 188) + poly(10.0, 7.0, 100.0, 18.0), // "Department" (item at page-pt x=107.1) + poly(110.0, 7.0, 200.0, 18.0), // "Core Courses" (item at page-pt x=199.0) + // Data row: y crop-relative (35, 60) → page-pt y (205, 230) + poly(10.0, 35.0, 100.0, 60.0), // "BIO" (item at page-pt x=104.1) + poly(110.0, 35.0, 200.0, 60.0), // "8.23" (item at page-pt x=221.2) + ]; + + // Minimal SLANet-style token stream: a 2-row table with a thead and tbody. + let tokens: Vec = [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(String::from) + .collect(); + + let inputs = vec![TsrTableInput { + page: 3, + crop_pdf_pt_bbox: crop, + render_dpi: dpi, + structure_tokens: tokens, + cell_bboxes, + }]; + + let mds = extract_tables_with_structure_mem(&buf, &inputs).unwrap(); + assert_eq!(mds.len(), 1); + let md = &mds[0]; + + // Hand-written gold standard for the rendered markdown. + let expected = "|Department|Core Courses|\n|---|---|\n|BIO|8.23|\n"; + assert_eq!( + md, expected, + "structured-table markdown should match the gold standard exactly\nactual: {md}" + ); +} + +#[test] +fn test_extract_tables_with_structure_input_order_preserved() { + use pdf_inspector::{extract_tables_with_structure_mem, TsrTableInput}; + let buf = std::fs::read("tests/fixtures/bits_pilani_feedback.pdf").unwrap(); + + // Two inputs; both target the same page but with different shapes. + // We just need to confirm we get 2 outputs in the same order. + let make_input = |toks: Vec<&str>, cells: Vec>| TsrTableInput { + page: 3, + crop_pdf_pt_bbox: [80.0, 170.0, 280.0, 240.0], + render_dpi: 72.0, + structure_tokens: toks.into_iter().map(String::from).collect(), + cell_bboxes: cells, + }; + + let inputs = vec![ + make_input( + vec!["", "", "", "", "
"], + vec![poly(10.0, 35.0, 100.0, 60.0)], + ), + make_input( + vec!["", "", "", "", "
"], + vec![poly(110.0, 35.0, 200.0, 60.0)], + ), + ]; + + let mds = extract_tables_with_structure_mem(&buf, &inputs).unwrap(); + assert_eq!(mds.len(), 2); + assert!( + mds[0].contains("BIO"), + "input 0 should pull 'BIO': {}", + mds[0] + ); + assert!( + mds[1].contains("8.23"), + "input 1 should pull '8.23': {}", + mds[1] + ); +} + +#[test] +fn test_extract_tables_with_structure_out_of_range_page() { + use pdf_inspector::{extract_tables_with_structure_mem, TsrTableInput}; + let buf = std::fs::read("tests/fixtures/bits_pilani_feedback.pdf").unwrap(); + + let inputs = vec![TsrTableInput { + page: 9999, + crop_pdf_pt_bbox: [0.0, 0.0, 100.0, 100.0], + render_dpi: 72.0, + structure_tokens: vec![ + "".into(), + "".into(), + "".into(), + "".into(), + "
".into(), + ], + cell_bboxes: vec![poly(0.0, 0.0, 50.0, 50.0)], + }]; + + let mds = extract_tables_with_structure_mem(&buf, &inputs).unwrap(); + assert_eq!(mds.len(), 1); + assert!( + mds[0].is_empty(), + "out-of-range page should yield empty string" + ); +} + +#[test] +fn test_extract_tables_with_structure_not_a_pdf() { + use pdf_inspector::extract_tables_with_structure_mem; + let result = extract_tables_with_structure_mem(b"not a pdf", &[]); + assert!(result.is_err()); +} + +#[test] +fn test_extract_tables_with_structure_empty_inputs() { + use pdf_inspector::extract_tables_with_structure_mem; + let buf = std::fs::read("tests/fixtures/bits_pilani_feedback.pdf").unwrap(); + let mds = extract_tables_with_structure_mem(&buf, &[]).unwrap(); + assert!(mds.is_empty()); +} + +#[test] +fn test_extract_tables_with_structure_cells_real_pdf_bits_pilani() { + use pdf_inspector::{extract_tables_with_structure_cells_mem, TsrTableInput}; + // Same fixture as test_extract_tables_with_structure_real_pdf_bits_pilani + // but exercising the cell-level API. Verifies that callers receive + // structured per-cell metadata (row/col/spans/is_header/page_pt_bbox) + // alongside the extracted text. + let buf = std::fs::read("tests/fixtures/bits_pilani_feedback.pdf").unwrap(); + + let crop = [80.0_f32, 170.0, 280.0, 240.0]; + let dpi = 72.0_f32; + let cell_bboxes = vec![ + poly(10.0, 7.0, 100.0, 18.0), + poly(110.0, 7.0, 200.0, 18.0), + poly(10.0, 35.0, 100.0, 60.0), + poly(110.0, 35.0, 200.0, 60.0), + ]; + let tokens: Vec = [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(String::from) + .collect(); + + let inputs = vec![TsrTableInput { + page: 3, + crop_pdf_pt_bbox: crop, + render_dpi: dpi, + structure_tokens: tokens, + cell_bboxes, + }]; + + let cells_lists = extract_tables_with_structure_cells_mem(&buf, &inputs).unwrap(); + assert_eq!(cells_lists.len(), 1); + let cells = &cells_lists[0]; + assert_eq!(cells.len(), 4); + + // Header row: both cells flagged as headers (they were in /). + assert!(cells[0].is_header); + assert!(cells[1].is_header); + assert_eq!((cells[0].row, cells[0].col), (0, 0)); + assert_eq!((cells[1].row, cells[1].col), (0, 1)); + assert_eq!(cells[0].text, "Department"); + assert_eq!(cells[1].text, "Core Courses"); + + // Data row: not flagged as header. + assert!(!cells[2].is_header); + assert!(!cells[3].is_header); + assert_eq!((cells[2].row, cells[2].col), (1, 0)); + assert_eq!((cells[3].row, cells[3].col), (1, 1)); + assert_eq!(cells[2].text, "BIO"); + assert_eq!(cells[3].text, "8.23"); + + // Every cell carries a non-degenerate page-pt bbox. + for c in cells { + let [x1, y1, x2, y2] = c.page_pt_bbox; + assert!( + x1 < x2 && y1 < y2, + "cell bbox should be non-empty: {:?}", + c.page_pt_bbox + ); + } +} + +#[test] +fn test_extract_tables_with_structure_separator_after_thead() { + use pdf_inspector::{extract_tables_with_structure_mem, TsrTableInput}; + // Re-run the same 2x2 fixture but assert exact markdown output: with + // + headers, the separator should land after the header + // row (which is also row 0 here, so the gold-standard hasn't changed). + let buf = std::fs::read("tests/fixtures/bits_pilani_feedback.pdf").unwrap(); + + let crop = [80.0_f32, 170.0, 280.0, 240.0]; + let dpi = 72.0_f32; + let cell_bboxes = vec![ + poly(10.0, 7.0, 100.0, 18.0), + poly(110.0, 7.0, 200.0, 18.0), + poly(10.0, 35.0, 100.0, 60.0), + poly(110.0, 35.0, 200.0, 60.0), + ]; + let tokens: Vec = [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + ] + .into_iter() + .map(String::from) + .collect(); + + let mds = extract_tables_with_structure_mem( + &buf, + &[TsrTableInput { + page: 3, + crop_pdf_pt_bbox: crop, + render_dpi: dpi, + structure_tokens: tokens, + cell_bboxes, + }], + ) + .unwrap(); + assert_eq!(mds.len(), 1); + assert_eq!(mds[0], "|Department|Core Courses|\n|---|---|\n|BIO|8.23|\n"); +} + // ========================================================================= // extract_pages_markdown_mem tests // =========================================================================