Files
pdf/gateway/app/services/convert/layout/ml_merge.py
T

397 lines
15 KiB
Python

"""Merge ML layout regions with heuristic glyph lines into IDM blocks."""
from __future__ import annotations
from app.services.convert.idm.model import BBox, Block, BlockType, PageKind
from app.services.convert.layout.blocks import (
estimate_body_size,
lines_to_block,
)
from app.services.convert.layout.glyphs import Line
from app.services.convert.layout.ml_regions import Region
from app.services.convert.layout.paragraphs import group_lines as group_paragraph_lines
from app.services.convert.layout.paragraphs import lexical_hyphen_pairs
from app.services.convert.layout.paragraphs import measure as measure_page
from app.services.convert.layout.reading_order import order_lines
from app.services.convert.layout.tables import coalesce_lines_by_y, extract_tables
from app.services.convert.options import get_options
_LABEL_TO_BLOCK = {
"title": BlockType.heading,
"text": BlockType.paragraph,
"list": BlockType.list_item,
"table": BlockType.table,
"figure": BlockType.figure,
"header": BlockType.header,
"footer": BlockType.footer,
}
def _line_center(line: Line) -> tuple[float, float]:
return ((line.x0 + line.x1) / 2.0, line.y)
def _contains(bbox: BBox, x: float, y: float, pad: float = 2.0) -> bool:
return (
bbox.x - pad <= x <= bbox.x + bbox.w + pad
and bbox.y - pad <= y <= bbox.y + bbox.h + pad
)
def _overlap_line(bbox: BBox, line: Line, *, y_pad: float = 0.0) -> bool:
cx, cy = _line_center(line)
if _contains(bbox, cx, cy, pad=max(2.0, y_pad)):
return True
# horizontal overlap + y near region vertical span
lx0, lx1 = line.x0, line.x1
y0, y1 = bbox.y - y_pad, bbox.y + bbox.h + y_pad
return not (lx1 < bbox.x or lx0 > bbox.x + bbox.w) and y0 <= cy <= y1
def lines_in_region(lines: list[Line], region: Region) -> list[Line]:
return [ln for ln in lines if _overlap_line(region.bbox_pdf, ln, y_pad=8.0)]
def _paragraph_blocks(
lines: list[Line],
start_order: int,
body: float,
page_width: float,
*,
force_type: BlockType | None = None,
) -> list[Block]:
"""Group ordered lines into paragraph blocks.
The ML path previously emitted one block per visual line, so a document
routed through layout detection came out with every line as its own
``<w:p>`` even though the deterministic path reconstructed paragraphs.
Both paths now share :mod:`layout.paragraphs`.
``force_type`` applies a region's label to the resulting blocks. Headings
and list items are never merged across lines: a heading region holds one
heading, and each list item is its own block.
"""
if not lines:
return []
unmergeable = force_type in (
BlockType.heading,
BlockType.list_item,
BlockType.header,
BlockType.footer,
)
metrics = measure_page(lines, body_size=body, page_width=page_width)
pairs = lexical_hyphen_pairs("\n".join(ln.text for ln in lines))
runs = [[ln] for ln in lines] if unmergeable else group_paragraph_lines(lines, metrics)
out: list[Block] = []
order = start_order
for run in runs:
block = lines_to_block(run, order, body, lexical_pairs=pairs)
if force_type is BlockType.heading:
block.type = BlockType.heading
block.level = block.level or 1
elif force_type is BlockType.paragraph:
# Adobe-style: a text region is body, not a size-based H1.
if block.type == BlockType.heading:
block.type = BlockType.paragraph
block.level = 0
elif force_type is not None:
block.type = force_type
out.append(block)
order += 1
return out
def _structure_rulings(
reg: Region,
*,
page_png: bytes | None,
page_width: float,
page_height: float,
page_index: int,
vertical_rulings: list[float] | None,
warnings: list[str],
) -> list[float] | None:
"""Column separators for one ML table region, or None to keep the heuristics.
Every failure path here returns None and leaves ``vertical_rulings``
untouched: the optional structure model may sharpen a grid, never break one.
"""
from app.services.convert.layout import ml_table_structure as mts
if not mts.table_structure_enabled() or not page_png or not page_height:
return None
try:
structure, err = mts.structure_for_region(
page_png, reg.bbox_pdf, page_width, page_height
)
except Exception as exc: # pragma: no cover - defensive; hook must never raise
warnings.append(f"Page {page_index + 1}: table_structure_fallback=heuristic ({exc})")
return None
if structure is None:
if err:
warnings.append(f"Page {page_index + 1}: table_structure_fallback=heuristic ({err})")
return None
merged = mts.merge_rulings(vertical_rulings, structure.column_x)
warnings.append(
f"Page {page_index + 1}: table_structure=onnx cols={structure.columns} "
f"score={structure.score:.2f}."
)
return merged
def merge_regions_with_lines(
lines: list[Line],
regions: list[Region],
*,
page_width: float,
page_index: int,
kind: PageKind,
warnings: list[str],
vertical_rulings: list[float] | None = None,
page_rects: list | None = None,
path_ops: list[dict] | None = None,
page_png: bytes | None = None,
page_height: float | None = None,
) -> list[Block]:
"""
Build blocks using ML region order + heuristic tables.
Order of operations (best practice fusion):
1. Coalesce gutter-split column fragments into logical rows.
2. Extract grids from explicit ML *table* regions (relaxed / ml_hinted).
3. Run page-level heuristic extract_tables on leftovers (catches pipe
tables that ML mislabeled as per-row ``text`` regions).
4. Type remaining lines via non-table ML regions (heading/list/header…).
5. Leftover lines → paragraphs.
"""
if not regions:
return []
working = coalesce_lines_by_y(lines) if lines else []
if not any(ln.text.strip() for ln in working):
# A full-page bitmap: the detector found rectangles, but there is no
# text layer for them to describe. Emitting region blocks here — empty
# tables in particular — puts structure on the page that nothing has
# read yet, and the OCR rebuild that follows has the actual detection
# boxes and assembles far better grids from them. The regions are not
# wasted: the OCR pass still uses the table rectangles as crop hints.
warnings.append(
f"Page {page_index + 1}: layout_ml regions on a raster page; "
"deferring structure to OCR."
)
return []
assigned: set[int] = set()
blocks: list[Block] = []
order = 0
body = estimate_body_size(working) if working else 12.0
detect_tables = get_options().detect_tables
table_regs = sorted(
[r for r in regions if r.label == "table"] if detect_tables else [],
key=lambda r: r.reading_index,
)
other_regs = sorted(
[r for r in regions if r.label != "table"],
key=lambda r: r.reading_index,
)
# --- Pass A: explicit ML table regions (relaxed grid builder) ---
for reg in table_regs:
idxs = [
i
for i, ln in enumerate(working)
if i not in assigned and _overlap_line(reg.bbox_pdf, ln, y_pad=10.0)
]
# Expand vertically: include nearby unassigned lines inside region Y span
y0, y1 = reg.bbox_pdf.y - 6.0, reg.bbox_pdf.y + reg.bbox_pdf.h + 6.0
for i, ln in enumerate(working):
if i in assigned or i in idxs:
continue
if y0 <= ln.y <= y1 and not (ln.x1 < reg.bbox_pdf.x or ln.x0 > reg.bbox_pdf.x + reg.bbox_pdf.w):
idxs.append(i)
idxs = sorted(set(idxs))
reg_lines = [working[i] for i in idxs]
for i in idxs:
assigned.add(i)
if not reg_lines:
# An empty table is not a table. It reached the document as a
# bordered box around nothing, and it displaced the grid that OCR
# would have built from the same rectangle.
warnings.append(
f"Page {page_index + 1}: ML table region with no text lines; skipped."
)
continue
# Optional structure model: column separators for THIS region only.
# None (disabled, no weights, bad output) leaves the heuristics alone.
region_rulings = _structure_rulings(
reg,
page_png=page_png,
page_width=page_width,
page_height=float(page_height or 0.0),
page_index=page_index,
vertical_rulings=vertical_rulings,
warnings=warnings,
)
table_blocks, remaining, conf = extract_tables(
order_lines(reg_lines, page_width),
start_order=order,
vertical_rulings=region_rulings if region_rulings is not None else vertical_rulings,
page_rects=page_rects,
path_ops=path_ops,
ml_hinted=True,
)
# Require a real multi-row grid; weak 1-row hits release lines for Pass B.
strong = bool(
table_blocks
and table_blocks[0].cells
and len(table_blocks[0].cells) >= 2
and max(len(r) for r in table_blocks[0].cells) >= 2
)
if strong:
for tb in table_blocks:
tb.reading_order = order
tb.bbox = reg.bbox_pdf
if tb.table_confidence < 0.5:
tb.table_confidence = max(tb.table_confidence, float(reg.score or 0.5))
blocks.append(tb)
order += 1
para_blocks = _paragraph_blocks(remaining, order, body, page_width)
blocks.extend(para_blocks)
order += len(para_blocks)
warnings.append(
f"Page {page_index + 1}: ML table->grid rows={len(table_blocks[0].cells)} "
f"cols={len(table_blocks[0].cells[0])} conf={conf:.2f}."
)
else:
# Release region lines so page-level heuristics can still build a grid
for i in idxs:
assigned.discard(i)
warnings.append(
f"Page {page_index + 1}: ML table region grid_failed conf={conf:.2f}; "
"deferring to heuristic."
)
# --- Pass B: page-level heuristic tables on unassigned lines ---
# Fixes ML that labels each table row as a separate "text" region.
leftover = [working[i] for i in range(len(working)) if i not in assigned]
if leftover and detect_tables:
table_blocks, remaining, conf = extract_tables(
order_lines(leftover, page_width),
start_order=order,
vertical_rulings=vertical_rulings,
page_rects=page_rects,
path_ops=path_ops,
ml_hinted=False,
)
if table_blocks and conf >= 0.60 and table_blocks[0].cells:
remaining_ids = {id(ln) for ln in remaining}
leftover_ids = {id(ln) for ln in leftover}
for i, ln in enumerate(working):
if id(ln) in leftover_ids and id(ln) not in remaining_ids:
assigned.add(i)
for tb in table_blocks:
tb.reading_order = order
blocks.append(tb)
order += 1
warnings.append(
f"Page {page_index + 1}: heuristic table under ML regions "
f"rows={len(table_blocks[0].cells)} conf={conf:.2f}."
)
# --- Pass C: non-table ML regions for typing ---
for reg in other_regs:
idxs = [
i
for i, ln in enumerate(working)
if i not in assigned and _overlap_line(reg.bbox_pdf, ln, y_pad=4.0)
]
reg_lines = [working[i] for i in idxs]
for i in idxs:
assigned.add(i)
btype = _LABEL_TO_BLOCK.get(reg.label, BlockType.paragraph)
if not reg_lines:
if btype in (BlockType.figure, BlockType.header, BlockType.footer):
blocks.append(
Block(type=btype, text="", bbox=reg.bbox_pdf, reading_order=order)
)
order += 1
continue
region_blocks = _paragraph_blocks(
order_lines(reg_lines, page_width),
order,
body,
page_width,
force_type=btype,
)
blocks.extend(region_blocks)
order += len(region_blocks)
# --- Pass D: true leftovers ---
leftover = [working[i] for i in range(len(working)) if i not in assigned]
if leftover:
ordered = order_lines(leftover, page_width)
table_blocks, remaining, conf = ([], ordered, 0.0)
if detect_tables:
table_blocks, remaining, conf = extract_tables(
ordered,
start_order=order,
vertical_rulings=vertical_rulings,
page_rects=page_rects,
path_ops=path_ops,
)
if table_blocks and conf >= 0.65 and table_blocks[0].cells:
for tb in table_blocks:
tb.reading_order = order
blocks.append(tb)
order += 1
para_blocks = _paragraph_blocks(remaining, order, body, page_width)
blocks.extend(para_blocks)
order += len(para_blocks)
else:
para_blocks = _paragraph_blocks(ordered, order, body, page_width)
blocks.extend(para_blocks)
order += len(para_blocks)
if kind == PageKind.blank and not blocks:
warnings.append(f"Page {page_index + 1}: ML regions produced no blocks.")
# Order blocks by page reading order (XY-cut of text lines) and geometry,
# rather than the order of detector passes (which previously put tables before titles).
ordered_page_lines = order_lines(lines, page_width, page_height=float(page_height or 792.0))
def _block_sort_key(b: Block) -> tuple:
if b.type == BlockType.header:
return (0, 0, 0.0, 0.0)
if b.type == BlockType.footer:
return (2, 0, 0.0, 0.0)
box = b.bbox
if box is not None:
inside = [
i
for i, ln in enumerate(ordered_page_lines)
if box.x - 3.0 <= (ln.x0 + ln.x1) / 2.0 <= box.x + box.w + 3.0
and box.y - 3.0 <= ln.y <= box.y + box.h + 3.0
]
if inside:
return (1, 0, min(inside), 0.0)
top_y = box.y + box.h
return (1, 1, -top_y, box.x)
return (1, 2, b.reading_order, 0.0)
blocks.sort(key=_block_sort_key)
for i, b in enumerate(blocks):
b.reading_order = i
return blocks
def ml_header_footer_hints(regions: list[Region]) -> tuple[list[BBox], list[BBox]]:
headers = [r.bbox_pdf for r in regions if r.label == "header"]
footers = [r.bbox_pdf for r in regions if r.label == "footer"]
return headers, footers