陈辅元
Cursor
@@ -0,0 +1,349 @@
|
||||
"""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,
|
||||
)
|
||||
Reference in New Issue
Block a user