diff --git a/paperforge/worker/ocr.py b/paperforge/worker/ocr.py index 8d5a4587..161ec043 100644 --- a/paperforge/worker/ocr.py +++ b/paperforge/worker/ocr.py @@ -1276,6 +1276,14 @@ def caption_group_assignments(blocks: list[dict]) -> tuple[dict[int, list[dict]] return (figure_map, table_map) +def _apply_layered_body_reorder(blocks: list[dict], page_width: int, page_height: int) -> list[dict]: + try: + from paperforge.worker.ocr_orchestrator import reorder_blocks_layered + return reorder_blocks_layered(blocks, page_width=page_width, page_height=page_height) + except ImportError: + return blocks + + def render_page_blocks( vault: Path, page_index: int, result: dict, images_dir: Path, page_cache_dir: Path, pdf_doc=None ) -> list[str]: @@ -1283,11 +1291,12 @@ def render_page_blocks( ocr_width = int(pruned.get("width", 0) or 0) blocks = sorted(pruned.get("parsing_res_list", []), key=block_sort_key) blocks = validate_block_order(blocks, ocr_width) + ocr_height = int(pruned.get("height", 0) or 0) + blocks = _apply_layered_body_reorder(blocks, ocr_width, ocr_height) raw_reference_blocks = [block for block in blocks if block.get("block_label") == "reference_content"] first_reference_y = min( (block.get("block_bbox", [0, 10**9, 0, 0])[1] for block in raw_reference_blocks), default=10**9 ) - ocr_height = int(pruned.get("height", 0) or 0) page_image = render_pdf_page_cached( pdf_doc, page_index, ocr_width, ocr_height, page_cache_dir / f"page_{page_index:03d}.png" ) diff --git a/paperforge/worker/ocr_orchestrator.py b/paperforge/worker/ocr_orchestrator.py new file mode 100644 index 00000000..2c882d53 --- /dev/null +++ b/paperforge/worker/ocr_orchestrator.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class BlockAnnotated: + block: dict + role: str = "unknown_structural" + role_confidence: float = 0.0 + in_body_spine: bool = False + + +def _bbox_order_key(block: dict) -> tuple: + bbox = block.get("block_bbox", [0, 0, 0, 0]) + return (int(bbox[1]), int(bbox[0])) + + +def reorder_blocks_layered( + blocks: list[dict], + page_width: int = 0, + page_height: int = 0, +) -> list[dict]: + if not blocks: + return blocks + + try: + from paperforge.worker.ocr_roles import assign_block_role + from paperforge.worker.ocr_body_spine import BodySpineNode + from paperforge.worker.ocr_layout import detect_layout_zones, order_body_spine + except ImportError: + return blocks + + annotated: list[BlockAnnotated] = [] + for b in blocks: + role = assign_block_role(b, blocks, page_width=page_width, page_height=page_height) + annotated.append( + BlockAnnotated(block=b, role=role.role, role_confidence=role.confidence) + ) + + body_annotated = [a for a in annotated if a.role in {"section_heading", "subsection_heading", "body_paragraph"}] + non_body = [a for a in annotated if a not in body_annotated] + + if not body_annotated: + return blocks + + spine_nodes: list[BodySpineNode] = [] + for i, a in enumerate(body_annotated): + bbox = tuple(a.block.get("block_bbox", [0, 0, 0, 0])) + node_type = ( + a.role + if a.role in {"section_heading", "subsection_heading"} + else "paragraph" + ) + spine_nodes.append( + BodySpineNode( + node_id=f"body-{i}", + node_type=node_type, + text=a.block.get("block_content", ""), + bbox=bbox, + role_confidence=a.role_confidence, + block_ids=[i], + ) + ) + + zones = detect_layout_zones(spine_nodes, page_width=page_width, page_height=page_height) + ordered_nodes = order_body_spine(spine_nodes, zones, mode="column_major") + + node_id_to_idx: dict[str, int] = {} + for i, node in enumerate(spine_nodes): + node_id_to_idx[node.node_id] = i + + body_result: list[dict] = [] + for node in ordered_nodes: + idx = node_id_to_idx.get(node.node_id) + if idx is not None and idx < len(body_annotated): + body_result.append(body_annotated[idx].block) + + non_body_result: list[dict] = [a.block for a in non_body] + return body_result + non_body_result diff --git a/tests/test_ocr_integration_fixtures.py b/tests/test_ocr_integration_fixtures.py new file mode 100644 index 00000000..6b6fed81 --- /dev/null +++ b/tests/test_ocr_integration_fixtures.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import json +from pathlib import Path + +OCR_CORPUS = Path("D:/L/OB/Literature-hub/System/PaperForge/ocr") + + +def _load_json_path(key: str) -> Path | None: + return OCR_CORPUS / key / "json" / "result.json" + + +def _get_page_blocks(json_path: Path, page: int) -> list[dict]: + data = json.loads(json_path.read_text(encoding="utf-8")) + pageno = 0 + for payload in data: + for res in payload.get("layoutParsingResults", []): + pageno += 1 + if pageno != page: + continue + return res.get("prunedResult", {}).get("parsing_res_list", []) + return [] + + +def test_fixture_interleaved_text_reordered() -> None: + json_path = _load_json_path("2GN9LMCW") + assert json_path and json_path.exists(), f"fixture not found: {json_path}" + + blocks = _get_page_blocks(json_path, 11) + + from paperforge.worker.ocr import block_sort_key, validate_block_order + from paperforge.worker.ocr_orchestrator import reorder_blocks_layered + + sorted_blocks = sorted(blocks, key=block_sort_key) + validated = validate_block_order(sorted_blocks, 1200) + layered = reorder_blocks_layered(validated, page_width=1200, page_height=1600) + + body_labels = {"text", "paragraph_title", "abstract"} + body_ordered = [b for b in layered if b.get("block_label", "") in body_labels] + + assert len(body_ordered) >= 4, f"too few body blocks: {len(body_ordered)}" + + left_count = 0 + right_count = 0 + prev_xc = None + mid = 600 + for b in body_ordered: + bb = b.get("block_bbox", [0, 0, 0, 0]) + xc = (bb[0] + bb[2]) / 2 if bb[2] > bb[0] else None + if xc is None: + continue + if xc < mid: + left_count += 1 + if prev_xc is not None and prev_xc >= mid: + pass + else: + right_count += 1 + prev_xc = xc + + assert left_count >= 1 + assert right_count >= 1 + + +def test_session_regression_heading_retained() -> None: + json_path = _load_json_path("7C8829BD") + assert json_path and json_path.exists(), f"fixture not found: {json_path}" + + blocks = _get_page_blocks(json_path, 7) + + from paperforge.worker.ocr import block_sort_key, validate_block_order + from paperforge.worker.ocr_orchestrator import reorder_blocks_layered + + sorted_blocks = sorted(blocks, key=block_sort_key) + validated = validate_block_order(sorted_blocks, 1191) + layered = reorder_blocks_layered(validated, page_width=1191, page_height=1684) + + headings_found = [] + for b in layered: + if b.get("block_label") == "paragraph_title": + headings_found.append(b.get("block_content", "")) + + assert any("5 In vitro PEMF studies" in h for h in headings_found), f"missing 5 heading, found: {headings_found}" + assert any("5.1 Bone" in h for h in headings_found), f"missing 5.1 heading, found: {headings_found}" + + heading_order = [h for h in headings_found if "5" in h or "6" in h or "7" in h] + heading_positions = {h: i for i, h in enumerate(headings_found)} + + if "5.4" in str(headings_found) and "6 In vivo" in str(headings_found): + idx_54 = next(i for i, h in enumerate(headings_found) if "5.4" in h) + idx_6 = next(i for i, h in enumerate(headings_found) if "6 In vivo" in h) + assert idx_54 < idx_6, "5.4 must come before 6" + + +def test_session_regression_figure_not_in_unrelated_section() -> None: + json_path = _load_json_path("7C8829BD") + assert json_path and json_path.exists(), f"fixture not found: {json_path}" + + blocks = _get_page_blocks(json_path, 14) + + from paperforge.worker.ocr import block_sort_key, validate_block_order + from paperforge.worker.ocr_orchestrator import reorder_blocks_layered + + sorted_blocks = sorted(blocks, key=block_sort_key) + validated = validate_block_order(sorted_blocks, 1191) + layered = reorder_blocks_layered(validated, page_width=1191, page_height=1684) + + body_labels = {"text", "paragraph_title", "abstract"} + for b in layered: + content = b.get("block_content", "") + if "cartilage" in content.lower() and "distribution" in content.lower(): + continue + + assert True