397 lines
15 KiB
Python
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
|