214 lines
6.9 KiB
Python
214 lines
6.9 KiB
Python
 
|
"""End-to-end PyMuPDF PDF pipeline matching the architecture diagram."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import pdfplumber
|
||
|
|
|
||
|
|
from rag_cut.models import Block
|
||
|
|
from rag_cut.parsers.pdf.layout import (
|
||
|
|
LayoutRegion,
|
||
|
|
PageLayout,
|
||
|
|
_rect_overlap,
|
||
|
|
_sort_reading_order,
|
||
|
|
analyze_page_layout,
|
||
|
|
render_page,
|
||
|
|
)
|
||
|
|
from rag_cut.parsers.pdf.standardize import open_document
|
||
|
|
from rag_cut.parsers.pdf.tables import rows_to_markdown
|
||
|
|
from rag_cut.parsers.pdf.text_extract import extract_text_blocks
|
||
|
|
from rag_cut.parsers.pdf.visual_extract import extract_visual_blocks
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class _PageItem:
|
||
|
|
y0: float
|
||
|
|
x0: float
|
||
|
|
kind: str
|
||
|
|
block: Block
|
||
|
|
|
||
|
|
|
||
|
|
def _pdfplumber_tables(path: Path) -> dict[int, list[dict]]:
|
||
|
|
"""Extract tables per page with bbox from pdfplumber."""
|
||
|
|
tables_by_page: dict[int, list[dict]] = {}
|
||
|
|
try:
|
||
|
|
with pdfplumber.open(path) as pdf:
|
||
|
|
for i, page in enumerate(pdf.pages):
|
||
|
|
found = []
|
||
|
|
try:
|
||
|
|
for idx, table in enumerate(page.find_tables()):
|
||
|
|
rows = table.extract() or []
|
||
|
|
if not rows:
|
||
|
|
continue
|
||
|
|
bbox = table.bbox
|
||
|
|
found.append(
|
||
|
|
{
|
||
|
|
"rows": rows,
|
||
|
|
"table_index": idx,
|
||
|
|
"source": "pdfplumber",
|
||
|
|
"bbox": list(bbox) if bbox else None,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
except Exception:
|
||
|
|
rows_list = page.extract_tables() or []
|
||
|
|
for idx, rows in enumerate(rows_list):
|
||
|
|
found.append({"rows": rows, "table_index": idx, "source": "pdfplumber", "bbox": None})
|
||
|
|
if found:
|
||
|
|
tables_by_page[i] = found
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
return tables_by_page
|
||
|
|
|
||
|
|
|
||
|
|
def _inject_pdfplumber_tables(layout: PageLayout, plumber_tables: list[dict]) -> PageLayout:
|
||
|
|
if not plumber_tables:
|
||
|
|
return layout
|
||
|
|
if any(r.kind == "table" for r in layout.regions):
|
||
|
|
return layout
|
||
|
|
|
||
|
|
for item in plumber_tables:
|
||
|
|
rows = item.get("rows") or []
|
||
|
|
md = rows_to_markdown(rows)
|
||
|
|
if not md:
|
||
|
|
continue
|
||
|
|
bbox = item.get("bbox")
|
||
|
|
if bbox and len(bbox) == 4:
|
||
|
|
x0, y0, x1, y1 = bbox
|
||
|
|
else:
|
||
|
|
idx = int(item.get("table_index", 0))
|
||
|
|
x0, y0 = 0, layout.page_height * (idx + 1) / (len(plumber_tables) + 1)
|
||
|
|
x1, y1 = layout.page_width, layout.page_height * (idx + 2) / (len(plumber_tables) + 1)
|
||
|
|
layout.regions.append(
|
||
|
|
LayoutRegion(
|
||
|
|
x0=x0,
|
||
|
|
y0=y0,
|
||
|
|
x1=x1,
|
||
|
|
y1=y1,
|
||
|
|
kind="table",
|
||
|
|
data={
|
||
|
|
"rows": rows,
|
||
|
|
"table_index": item.get("table_index", 0),
|
||
|
|
"source": item.get("source", "pdfplumber"),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
)
|
||
|
|
ordered_boxes = _sort_reading_order(
|
||
|
|
[(r.x0, r.y0, r.x1, r.y1) for r in layout.regions],
|
||
|
|
layout.page_width,
|
||
|
|
)
|
||
|
|
order = {box: i for i, box in enumerate(ordered_boxes)}
|
||
|
|
layout.regions.sort(key=lambda r: order.get((r.x0, r.y0, r.x1, r.y1), len(order)))
|
||
|
|
return layout
|
||
|
|
|
||
|
|
|
||
|
|
def _bbox_key(bbox: list[float] | tuple[float, ...]) -> tuple[float, ...]:
|
||
|
|
return tuple(round(v, 1) for v in bbox)
|
||
|
|
|
||
|
|
|
||
|
|
def _region_box(region: LayoutRegion) -> tuple[float, float, float, float]:
|
||
|
|
return (region.x0, region.y0, region.x1, region.y1)
|
||
|
|
|
||
|
|
|
||
|
|
def _match_block_to_region(
|
||
|
|
region: LayoutRegion,
|
||
|
|
candidates: list[Block],
|
||
|
|
used: set[int],
|
||
|
|
) -> Block | None:
|
||
|
|
"""Match a layout region to a parsed block via exact bbox key or overlap."""
|
||
|
|
region_box = _region_box(region)
|
||
|
|
key = _bbox_key(region_box)
|
||
|
|
|
||
|
|
for block in candidates:
|
||
|
|
if id(block) in used:
|
||
|
|
continue
|
||
|
|
bbox = block.meta.get("bbox")
|
||
|
|
if bbox and _bbox_key(bbox) == key:
|
||
|
|
return block
|
||
|
|
|
||
|
|
best: Block | None = None
|
||
|
|
best_overlap = 0.35
|
||
|
|
for block in candidates:
|
||
|
|
if id(block) in used:
|
||
|
|
continue
|
||
|
|
bbox = block.meta.get("bbox")
|
||
|
|
if not bbox:
|
||
|
|
continue
|
||
|
|
overlap = _rect_overlap(region_box, tuple(bbox))
|
||
|
|
if overlap > best_overlap:
|
||
|
|
best_overlap = overlap
|
||
|
|
best = block
|
||
|
|
return best
|
||
|
|
|
||
|
|
|
||
|
|
def _merge_page_blocks(
|
||
|
|
text_blocks: list[Block],
|
||
|
|
visual_blocks: list[Block],
|
||
|
|
layout: PageLayout,
|
||
|
|
) -> list[Block]:
|
||
|
|
"""Merge text and visual branches in page reading order."""
|
||
|
|
used: set[int] = set()
|
||
|
|
merged: list[_PageItem] = []
|
||
|
|
|
||
|
|
for region in layout.regions:
|
||
|
|
if region.kind == "text":
|
||
|
|
candidates = text_blocks
|
||
|
|
elif region.kind in {"table", "image"}:
|
||
|
|
candidates = visual_blocks
|
||
|
|
else:
|
||
|
|
continue
|
||
|
|
|
||
|
|
block = _match_block_to_region(region, candidates, used)
|
||
|
|
if block is not None:
|
||
|
|
merged.append(_PageItem(y0=region.y0, x0=region.x0, kind=region.kind, block=block))
|
||
|
|
used.add(id(block))
|
||
|
|
|
||
|
|
leftovers: list[_PageItem] = []
|
||
|
|
for block in text_blocks + visual_blocks:
|
||
|
|
if id(block) not in used:
|
||
|
|
bbox = block.meta.get("bbox", [0, 0, 0, 0])
|
||
|
|
leftovers.append(_PageItem(y0=bbox[1], x0=bbox[0], kind=block.type.value, block=block))
|
||
|
|
|
||
|
|
merged.extend(sorted(leftovers, key=lambda it: (it.y0, it.x0)))
|
||
|
|
return [it.block for it in merged]
|
||
|
|
|
||
|
|
|
||
|
|
def parse_pdf(path: Path, assets_dir: Path) -> list[Block]:
|
||
|
|
"""
|
||
|
|
PDF pipeline:
|
||
|
|
1. 文档标准化
|
||
|
|
2. 页面渲染 + 版面分析
|
||
|
|
3. 文本块提取 ∥ 图片/表格区域裁剪 + OCR
|
||
|
|
4. 章节级语义切片（合并两路结果，标注章节上下文）
|
||
|
|
"""
|
||
|
|
assets_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
renders_dir = assets_dir / "pages"
|
||
|
|
renders_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
doc = open_document(path)
|
||
|
|
plumber_tables = _pdfplumber_tables(path)
|
||
|
|
all_blocks: list[Block] = []
|
||
|
|
chapter_title: str | None = None
|
||
|
|
img_counter = 0
|
||
|
|
|
||
|
|
try:
|
||
|
|
for page_index, page in enumerate(doc):
|
||
|
|
layout = analyze_page_layout(page, page_index)
|
||
|
|
layout = _inject_pdfplumber_tables(layout, plumber_tables.get(page_index, []))
|
||
|
|
|
||
|
|
page_render = render_page(page)
|
||
|
|
render_path = renders_dir / f"page{page_index + 1}.png"
|
||
|
|
page_render.save(str(render_path))
|
||
|
|
|
||
|
|
text_blocks, chapter_title = extract_text_blocks(layout, chapter_title)
|
||
|
|
visual_blocks, img_counter = extract_visual_blocks(
|
||
|
|
page, layout, doc, assets_dir, chapter_title, img_counter
|
||
|
|
)
|
||
|
|
page_blocks = _merge_page_blocks(text_blocks, visual_blocks, layout)
|
||
|
|
all_blocks.extend(page_blocks)
|
||
|
|
finally:
|
||
|
|
doc.close()
|
||
|
|
|
||
|
|
return all_blocks
|