From 6bbe00204ce2ea430d1ca85161b87401ec6dabeb Mon Sep 17 00:00:00 2001 From: Research Assistant Date: Sat, 6 Jun 2026 12:34:53 +0800 Subject: [PATCH] feat: add section-aware OCR rescue pass --- paperforge/worker/ocr_blocks.py | 23 ++-- paperforge/worker/ocr_document.py | 86 +++++++++++++++ tests/test_ocr_document.py | 167 +++++++++++++++++++++++++++++- tests/test_ocr_roles.py | 2 +- 4 files changed, 260 insertions(+), 18 deletions(-) diff --git a/paperforge/worker/ocr_blocks.py b/paperforge/worker/ocr_blocks.py index 58f960af..dff3329a 100644 --- a/paperforge/worker/ocr_blocks.py +++ b/paperforge/worker/ocr_blocks.py @@ -70,24 +70,15 @@ def build_structured_blocks(raw_blocks: list[dict]) -> list[dict]: paper_context["role_profiles"] = build_role_span_profiles(rows) - # Second pass: cross-validate low-confidence blocks + # Second pass: section-aware role rescue if paper_context.get("role_profiles"): - for row in rows: - if row.get("role_confidence", 1.0) < 0.7: - from paperforge.worker.ocr_roles import second_pass_cross_validate + from paperforge.worker.ocr_document import ( + analyze_document_structure, + rescue_roles_with_document_context, + ) - xv_result = second_pass_cross_validate(row, paper_context["role_profiles"]) - row["role_confidence"] = round( - max(0.0, min(1.0, row.get("role_confidence", 0.5) + xv_result["confidence_adjustment"])), - 4, - ) - if xv_result["role_changed"]: - row["role"] = xv_result["role"] - row.setdefault("evidence", []).append(f"second_pass: {xv_result['role']} (span cross-validation)") - if xv_result["suggested_roles"]: - row.setdefault("evidence", []).append( - f"span_alternatives: {','.join(xv_result['suggested_roles'])}" - ) + document_structure = analyze_document_structure(rows) + rows = rescue_roles_with_document_context(rows, paper_context["role_profiles"], document_structure) return rows diff --git a/paperforge/worker/ocr_document.py b/paperforge/worker/ocr_document.py index d7987f56..d6d22b80 100644 --- a/paperforge/worker/ocr_document.py +++ b/paperforge/worker/ocr_document.py @@ -446,6 +446,92 @@ def _page_still_frontmatter(page_blocks: list[dict], page_num: int, page_height: return True +def rescue_roles_with_document_context( + blocks: list[dict], + role_profiles: dict, + document_structure: DocumentStructure | None = None, +) -> list[dict]: + """Apply section-context-aware role rescue rules using document structure. + + Uses the previously-built role style profiles and document boundaries to + correct common role-assignment errors: + + 1. ``frontmatter_noise`` in the body section with body-like font + → ``body_paragraph`` + 2. ``body_paragraph`` in the references section with reference-like font + → ``reference_item`` (only when confidence < 0.7) + 3. Weak heading (confidence < 0.6) with body-like font → ``body_paragraph`` + + Never overrides: strong formal prefixes (Figure, Table), strong numbering, + or explicit boundary-heading logic. + + Returns a new list of blocks with corrected roles. + """ + from paperforge.worker.ocr_profiles import ( + compare_against_role_family, + extract_block_span_profile, + ) + from paperforge.worker.ocr_roles import _has_heading_numbering + + if document_structure is None: + document_structure = analyze_document_structure(blocks) + + body_end_page = document_structure.body_end_page + refs_start = document_structure.references_start + refs_start_page = refs_start.page if refs_start else None + + result = list(blocks) + + for block in result: + # --- Rule 1: frontmatter_noise → body_paragraph (body section + body font) + if block.get("role") == "frontmatter_noise": + page = block.get("page", 1) or 1 + if body_end_page is not None and page <= body_end_page: + bp = extract_block_span_profile(block) + if bp: + body_fam = role_profiles.get("body_paragraph", {}) + if body_fam: + match = compare_against_role_family(bp, body_fam) + if match["size_compatible"] and match["match_score"] > 0.5: + block["role"] = "body_paragraph" + block["role_confidence"] = min(block.get("role_confidence", 0.5) + 0.1, 1.0) + block.setdefault("evidence", []).append("rescue: frontmatter_noise → body_paragraph") + + # --- Rule 2: body_paragraph → reference_item (refs section + ref font) + role = block.get("role", "") + if role == "body_paragraph" and block.get("role_confidence", 1.0) < 0.7: + page = block.get("page", 1) or 1 + if refs_start_page is not None and page >= refs_start_page: + bp = extract_block_span_profile(block) + if bp: + ref_fam = role_profiles.get("reference_item", {}) + if ref_fam: + match = compare_against_role_family(bp, ref_fam) + if match["size_compatible"] and match["match_score"] > 0.5: + block["role"] = "reference_item" + block["role_confidence"] = min(block.get("role_confidence", 0.5) + 0.2, 1.0) + block.setdefault("evidence", []).append("rescue: body_paragraph → reference_item") + + # --- Rule 3: weak heading with body font → body_paragraph + if ( + role in {"section_heading", "subsection_heading", "sub_subsection_heading"} + and block.get("role_confidence", 1.0) < 0.6 + ): + text = str(block.get("text", "") or block.get("block_content", "") or "") + if _has_heading_numbering(text): + continue + bp = extract_block_span_profile(block) + if bp: + body_fam = role_profiles.get("body_paragraph", {}) + if body_fam: + match = compare_against_role_family(bp, body_fam) + if match["size_compatible"] and match["match_score"] > 0.5: + block["role"] = "body_paragraph" + block.setdefault("evidence", []).append("rescue: heading → body_paragraph") + + return result + + def analyze_document_structure(blocks: list[dict]) -> DocumentStructure: """Produce a structured document boundary object. diff --git a/tests/test_ocr_document.py b/tests/test_ocr_document.py index cdeadd49..f47461dd 100644 --- a/tests/test_ocr_document.py +++ b/tests/test_ocr_document.py @@ -68,7 +68,172 @@ def test_analyze_document_structure_no_tail() -> None: assert doc.backmatter_form == "flat" -def test_analyze_document_structure_mixed_tail_spread() -> None: +def test_rescue_frontmatter_noise_to_body_paragraph() -> None: + """frontmatter_noise in body section with body-like font → body_paragraph.""" + from paperforge.worker.ocr_document import ( + DocumentStructure, + rescue_roles_with_document_context, + ) + + blocks = [ + {"page": 1, "role": "body_paragraph", "role_confidence": 0.7, "text": "Intro body paragraph with enough text to establish a profile.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 2, "role": "body_paragraph", "role_confidence": 0.7, "text": "Methods section describing the experimental setup in full detail.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 3, "role": "frontmatter_noise", "role_confidence": 0.6, "text": "This block was misclassified as frontmatter noise but is actually body text.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 4, "role": "reference_heading", "role_confidence": 0.9, "text": "References"}, + {"page": 4, "role": "reference_item", "role_confidence": 0.7, "text": "[1] Author A, Journal, 2025"}, + ] + + role_profiles = { + "body_paragraph": {"block_count": 2, "mean_size": 10.0, "max_size": 10.5, "min_size": 9.5, + "dispersion": 0.05, "quality": "strong", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + "frontmatter_noise": {"block_count": 1, "mean_size": 8.0, "max_size": 8.5, "min_size": 7.5, + "dispersion": 0.05, "quality": "weak", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + } + + ds = DocumentStructure(body_end_page=3) + result = rescue_roles_with_document_context(blocks, role_profiles, ds) + + rescued = next(b for b in result if b["text"].startswith("This block was misclassified")) + assert rescued["role"] == "body_paragraph", ( + f"Expected body_paragraph, got {rescued['role']}" + ) + + +def test_rescue_body_paragraph_to_reference_item() -> None: + """body_paragraph with ref-like font in references section → reference_item.""" + from paperforge.worker.ocr_document import ( + DocumentStructure, + PagePosition, + rescue_roles_with_document_context, + ) + + blocks = [ + {"page": 1, "role": "body_paragraph", "role_confidence": 0.7, "text": "Main body paragraph that is long enough for context.", + "span_metadata": {"size": 10.5, "flags": "normal"}}, + {"page": 5, "role": "reference_heading", "role_confidence": 0.9, "text": "References"}, + {"page": 5, "role": "body_paragraph", "role_confidence": 0.5, "text": "1. Smith J, Johnson K. A study on something. Journal, 2025.", + "span_metadata": {"size": 9.0, "flags": "normal"}}, + ] + + role_profiles = { + "body_paragraph": {"block_count": 1, "mean_size": 10.5, "max_size": 10.5, "min_size": 10.5, + "dispersion": 0.0, "quality": "strong", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + "reference_item": {"block_count": 0, "mean_size": 9.0, "max_size": 9.0, "min_size": 9.0, + "dispersion": 0.0, "quality": "weak", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + } + + ds = DocumentStructure(body_end_page=4, references_start=PagePosition(page=5, y=0)) + result = rescue_roles_with_document_context(blocks, role_profiles, ds) + + promoted = next(b for b in result if b["text"].startswith("1. Smith")) + assert promoted["role"] == "reference_item", ( + f"Expected reference_item, got {promoted['role']}" + ) + + +def test_rescue_weak_heading_demoted_to_body() -> None: + """Weak heading (confidence < 0.6) with body-like font → body_paragraph.""" + from paperforge.worker.ocr_document import ( + DocumentStructure, + rescue_roles_with_document_context, + ) + + blocks = [ + {"page": 1, "role": "body_paragraph", "role_confidence": 0.7, "text": "Some intro body text that provides a long enough context for profiling.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 2, "role": "body_paragraph", "role_confidence": 0.7, "text": "More body text to strengthen the body profile for better font matching.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 2, "role": "section_heading", "role_confidence": 0.5, "text": "Methods", + "span_metadata": {"size": 10, "flags": "normal"}}, + ] + + role_profiles = { + "body_paragraph": {"block_count": 2, "mean_size": 10.0, "max_size": 10.5, "min_size": 9.5, + "dispersion": 0.05, "quality": "strong", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + "section_heading": {"block_count": 0, "mean_size": 12.0, "max_size": 12.0, "min_size": 12.0, + "dispersion": 0.0, "quality": "no_data", "bold_ratio": 1.0, + "italic_ratio": 0.0, "font_families": []}, + } + + ds = DocumentStructure(body_end_page=2) + result = rescue_roles_with_document_context(blocks, role_profiles, ds) + + demoted = next(b for b in result if b["text"] == "Methods") + assert demoted["role"] == "body_paragraph", ( + f"Expected body_paragraph, got {demoted['role']}" + ) + + +def test_rescue_strong_numbered_heading_not_demoted() -> None: + """Strong numbered heading (e.g. '5.1 Results') should NOT be demoted.""" + from paperforge.worker.ocr_document import ( + DocumentStructure, + rescue_roles_with_document_context, + ) + + blocks = [ + {"page": 1, "role": "body_paragraph", "role_confidence": 0.7, "text": "Body paragraph one with enough text to contribute to a font profile.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 2, "role": "section_heading", "role_confidence": 0.5, "text": "5.1 Results", + "span_metadata": {"size": 10, "flags": "normal"}}, + ] + + role_profiles = { + "body_paragraph": {"block_count": 1, "mean_size": 10.0, "max_size": 10.5, "min_size": 9.5, + "dispersion": 0.05, "quality": "strong", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + } + + ds = DocumentStructure(body_end_page=2) + result = rescue_roles_with_document_context(blocks, role_profiles, ds) + + heading = next(b for b in result if b["text"] == "5.1 Results") + assert heading["role"] == "section_heading", ( + f"Expected section_heading preserved, got {heading['role']}" + ) + + +def test_rescue_no_document_structure_derived() -> None: + """rescue_roles_with_document_context should derive structure when not provided.""" + from paperforge.worker.ocr_document import rescue_roles_with_document_context + + blocks = [ + {"page": 1, "role": "body_paragraph", "role_confidence": 0.7, "text": "Body text that is long enough for establishing a font profile here.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 2, "role": "body_paragraph", "role_confidence": 0.7, "text": "More body text here to help build a solid body paragraph profile.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 2, "role": "frontmatter_noise", "role_confidence": 0.6, "text": "Misclassified text block that should be rescued to body paragraph.", + "span_metadata": {"size": 10, "flags": "normal"}}, + {"page": 3, "role": "reference_heading", "role_confidence": 0.9, "text": "References"}, + {"page": 3, "role": "reference_item", "role_confidence": 0.7, "text": "[1] Author A, Journal, 2025"}, + ] + + role_profiles = { + "body_paragraph": {"block_count": 2, "mean_size": 10.0, "max_size": 10.5, "min_size": 9.5, + "dispersion": 0.05, "quality": "strong", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + "frontmatter_noise": {"block_count": 1, "mean_size": 8.0, "max_size": 8.5, "min_size": 7.5, + "dispersion": 0.05, "quality": "weak", "bold_ratio": 0.0, + "italic_ratio": 0.0, "font_families": []}, + } + + result = rescue_roles_with_document_context(blocks, role_profiles) + + rescued = next(b for b in result if b["text"].startswith("Misclassified")) + assert rescued["role"] == "body_paragraph", ( + f"Expected body_paragraph, got {rescued['role']}" + ) + + +def test_rescue_analyze_document_structure_mixed_tail_spread() -> None: from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure blocks = [ diff --git a/tests/test_ocr_roles.py b/tests/test_ocr_roles.py index 22cbf9b4..f79d50d5 100644 --- a/tests/test_ocr_roles.py +++ b/tests/test_ocr_roles.py @@ -353,4 +353,4 @@ def test_backmatter_boundary_detects_on_early_page() -> None: "page": 5, } result = _is_backmatter_boundary_heading(block, 5, 10) - assert result == True # After fix: 5/10 = 50% in second half, has container words + bold + assert result # After fix: 5/10 = 50% in second half, has container words + bold