350 lines
12 KiB
Python
350 lines
12 KiB
Python
"""Page rendering and layout analysis."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
import fitz
|
||
|
||
from rag_cut.parsers.pdf.tables import looks_like_table
|
||
|
||
|
||
_NUMBERED_LINE_RE = re.compile(
|
||
r"^\s*(\d+(?:\.\d+)*)(?:\.|．)?\s*([A-Za-z0-9\u4e00-\u9fff].{2,})$"
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class TextLine:
|
||
x0: float
|
||
y0: float
|
||
x1: float
|
||
y1: float
|
||
text: str
|
||
font_size: float
|
||
|
||
|
||
@dataclass
|
||
class LayoutRegion:
|
||
x0: float
|
||
y0: float
|
||
x1: float
|
||
y1: float
|
||
kind: str # text | table | image
|
||
data: dict[str, Any] = field(default_factory=dict)
|
||
|
||
|
||
@dataclass
|
||
class PageLayout:
|
||
page_index: int
|
||
page_width: float
|
||
page_height: float
|
||
body_font_size: float
|
||
regions: list[LayoutRegion] = field(default_factory=list)
|
||
|
||
|
||
def render_page(page: fitz.Page, zoom: float = 2.0) -> fitz.Pixmap:
|
||
"""Render page to pixmap for region cropping and OCR."""
|
||
return page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
|
||
|
||
|
||
def _median_body_size(text_dict: dict) -> float:
|
||
sizes: list[float] = []
|
||
for block in text_dict.get("blocks", []):
|
||
if block.get("type") != 0:
|
||
continue
|
||
for line in block.get("lines", []):
|
||
for span in line.get("spans", []):
|
||
sizes.append(span.get("size", 12))
|
||
return sorted(sizes)[len(sizes) // 2] if sizes else 12
|
||
|
||
|
||
def _collect_text_lines(page: fitz.Page, body_size: float) -> list[TextLine]:
|
||
text_dict = page.get_text("dict")
|
||
lines: list[TextLine] = []
|
||
for block in text_dict.get("blocks", []):
|
||
if block.get("type") != 0:
|
||
continue
|
||
for line in block.get("lines", []):
|
||
text = "".join(span.get("text", "") for span in line.get("spans", [])).strip()
|
||
if not text:
|
||
continue
|
||
x0, y0, x1, y1 = line["bbox"]
|
||
max_size = body_size
|
||
for span in line.get("spans", []):
|
||
max_size = max(max_size, span.get("size", body_size))
|
||
lines.append(TextLine(x0=x0, y0=y0, x1=x1, y1=y1, text=text, font_size=max_size))
|
||
return lines
|
||
|
||
|
||
def _detect_columns(
|
||
items: list[tuple[float, float, float, float]],
|
||
page_width: float,
|
||
) -> list[tuple[float, float, float, float]]:
|
||
if not items:
|
||
return items
|
||
mid = page_width / 2
|
||
left = [b for b in items if (b[0] + b[2]) / 2 < mid]
|
||
right = [b for b in items if (b[0] + b[2]) / 2 >= mid]
|
||
if len(left) >= 2 and len(right) >= 2:
|
||
left.sort(key=lambda b: (b[1], b[0]))
|
||
right.sort(key=lambda b: (b[1], b[0]))
|
||
return left + right
|
||
return sorted(items, key=lambda b: (b[1], b[0]))
|
||
|
||
|
||
def _sort_reading_order(
|
||
items: list[tuple[float, float, float, float]],
|
||
page_width: float,
|
||
) -> list[tuple[float, float, float, float]]:
|
||
"""Sort page boxes in a human reading order, preserving two-column flows."""
|
||
if not items:
|
||
return []
|
||
|
||
mid = page_width / 2
|
||
full_width: list[tuple[float, float, float, float]] = []
|
||
column_items: list[tuple[float, float, float, float]] = []
|
||
for box in items:
|
||
width = box[2] - box[0]
|
||
spans_mid = box[0] < mid < box[2]
|
||
if width >= page_width * 0.60 or (spans_mid and width >= page_width * 0.35):
|
||
full_width.append(box)
|
||
else:
|
||
column_items.append(box)
|
||
|
||
left = [b for b in column_items if (b[0] + b[2]) / 2 < mid]
|
||
right = [b for b in column_items if (b[0] + b[2]) / 2 >= mid]
|
||
if len(left) < 2 or len(right) < 2:
|
||
return sorted(items, key=lambda b: (b[1], b[0]))
|
||
|
||
full_width.sort(key=lambda b: (b[1], b[0]))
|
||
ordered: list[tuple[float, float, float, float]] = []
|
||
segment_top = float("-inf")
|
||
|
||
def add_columns_between(top: float, bottom: float) -> None:
|
||
segment = [b for b in column_items if b[1] >= top and b[1] < bottom]
|
||
segment_left = sorted([b for b in segment if (b[0] + b[2]) / 2 < mid], key=lambda b: (b[1], b[0]))
|
||
segment_right = sorted([b for b in segment if (b[0] + b[2]) / 2 >= mid], key=lambda b: (b[1], b[0]))
|
||
ordered.extend(segment_left)
|
||
ordered.extend(segment_right)
|
||
|
||
for box in full_width:
|
||
add_columns_between(segment_top, box[1])
|
||
ordered.append(box)
|
||
segment_top = box[3]
|
||
add_columns_between(segment_top, float("inf"))
|
||
|
||
seen: set[tuple[float, float, float, float]] = set(ordered)
|
||
ordered.extend(b for b in sorted(items, key=lambda b: (b[1], b[0])) if b not in seen)
|
||
return ordered
|
||
|
||
|
||
def _rect_overlap(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float:
|
||
x0 = max(a[0], b[0])
|
||
y0 = max(a[1], b[1])
|
||
x1 = min(a[2], b[2])
|
||
y1 = min(a[3], b[3])
|
||
if x1 <= x0 or y1 <= y0:
|
||
return 0.0
|
||
inter = (x1 - x0) * (y1 - y0)
|
||
area_a = max((a[2] - a[0]) * (a[3] - a[1]), 1e-6)
|
||
return inter / area_a
|
||
|
||
|
||
def _center_inside(inner: tuple[float, float, float, float], outer: tuple[float, float, float, float]) -> bool:
|
||
cx = (inner[0] + inner[2]) / 2
|
||
cy = (inner[1] + inner[3]) / 2
|
||
return outer[0] <= cx <= outer[2] and outer[1] <= cy <= outer[3]
|
||
|
||
|
||
def _horizontal_overlap_ratio(a: TextLine, b: TextLine) -> float:
|
||
overlap = min(a.x1, b.x1) - max(a.x0, b.x0)
|
||
if overlap <= 0:
|
||
return 0.0
|
||
return overlap / max(min(a.x1 - a.x0, b.x1 - b.x0), 1e-6)
|
||
|
||
|
||
def _line_on_image(
|
||
line: TextLine,
|
||
image_boxes: list[tuple[float, float, float, float]],
|
||
large_image_boxes: list[tuple[float, float, float, float]],
|
||
) -> bool:
|
||
# Never drop numbered section titles even if they sit near a screenshot.
|
||
if _NUMBERED_LINE_RE.match(line.text.strip()):
|
||
return False
|
||
bbox = (line.x0, line.y0, line.x1, line.y1)
|
||
# Large background/screenshot: require near-full coverage of the line, not
|
||
# merely center-inside (which wiped text sitting in margins of wide figures).
|
||
if large_image_boxes and any(
|
||
_center_inside(bbox, box) and _rect_overlap(bbox, box) >= 0.85 for box in large_image_boxes
|
||
):
|
||
return True
|
||
return any(_rect_overlap(bbox, box) >= 0.55 for box in image_boxes)
|
||
|
||
|
||
def _area(box: tuple[float, float, float, float]) -> float:
|
||
return max(box[2] - box[0], 0.0) * max(box[3] - box[1], 0.0)
|
||
|
||
|
||
def _filter_nested_image_regions(regions: list[LayoutRegion]) -> list[LayoutRegion]:
|
||
"""Drop image fragments that are already contained in a larger screenshot."""
|
||
result: list[LayoutRegion] = []
|
||
boxes = [(r.x0, r.y0, r.x1, r.y1) for r in regions]
|
||
for region, box in zip(regions, boxes):
|
||
box_area = _area(box)
|
||
nested = False
|
||
for other in boxes:
|
||
other_area = _area(other)
|
||
if other == box or other_area <= box_area * 1.5:
|
||
continue
|
||
if _center_inside(box, other) and _rect_overlap(box, other) >= 0.85:
|
||
nested = True
|
||
break
|
||
if not nested:
|
||
result.append(region)
|
||
return result
|
||
|
||
|
||
def _merge_lines_to_paragraphs(lines: list[TextLine], page_width: float) -> list[TextLine]:
|
||
if not lines:
|
||
return []
|
||
|
||
ordered_boxes = _sort_reading_order([(ln.x0, ln.y0, ln.x1, ln.y1) for ln in lines], page_width)
|
||
order = {(b[0], b[1], b[2], b[3]): i for i, b in enumerate(ordered_boxes)}
|
||
ordered = sorted(lines, key=lambda ln: order.get((ln.x0, ln.y0, ln.x1, ln.y1), (ln.y0, ln.x0)))
|
||
|
||
paragraphs: list[TextLine] = []
|
||
current = ordered[0]
|
||
for nxt in ordered[1:]:
|
||
vgap = nxt.y0 - current.y1
|
||
line_h = max(current.y1 - current.y0, nxt.y1 - nxt.y0, 8.0)
|
||
size_gap = abs(current.font_size - nxt.font_size)
|
||
size_ratio = max(current.font_size, nxt.font_size) / max(
|
||
min(current.font_size, nxt.font_size), 1e-6
|
||
)
|
||
# Keep title vs body separate: small absolute gap is enough when ratio is large.
|
||
same_style = size_gap <= 1.0 and size_ratio <= 1.15
|
||
either_numbered = bool(
|
||
_NUMBERED_LINE_RE.match(current.text.strip()) or _NUMBERED_LINE_RE.match(nxt.text.strip())
|
||
)
|
||
if (
|
||
same_style
|
||
and not either_numbered
|
||
and vgap <= line_h * 2.2
|
||
and _horizontal_overlap_ratio(current, nxt) >= 0.25
|
||
):
|
||
joiner = "" if current.text.endswith("-") or current.text.endswith(" ") else " "
|
||
current = TextLine(
|
||
x0=min(current.x0, nxt.x0),
|
||
y0=current.y0,
|
||
x1=max(current.x1, nxt.x1),
|
||
y1=nxt.y1,
|
||
text=f"{current.text}{joiner}{nxt.text}",
|
||
font_size=max(current.font_size, nxt.font_size),
|
||
)
|
||
else:
|
||
paragraphs.append(current)
|
||
current = nxt
|
||
paragraphs.append(current)
|
||
return paragraphs
|
||
|
||
|
||
def analyze_page_layout(page: fitz.Page, page_index: int) -> PageLayout:
|
||
"""Layout analysis: text paragraphs, table regions, image regions."""
|
||
text_dict = page.get_text("dict")
|
||
body_size = _median_body_size(text_dict)
|
||
page_width = page.rect.width
|
||
page_height = page.rect.height
|
||
page_area = page_width * page_height
|
||
|
||
table_regions: list[LayoutRegion] = []
|
||
try:
|
||
table_finder = page.find_tables()
|
||
for idx, table in enumerate(table_finder.tables):
|
||
bbox = table.bbox
|
||
rows = table.extract() or []
|
||
if not looks_like_table(rows):
|
||
continue
|
||
table_regions.append(
|
||
LayoutRegion(
|
||
x0=bbox[0],
|
||
y0=bbox[1],
|
||
x1=bbox[2],
|
||
y1=bbox[3],
|
||
kind="table",
|
||
data={"rows": rows, "table_index": idx, "source": "pymupdf"},
|
||
)
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
image_regions: list[LayoutRegion] = []
|
||
for img_info in page.get_images(full=True):
|
||
xref = img_info[0]
|
||
try:
|
||
rects = page.get_image_rects(xref)
|
||
except Exception:
|
||
rects = []
|
||
if not rects:
|
||
continue
|
||
for rect_idx, rect in enumerate(rects):
|
||
image_regions.append(
|
||
LayoutRegion(
|
||
x0=rect.x0,
|
||
y0=rect.y0,
|
||
x1=rect.x1,
|
||
y1=rect.y1,
|
||
kind="image",
|
||
data={"xref": xref, "rect_index": rect_idx},
|
||
)
|
||
)
|
||
image_regions = _filter_nested_image_regions(image_regions)
|
||
|
||
image_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in image_regions]
|
||
large_image_boxes = [
|
||
box
|
||
for box in image_boxes
|
||
if (box[2] - box[0]) * (box[3] - box[1]) >= page_area * 0.12
|
||
]
|
||
table_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in table_regions]
|
||
|
||
raw_lines = _collect_text_lines(page, body_size)
|
||
filtered_lines = [
|
||
ln
|
||
for ln in raw_lines
|
||
if not _line_on_image(ln, image_boxes, large_image_boxes)
|
||
and not any(_rect_overlap((ln.x0, ln.y0, ln.x1, ln.y1), box) >= 0.55 for box in table_boxes)
|
||
]
|
||
paragraphs = _merge_lines_to_paragraphs(filtered_lines, page_width)
|
||
|
||
text_regions: list[LayoutRegion] = []
|
||
for para in paragraphs:
|
||
text_regions.append(
|
||
LayoutRegion(
|
||
x0=para.x0,
|
||
y0=para.y0,
|
||
x1=para.x1,
|
||
y1=para.y1,
|
||
kind="text",
|
||
data={"text": para.text, "font_size": para.font_size},
|
||
)
|
||
)
|
||
|
||
regions = text_regions + table_regions + image_regions
|
||
ordered_boxes = _sort_reading_order(
|
||
[(r.x0, r.y0, r.x1, r.y1) for r in regions],
|
||
page_width,
|
||
)
|
||
order = {box: i for i, box in enumerate(ordered_boxes)}
|
||
regions.sort(key=lambda r: order.get((r.x0, r.y0, r.x1, r.y1), len(order)))
|
||
|
||
return PageLayout(
|
||
page_index=page_index,
|
||
page_width=page_width,
|
||
page_height=page_height,
|
||
body_font_size=body_size,
|
||
regions=regions,
|
||
)
|