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
|