105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
"""Text block extraction from analyzed layout."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
from rag_cut.models import Block, BlockType
|
||
from rag_cut.parsers.pdf.layout import LayoutRegion, PageLayout
|
||
|
||
_UI_LABEL_RE = re.compile(r"^[\d\s\W]{0,6}[\u4e00-\u9fff]{1,6}$")
|
||
# "10.上升三角形態" / "1.2 ACCOUNT STATUS" (space after number optional)
|
||
_NUMBERED_HEADING_RE = re.compile(
|
||
r"^\s*(\d+(?:\.\d+)*)(?:\.|．)?\s*([A-Za-z0-9\u4e00-\u9fff][A-Za-z0-9\u4e00-\u9fff&/ \-_:：]{2,})\s*$"
|
||
)
|
||
|
||
# Title+body merges must not become headings; real section titles stay shorter.
|
||
MAX_HEADING_CHARS = 100
|
||
|
||
|
||
def _numbered_heading_level(text: str) -> int | None:
|
||
match = _NUMBERED_HEADING_RE.match(text.strip())
|
||
if not match or len(match.group(2).strip()) < 3:
|
||
return None
|
||
return match.group(1).count(".") + 1
|
||
|
||
|
||
def _looks_like_ui_label(text: str) -> bool:
|
||
"""True for short UI chips; exclude numbered / CJK section titles."""
|
||
stripped = text.strip()
|
||
if _numbered_heading_level(stripped) is not None:
|
||
return False
|
||
if re.match(r"^\d+(?:\.\d+)*(?:\.|．)", stripped):
|
||
return False
|
||
# Short Chinese section banners like 「形態指標」are not toolbar labels.
|
||
if re.fullmatch(r"[\u4e00-\u9fff]{2,12}", stripped):
|
||
return False
|
||
return bool(_UI_LABEL_RE.match(stripped))
|
||
|
||
|
||
def _font_heading_level(size: float, body_size: float, text: str) -> int | None:
|
||
stripped = text.strip()
|
||
numbered = _numbered_heading_level(stripped)
|
||
if numbered and len(stripped) <= MAX_HEADING_CHARS:
|
||
return numbered
|
||
if len(stripped) < 4 or len(stripped) > MAX_HEADING_CHARS:
|
||
return None
|
||
if _looks_like_ui_label(stripped):
|
||
return None
|
||
|
||
# Compact CJK section titles such as 「形態指標」.
|
||
if (
|
||
size >= body_size + 3
|
||
and re.fullmatch(r"[\u4e00-\u9fff]{2,12}", stripped)
|
||
and not _numbered_heading_level(stripped)
|
||
):
|
||
return 1
|
||
|
||
if size >= body_size + 6:
|
||
return 1 if len(stripped) >= 10 else 2
|
||
if size >= body_size + 3:
|
||
return 2 if len(stripped) >= 8 else 3
|
||
if size >= body_size + 1.5:
|
||
return 3
|
||
return None
|
||
|
||
|
||
def extract_text_blocks(layout: PageLayout, chapter_title: str | None = None) -> tuple[list[Block], str | None]:
|
||
"""Extract heading/paragraph blocks; update chapter title when headings appear."""
|
||
blocks: list[Block] = []
|
||
current_chapter = chapter_title
|
||
|
||
for region in layout.regions:
|
||
if region.kind != "text":
|
||
continue
|
||
|
||
text = region.data.get("text", "").strip()
|
||
if not text:
|
||
continue
|
||
|
||
font_size = float(region.data.get("font_size", layout.body_font_size))
|
||
level = _font_heading_level(font_size, layout.body_font_size, text)
|
||
page_no = layout.page_index + 1
|
||
meta = {
|
||
"page": page_no,
|
||
"bbox": [region.x0, region.y0, region.x1, region.y1],
|
||
"font_size": font_size,
|
||
"body_font_size": layout.body_font_size,
|
||
"page_height": layout.page_height,
|
||
}
|
||
if current_chapter:
|
||
meta["chapter"] = current_chapter
|
||
|
||
if level:
|
||
if level <= 2:
|
||
current_chapter = text
|
||
if current_chapter:
|
||
meta["chapter"] = current_chapter
|
||
blk = Block(type=BlockType.HEADING, text=text, level=level, meta=meta)
|
||
else:
|
||
blk = Block(type=BlockType.PARAGRAPH, text=text, meta=meta)
|
||
|
||
blocks.append(blk)
|
||
|
||
return blocks, current_chapter
|