diff --git a/paperforge/worker/ocr_document.py b/paperforge/worker/ocr_document.py index ea53ae5d..1328a783 100644 --- a/paperforge/worker/ocr_document.py +++ b/paperforge/worker/ocr_document.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from paperforge.worker.ocr_families import ( discover_body_family_anchor, + partition_zone_families, discover_reference_family_anchor, ) from paperforge.worker.ocr_decisions import record_decision @@ -1627,6 +1628,14 @@ def analyze_document_structure(blocks: list[dict]) -> DocumentStructure: "reference_family_anchor": reference_family_anchor, }, ) + _apply_zone_labels(blocks, region_bus) + partition_zone_families( + blocks, + { + "body_family_anchor": body_family_anchor, + "reference_family_anchor": reference_family_anchor, + }, + ) tail_spread = _reconcile_tail_spread(blocks, page_layouts) if tail_spread is not None: backmatter_form = _classify_backmatter_form(tail_spread, blocks) @@ -1664,6 +1673,24 @@ def analyze_document_structure(blocks: list[dict]) -> DocumentStructure: return ds +def _apply_zone_labels(blocks: list[dict], region_bus: dict[str, dict] | None) -> None: + if not region_bus: + return + + block_ids_to_zone: dict[str, str] = {} + for zone_name, zone in region_bus.items(): + for block_id in zone.get("block_ids") or []: + block_ids_to_zone[str(block_id)] = zone_name + + for block in blocks: + block_id = block.get("block_id") + if block_id is None: + continue + zone_name = block_ids_to_zone.get(str(block_id)) + if zone_name: + block["zone"] = zone_name + + def _get_column(block: dict, page_width: float = 1200) -> int: bbox = block.get("bbox") or block.get("block_bbox") if bbox and len(bbox) >= 4: @@ -3085,6 +3112,14 @@ def normalize_document_structure(blocks: list[dict]) -> tuple[DocumentStructure, "reference_family_anchor": reference_family_anchor, }, ) + _apply_zone_labels(blocks, region_bus) + partition_zone_families( + blocks, + { + "body_family_anchor": body_family_anchor, + "reference_family_anchor": reference_family_anchor, + }, + ) tail_spread = _reconcile_tail_spread(blocks, page_layouts) if tail_spread is not None: diff --git a/paperforge/worker/ocr_families.py b/paperforge/worker/ocr_families.py index f9ca40c3..6de3c109 100644 --- a/paperforge/worker/ocr_families.py +++ b/paperforge/worker/ocr_families.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections import Counter, defaultdict +import re from typing import Any @@ -28,6 +29,8 @@ _REFERENCE_MARKER_TYPES = { } _REFERENCE_HEADING_TEXTS = {"references", "bibliography"} +_HEADING_MARKER_TYPES = {"canonical_section_name"} +_SUPPORT_ZONE_HINTS = ("support", "insert", "frontmatter", "margin", "sidebar") def discover_body_family_anchor(blocks: list[dict], page_count: int | None = None) -> dict[str, Any]: @@ -163,6 +166,29 @@ def discover_reference_family_anchor(blocks: list[dict], page_count: int | None } +def partition_zone_families( + blocks: list[dict], + anchors: dict[str, dict[str, Any]] | None = None, +) -> dict[str, dict[str, Any]]: + """Assign in-zone style/layout family artifacts without resolving final roles.""" + anchors = anchors or {} + body_anchor = anchors.get("body_family_anchor") or {} + reference_anchor = anchors.get("reference_family_anchor") or {} + + partitioned: dict[str, dict[str, Any]] = {} + for index, block in enumerate(blocks): + block_id = str(block.get("block_id") or f"block_{index}") + style_family, authority = _classify_style_family(block, body_anchor, reference_anchor) + block["style_family"] = style_family + block["style_family_authority"] = authority + partitioned[block_id] = { + "style_family": style_family, + "authority": authority, + "zone": block.get("zone"), + } + return partitioned + + def _select_middle_sample_pages(blocks: list[dict], page_count: int) -> set[int]: if page_count <= 0: return set() @@ -203,6 +229,100 @@ def _select_middle_sample_pages(blocks: list[dict], page_count: int) -> set[int] return stable_pages +def _classify_style_family( + block: dict, + body_anchor: dict[str, Any], + reference_anchor: dict[str, Any], +) -> tuple[str, str]: + marker_type = ((block.get("marker_signature") or {}).get("type") or "none") + text = str(block.get("text") or block.get("block_content") or "").strip() + zone = str(block.get("zone") or "") + text_lower = text.lower() + + if _reference_anchor_matches_block(reference_anchor, block): + return "reference_like", "reference_family_anchor" + if marker_type in _REFERENCE_MARKER_TYPES: + return "reference_like", "reference_marker" + if marker_type == "table_number" or text_lower.startswith("table "): + return "table_caption_like", "table_marker" + if marker_type in {"figure_number", "panel_label"} or text_lower.startswith("figure "): + return "legend_like", "figure_marker" + if marker_type in _HEADING_MARKER_TYPES: + return "heading_like", "heading_marker" + if _anchor_matches_block(body_anchor, block, _body_family_key): + return "body_like", "body_family_anchor" + if body_anchor.get("status") == "ACCEPT" and zone == "body_zone" and _is_body_zone_body_like(block): + return "body_like", "body_zone_with_anchor" + if zone == "body_zone" and _is_body_family_candidate(block): + return "body_like", "body_zone_candidate" + if any(hint in zone for hint in _SUPPORT_ZONE_HINTS): + return "support_like", "zone_context" + return "unknown_like", "fallback" + + +def _anchor_matches_block( + anchor: dict[str, Any], + block: dict, + key_builder: Any, +) -> bool: + if anchor.get("status") != "ACCEPT": + return False + anchor_key = ( + anchor.get("font_family_norm"), + anchor.get("font_size_bucket"), + anchor.get("width_bucket"), + anchor.get("x_center_bucket"), + ) + if any(value is not None for value in anchor_key): + return key_builder(block) == anchor_key + return False + + +def _reference_anchor_matches_block(anchor: dict[str, Any], block: dict) -> bool: + if not _anchor_matches_block(anchor, block, _reference_family_key): + return False + marker_type = ((block.get("marker_signature") or {}).get("type") or "none") + if marker_type in _REFERENCE_MARKER_TYPES: + return True + + zone = str(block.get("zone") or "") + if zone == "reference_zone": + return True + if zone == "body_zone": + return False + if _has_reference_text_structure(block): + return True + if zone: + return False + + sample_pages = [int(page) for page in anchor.get("sample_pages") or [] if page is not None] + page = int(block.get("page", 0) or 0) + if sample_pages and page > 0: + if page >= min(sample_pages) - 1: + return _has_reference_text_structure(block) + return False + + +def _has_reference_text_structure(block: dict) -> bool: + text = str(block.get("text") or block.get("block_content") or "").strip() + if not text: + return False + + if re.match(r"^\[\d+\]", text): + return True + if re.match(r"^\d+[\.)]", text): + return True + + citation_years = re.findall(r"\b(?:19|20)\d{2}[a-z]?\b", text) + if len(citation_years) >= 2: + return True + + if text.count(";") >= 2 and text.count(",") >= 3: + return True + + return False + + def _page_is_contaminated(page_blocks: list[dict]) -> bool: if not page_blocks: return True @@ -259,6 +379,19 @@ def _is_body_family_candidate(block: dict) -> bool: return True +def _is_body_zone_body_like(block: dict) -> bool: + marker_type = ((block.get("marker_signature") or {}).get("type") or "none") + if marker_type in _STRONG_MARKER_TYPES | {"panel_label"}: + return False + span_signature = block.get("span_signature") or {} + layout_signature = block.get("layout_signature") or {} + return bool( + span_signature.get("font_family_norm") + and (span_signature.get("font_size_median") is not None or span_signature.get("font_size_bucket") is not None) + and layout_signature.get("width") + ) + + def _body_family_key(block: dict) -> tuple[Any, ...] | None: span_signature = block.get("span_signature") or {} layout_signature = block.get("layout_signature") or {} diff --git a/paperforge/worker/ocr_profiles.py b/paperforge/worker/ocr_profiles.py index 3520a3b0..edbee024 100644 --- a/paperforge/worker/ocr_profiles.py +++ b/paperforge/worker/ocr_profiles.py @@ -245,6 +245,16 @@ FAMILY_DEFINITIONS: dict[str, list[str]] = { "caption_family": ["figure_caption", "table_caption"], } +STYLE_FAMILY_TO_PROFILE_FAMILY: dict[str, str] = { + "body_like": "body_family", + "heading_like": "heading_family", + "legend_like": "legend_family", + "table_caption_like": "table_caption_family", + "reference_like": "reference_family", + "support_like": "support_family", + "unknown_like": "unknown_family", +} + def build_family_profiles(blocks: list[dict]) -> dict: """Derive family-level profiles from block role/span data. @@ -259,7 +269,30 @@ def build_family_profiles(blocks: list[dict]) -> dict: """ role_profiles = build_role_span_profiles(blocks) families: dict = {} + + style_family_blocks: dict[str, list[dict]] = {} + for block in blocks: + style_family = block.get("style_family") + if not style_family: + continue + style_family_blocks.setdefault(str(style_family), []).append(block) + + for style_family, family_name in STYLE_FAMILY_TO_PROFILE_FAMILY.items(): + members = style_family_blocks.get(style_family, []) + if not members: + continue + profile = _build_profile_from_blocks(members) + if profile is None: + continue + families[family_name] = { + **profile, + "member_roles": sorted({str(block.get("role") or "") for block in members if block.get("role")}), + "source_style_family": style_family, + } + for family_name, member_roles in FAMILY_DEFINITIONS.items(): + if family_name in families: + continue available = { role: role_profiles[role] for role in member_roles @@ -306,6 +339,46 @@ def build_family_profiles(blocks: list[dict]) -> dict: return families +def _build_profile_from_blocks(blocks: list[dict]) -> dict | None: + profiles = [extract_block_span_profile(block) for block in blocks] + usable_profiles = [profile for profile in profiles if profile is not None] + if not usable_profiles: + return None + + sizes: list[float] = [] + font_families: set[str] = set() + bold_count = 0 + italic_count = 0 + colored_count = 0 + for profile in usable_profiles: + sizes.extend([profile["mean_size"], profile["max_size"]]) + font_families.update(profile["font_families"]) + if profile["is_bold"]: + bold_count += 1 + if profile["is_italic"]: + italic_count += 1 + if profile["is_colored"]: + colored_count += 1 + + block_count = len(usable_profiles) + mean_size = sum(sizes) / len(sizes) if sizes else 0.0 + max_size = max(sizes) if sizes else 0.0 + min_size = min(sizes) if sizes else 0.0 + dispersion = (max_size - min_size) / max_size if max_size > 0 else 0.0 + return { + "block_count": block_count, + "mean_size": round(mean_size, 2), + "max_size": round(max_size, 2), + "min_size": round(min_size, 2), + "dispersion": round(dispersion, 4), + "quality": _profile_quality(block_count, dispersion), + "bold_ratio": round(bold_count / block_count, 2) if block_count else 0.0, + "italic_ratio": round(italic_count / block_count, 2) if block_count else 0.0, + "font_families": list(font_families), + "colored_ratio": round(colored_count / block_count, 2) if block_count else 0.0, + } + + def compare_against_family( block_profile: dict, family_profile: dict, diff --git a/tests/test_ocr_document.py b/tests/test_ocr_document.py index e809ca13..155f1276 100644 --- a/tests/test_ocr_document.py +++ b/tests/test_ocr_document.py @@ -163,6 +163,55 @@ def test_analyze_document_structure_discovers_reference_family_anchor_before_tai assert blocks[4]["role"] == "body_paragraph" +def test_analyze_document_structure_wires_style_family_artifacts_into_blocks() -> None: + from paperforge.worker.ocr_document import analyze_document_structure + + blocks = [ + { + "block_id": "p3_b1", + "page": 3, + "role": "body_paragraph", + "text": "Long narrative body paragraph with enough text to be treated as core body evidence across the main article flow and stable typography. " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_size_bucket": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 260, "width_bucket": 250, "x_center": 240, "x_center_bucket": 250}, + }, + { + "block_id": "p5_b1", + "page": 4, + "role": "body_paragraph", + "text": "Long narrative body paragraph continuing the main article flow with matching typography and substantial prose for body-anchor detection. " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_size_bucket": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 262, "width_bucket": 250, "x_center": 242, "x_center_bucket": 250}, + }, + { + "block_id": "p5_b1b", + "page": 5, + "role": "body_paragraph", + "text": "Long narrative body paragraph extending the article with the same body typography and enough tokens to remain a strong body candidate. " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_size_bucket": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 261, "width_bucket": 250, "x_center": 241, "x_center_bucket": 250}, + }, + { + "block_id": "p5_b2", + "page": 5, + "role": "body_paragraph", + "text": "Figure 2. Compact legend text", + "marker_signature": {"type": "figure_number", "number": 2}, + "span_signature": {"font_size_median": 8.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 220, "x_center": 250}, + }, + ] + + analyze_document_structure(blocks) + + assert blocks[0]["style_family"] == "body_like" + assert blocks[0]["style_family_authority"] in {"body_family_anchor", "body_zone_with_anchor", "body_zone_candidate"} + assert blocks[3]["style_family"] == "legend_like" + + def test_normalize_document_structure_keeps_reference_family_anchor_anchor_first() -> None: from paperforge.worker.ocr_document import normalize_document_structure @@ -203,6 +252,59 @@ def test_normalize_document_structure_keeps_reference_family_anchor_anchor_first assert normalized_blocks[4]["role"] == "body_paragraph" +def test_normalize_document_structure_wires_style_family_artifacts_into_blocks() -> None: + from paperforge.worker.ocr_document import normalize_document_structure + + blocks = [ + { + "block_id": "p3_b1", + "page": 3, + "role": "body_paragraph", + "text": "Long narrative body paragraph with enough text to be treated as core body evidence across the main article flow and stable typography. " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_size_bucket": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 260, "width_bucket": 250, "x_center": 240, "x_center_bucket": 250}, + }, + { + "block_id": "p4_b1", + "page": 4, + "role": "body_paragraph", + "text": "Long narrative body paragraph continuing the main article flow with matching typography and substantial prose for body-anchor detection. " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_size_bucket": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 262, "width_bucket": 250, "x_center": 242, "x_center_bucket": 250}, + }, + { + "block_id": "p5_b1", + "page": 5, + "role": "body_paragraph", + "text": "Long narrative body paragraph extending the article with the same body typography and enough tokens to remain a strong body candidate. " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_size_bucket": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 261, "width_bucket": 250, "x_center": 241, "x_center_bucket": 250}, + }, + { + "block_id": "p5_b2", + "page": 5, + "role": "body_paragraph", + "text": "Figure 2. Compact legend text", + "marker_signature": {"type": "figure_number", "number": 2}, + "span_signature": {"font_size_median": 8.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 220, "x_center": 250}, + }, + ] + + _, normalized_blocks = normalize_document_structure(blocks) + + assert normalized_blocks[0]["style_family"] == "body_like" + assert normalized_blocks[0]["style_family_authority"] in { + "body_family_anchor", + "body_zone_with_anchor", + "body_zone_candidate", + } + assert normalized_blocks[3]["style_family"] == "legend_like" + + def test_reference_zone_is_inferred_from_reference_family_anchor_not_preexisting_roles() -> None: from paperforge.worker.ocr_document import infer_zones diff --git a/tests/test_ocr_families.py b/tests/test_ocr_families.py index 84ce7501..e10a00c4 100644 --- a/tests/test_ocr_families.py +++ b/tests/test_ocr_families.py @@ -589,3 +589,180 @@ def test_reference_family_anchor_accepts_single_page_penultimate_tail_family() - assert anchor["status"] == "ACCEPT" assert anchor["family_name"] == "reference_family" assert anchor["sample_pages"] == [9] + + +def test_partition_zone_families_separates_legend_like_from_body_like() -> None: + from paperforge.worker.ocr_families import partition_zone_families + + blocks = [ + { + "block_id": "p4_b1", + "zone": "body_zone", + "text": "Long narrative paragraph " * 6, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 260}, + }, + { + "block_id": "p4_b2", + "zone": "body_zone", + "text": "Figure 2. Long legend text", + "marker_signature": {"type": "figure_number", "number": 2}, + "span_signature": {"font_size_median": 8.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 220}, + }, + ] + anchors = {"body_family_anchor": {"status": "ACCEPT", "family_name": "body_family"}} + + partitioned = partition_zone_families(blocks, anchors) + assert partitioned["p4_b1"]["style_family"] == "body_like" + assert partitioned["p4_b2"]["style_family"] == "legend_like" + + +def test_partition_zone_families_marks_reference_like_from_reference_anchor() -> None: + from paperforge.worker.ocr_families import partition_zone_families + + blocks = [ + { + "block_id": "p9_b1", + "zone": "backmatter_zone", + "text": "[1] Example reference entry with authors and journal", + "marker_signature": {"type": "reference_numeric_bracket", "number": 1}, + "span_signature": {"font_size_median": 8.5, "font_size_bucket": 8.5, "font_family_norm": "Times"}, + "layout_signature": {"width": 250, "width_bucket": 250, "x_center": 240, "x_center_bucket": 250}, + } + ] + anchors = { + "reference_family_anchor": { + "status": "ACCEPT", + "family_name": "reference_family", + "font_family_norm": "Times", + "font_size_bucket": 8.5, + "width_bucket": 250, + "x_center_bucket": 250, + } + } + + partitioned = partition_zone_families(blocks, anchors) + assert partitioned["p9_b1"]["style_family"] == "reference_like" + + +def test_partition_zone_families_surfaces_style_family_on_blocks_for_downstream_use() -> None: + from paperforge.worker.ocr_families import partition_zone_families + + blocks = [ + { + "block_id": "p4_b1", + "zone": "body_zone", + "text": "Long narrative paragraph " * 6, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 9.0, "font_family_norm": "Times"}, + "layout_signature": {"width": 260}, + } + ] + + partitioned = partition_zone_families( + blocks, + {"body_family_anchor": {"status": "ACCEPT", "family_name": "body_family"}}, + ) + + assert partitioned["p4_b1"]["style_family"] == "body_like" + assert blocks[0]["style_family"] == "body_like" + assert blocks[0]["style_family_authority"] == "body_zone_with_anchor" + + +def test_partition_zone_families_does_not_mark_body_zone_style_match_as_reference_like() -> None: + from paperforge.worker.ocr_families import partition_zone_families + + blocks = [ + { + "block_id": "p4_b1", + "zone": "body_zone", + "text": "Long narrative paragraph without citation list structure " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 8.5, "font_size_bucket": 8.5, "font_family_norm": "Times"}, + "layout_signature": {"width": 250, "width_bucket": 250, "x_center": 240, "x_center_bucket": 250}, + } + ] + anchors = { + "body_family_anchor": {"status": "ACCEPT", "family_name": "body_family"}, + "reference_family_anchor": { + "status": "ACCEPT", + "family_name": "reference_family", + "font_family_norm": "Times", + "font_size_bucket": 8.5, + "width_bucket": 250, + "x_center_bucket": 250, + "sample_pages": [9, 10], + }, + } + + partitioned = partition_zone_families(blocks, anchors) + + assert partitioned["p4_b1"]["style_family"] == "body_like" + assert partitioned["p4_b1"]["authority"] != "reference_family_anchor" + + +def test_partition_zone_families_does_not_mark_late_body_zone_style_match_as_reference_like() -> None: + from paperforge.worker.ocr_families import partition_zone_families + + blocks = [ + { + "block_id": "p9_b1", + "page": 9, + "zone": "body_zone", + "text": "Long late-page narrative paragraph without numbered reference structure " * 3, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 8.5, "font_size_bucket": 8.5, "font_family_norm": "Times"}, + "layout_signature": {"width": 250, "width_bucket": 250, "x_center": 240, "x_center_bucket": 250}, + } + ] + anchors = { + "body_family_anchor": {"status": "ACCEPT", "family_name": "body_family"}, + "reference_family_anchor": { + "status": "ACCEPT", + "family_name": "reference_family", + "font_family_norm": "Times", + "font_size_bucket": 8.5, + "width_bucket": 250, + "x_center_bucket": 250, + "sample_pages": [9, 10], + }, + } + + partitioned = partition_zone_families(blocks, anchors) + + assert partitioned["p9_b1"]["style_family"] == "body_like" + assert partitioned["p9_b1"]["authority"] != "reference_family_anchor" + + +def test_partition_zone_families_does_not_mark_backmatter_prose_style_match_as_reference_like() -> None: + from paperforge.worker.ocr_families import partition_zone_families + + blocks = [ + { + "block_id": "p10_b1", + "page": 10, + "zone": "backmatter_zone", + "text": "Appendix discussion prose continues the argument in full sentences without numbered entries, author-year formatting, or list-like citation structure. " * 2, + "marker_signature": {"type": "none"}, + "span_signature": {"font_size_median": 8.5, "font_size_bucket": 8.5, "font_family_norm": "Times"}, + "layout_signature": {"width": 250, "width_bucket": 250, "x_center": 240, "x_center_bucket": 250}, + } + ] + anchors = { + "reference_family_anchor": { + "status": "ACCEPT", + "family_name": "reference_family", + "font_family_norm": "Times", + "font_size_bucket": 8.5, + "width_bucket": 250, + "x_center_bucket": 250, + "sample_pages": [9, 10], + } + } + + partitioned = partition_zone_families(blocks, anchors) + + assert partitioned["p10_b1"]["style_family"] != "reference_like" + assert partitioned["p10_b1"]["authority"] != "reference_family_anchor" diff --git a/tests/test_ocr_profiles.py b/tests/test_ocr_profiles.py index f333d752..ced3492b 100644 --- a/tests/test_ocr_profiles.py +++ b/tests/test_ocr_profiles.py @@ -214,3 +214,27 @@ def test_cross_validate_with_span_suggests_alternative() -> None: result = cross_validate_with_span(block, "body_paragraph", profiles) assert "section_heading" in result["suggested_roles"] assert result["adjustment"] < 0 + + +def test_build_family_profiles_prefers_style_family_partition_artifacts() -> None: + from paperforge.worker.ocr_profiles import build_family_profiles + + blocks = [ + { + "role": "body_paragraph", + "style_family": "body_like", + "span_metadata": {"size": 10.0, "font": "Times", "flags": "normal"}, + }, + { + "role": "body_paragraph", + "style_family": "legend_like", + "span_metadata": {"size": 8.0, "font": "Times", "flags": "normal"}, + }, + ] + + families = build_family_profiles(blocks) + + assert families["body_family"]["block_count"] == 1 + assert families["body_family"]["mean_size"] == 10.0 + assert families["legend_family"]["block_count"] == 1 + assert families["legend_family"]["mean_size"] == 8.0