From 29fc619d4de218f29331e2cac9052d8846020fcc Mon Sep 17 00:00:00 2001 From: LLLin000 <809867916@qq.com> Date: Fri, 3 Jul 2026 12:35:33 +0800 Subject: [PATCH] feat(ocr): add vnext same-page figure pass --- paperforge/worker/ocr_figure_vnext_passes.py | 104 +++++++++++++++++++ paperforge/worker/ocr_figures.py | 31 ++++-- tests/test_ocr_figure_vnext_compare.py | 17 +++ tests/test_ocr_figure_vnext_passes.py | 47 +++++++++ 4 files changed, 191 insertions(+), 8 deletions(-) create mode 100644 paperforge/worker/ocr_figure_vnext_passes.py create mode 100644 tests/test_ocr_figure_vnext_compare.py create mode 100644 tests/test_ocr_figure_vnext_passes.py diff --git a/paperforge/worker/ocr_figure_vnext_passes.py b/paperforge/worker/ocr_figure_vnext_passes.py new file mode 100644 index 00000000..1fca8dd6 --- /dev/null +++ b/paperforge/worker/ocr_figure_vnext_passes.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from .ocr_figure_vnext_types import ClaimProposal, PassReport, ResourceRef + + +def _resource_page(block: dict) -> int | None: + page = block.get("page") + if page is None: + page = block.get("page_num") + return int(page) if page is not None else None + + +class PrimarySamePagePass: + name = "primary_same_page" + + def _collect_proposals(self, state): + from . import ocr_figures + + proposals = [] + for legend in state.candidate_index.deduped_legends: + page = _resource_page(legend) + if page is None: + continue + page_groups = [g for g in state.candidate_index.candidate_groups if _resource_page(g) == page] + for group in page_groups: + score = ocr_figures._score_legend_to_group( + legend, + group, + caption_score=ocr_figures.score_figure_caption( + legend, + nearby_media=True, + caption_style_match=False, + body_prose_likelihood=False, + ), + page_width=state.corpus.page_width, + ) + if score.get("decision") != "matched": + continue + figure_no = ocr_figures._extract_figure_number(str(legend.get("text", ""))) + proposals.append(ClaimProposal( + pass_name=self.name, + figure_no=figure_no, + claim_type="match", + legends=[ResourceRef(kind="legend", page=page, block_id=legend.get("block_id"), figure_no=figure_no)], + assets=[ResourceRef(kind="asset", page=page, block_id=bid) for bid in group.get("asset_block_ids", [])], + groups=[ResourceRef(kind="group", page=page, block_id=None, group_id=group.get("group_id"))], + confidence=float(score.get("score", 0.0)), + evidence_rank=1, + reason="same_page_primary", + diagnostics={ + "evidence": list(score.get("evidence", [])), + "legend_block_id": str(legend.get("block_id", "")), + }, + )) + return proposals + + def _materialize_match(self, state, proposal): + from . import ocr_figures + + legend = proposal.legends[0] + page = legend.page + asset_ids = {str(r.block_id) for r in proposal.assets} + matched_assets = [ + ocr_figures._project_asset_record(a) + for a in state.corpus.raw_assets + if _resource_page(a) == page and str(a.get("block_id", "")) in asset_ids + ] + legend_text = next( + str(b.get("text", "")) + for b in state.candidate_index.deduped_legends + if str(b.get("block_id", "")) == legend.block_id + ) + namespace = ocr_figures._extract_figure_namespace(legend_text) + return { + "figure_id": ocr_figures._format_figure_id(namespace, proposal.figure_no), + "figure_namespace": namespace, + "figure_number": proposal.figure_no, + "legend_block_id": legend.block_id, + "page": page, + "text": legend_text, + "matched_assets": matched_assets, + "asset_block_ids": sorted(asset_ids), + "settlement_type": "same_page", + "confidence": proposal.confidence, + "match_score": {"score": proposal.confidence, "decision": "matched", "evidence": proposal.diagnostics["evidence"]}, + "flags": [], + "bridge_block_ids": [], + } + + def run(self, state): + report = PassReport(pass_name=self.name) + proposals = self._collect_proposals(state) + report.proposals.extend(proposals) + + for proposal in sorted(proposals, key=lambda p: (p.evidence_rank, -p.confidence, -(p.figure_no or -1))): + conflict = state.ledger.try_claim_assets(proposal.assets, owner=proposal.legends[0], reason=proposal.reason) + if conflict is not None: + report.conflicts.append(conflict) + report.rejected.append(proposal) + continue + state.accept_match(proposal, self._materialize_match(state, proposal)) + report.accepted.append(proposal) + + return report diff --git a/paperforge/worker/ocr_figures.py b/paperforge/worker/ocr_figures.py index 6a6ffe0f..e17fd43e 100644 --- a/paperforge/worker/ocr_figures.py +++ b/paperforge/worker/ocr_figures.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextlib import itertools import re +from dataclasses import asdict from pathlib import Path from typing import Any @@ -2980,23 +2981,37 @@ def _infer_missing_main_figure_numbers( def build_figure_inventory(structured_blocks: list[dict], page_width: float = 1200, page_pdf_lines_by_page: dict[int, list[dict]] | None = None) -> dict[str, Any]: return build_figure_inventory_legacy(structured_blocks, page_width, page_pdf_lines_by_page) +def build_figure_inventory_vnext(structured_blocks: list[dict], page_width: float = 1200) -> dict[str, Any]: + from .ocr_figure_vnext_corpus import FigureCandidateIndex, FigureCorpus + from .ocr_figure_vnext_passes import PrimarySamePagePass, _resource_page + from .ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger + + corpus = FigureCorpus.from_blocks(structured_blocks, page_width=page_width) + candidate_index = FigureCandidateIndex.from_corpus(corpus) + state = FigurePipelineState(corpus=corpus, candidate_index=candidate_index, ledger=OwnershipLedger()) + report = PrimarySamePagePass().run(state) + matched_ids = {str(m.get("legend_block_id", "")) for m in state.matches} -def build_figure_inventory_vnext(structured_blocks: list[dict], page_width: float = 1200, page_pdf_lines_by_page: dict[int, list[dict]] | None = None) -> dict[str, Any]: return { "pipeline_mode": "vnext", - "matched_figures": [], + "matched_figures": state.matches, "ambiguous_figures": [], - "unmatched_legends": [], - "unmatched_assets": [], + "unmatched_legends": [b for b in candidate_index.deduped_legends if str(b.get("block_id", "")) not in matched_ids], + "unmatched_assets": [ + a for a in corpus.raw_assets + if (_resource_page(a) is not None + and state.ledger.owner_of_asset(page=_resource_page(a), block_id=a.get("block_id")) is None) + ], "unresolved_clusters": [], - "held_figures": [], - "rejected_legends": [], + "held_figures": list(candidate_index.held_legends), + "rejected_legends": list(candidate_index.rejected_legends), "page_ledger": {}, "residual_ledger": {}, "local_pairing_hypotheses": [], + "pass_reports": [asdict(report)], "completeness": { - "total_numbered_legends": 0, - "accounted_for": 0, + "total_numbered_legends": len(candidate_index.deduped_legends), + "accounted_for": len(state.matches), "details": [], }, } diff --git a/tests/test_ocr_figure_vnext_compare.py b/tests/test_ocr_figure_vnext_compare.py new file mode 100644 index 00000000..3b438fb1 --- /dev/null +++ b/tests/test_ocr_figure_vnext_compare.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from paperforge.worker.ocr_figures import build_figure_inventory_legacy, build_figure_inventory_vnext +from scripts.dev.compare_figure_inventory_legacy_vs_vnext import compare_inventories + + +def test_compare_inventories_reports_counts_for_same_page_case(): + blocks = [ + {"block_id": "c1", "page": 1, "role": "figure_caption", "text": "Figure 1. Caption", "bbox": [0, 100, 200, 150]}, + {"block_id": "a1", "page": 1, "role": "figure_asset", "bbox": [0, 0, 200, 90], "raw_label": "image"}, + ] + legacy = build_figure_inventory_legacy(blocks, 1200) + vnext = build_figure_inventory_vnext(blocks, 1200) + diff = compare_inventories(legacy, vnext) + + assert diff["vnext_matched_count"] >= 1 + assert "vnext_consumed_block_ids" in diff diff --git a/tests/test_ocr_figure_vnext_passes.py b/tests/test_ocr_figure_vnext_passes.py new file mode 100644 index 00000000..8747d5af --- /dev/null +++ b/tests/test_ocr_figure_vnext_passes.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from paperforge.worker.ocr_figure_vnext_corpus import FigureCandidateIndex, FigureCorpus +from paperforge.worker.ocr_figure_vnext_passes import PrimarySamePagePass +from paperforge.worker.ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger + + +def test_primary_same_page_pass_matches_single_safe_group(): + blocks = [ + {"block_id": "c1", "page": 1, "role": "figure_caption", "text": "Figure 1. Caption", "bbox": [0, 100, 200, 150]}, + {"block_id": "a1", "page": 1, "role": "figure_asset", "bbox": [0, 0, 200, 90], "raw_label": "image"}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + + report = PrimarySamePagePass().run(state) + + assert len(report.accepted) == 1 + assert report.accepted[0].claim_type == "match" + assert len(state.matches) == 1 + + +def test_primary_same_page_pass_prefers_higher_score_when_two_legends_compete_for_one_asset(monkeypatch): + blocks = [ + {"block_id": "c1", "page": 1, "role": "figure_caption", "text": "Figure 1. Caption", "bbox": [0, 100, 200, 150]}, + {"block_id": "c2", "page": 1, "role": "figure_caption", "text": "Figure 2. Caption", "bbox": [0, 160, 200, 210]}, + {"block_id": "a1", "page": 1, "role": "figure_asset", "bbox": [0, 0, 200, 90], "raw_label": "image"}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + if not index.candidate_groups: + index.candidate_groups = [{"group_id": "g1", "page": 1, "asset_block_ids": ["a1"], "media_blocks": [{"block_id": "a1"}], "group_type": "single_asset", "cluster_bbox": [0, 0, 200, 90]}] + + scores = [{"score": 0.4, "decision": "matched", "evidence": ["low"]}, {"score": 0.9, "decision": "matched", "evidence": ["high"]}] + + def fake_score(*args, **kwargs): + return scores.pop(0) + + monkeypatch.setattr("paperforge.worker.ocr_figures._score_legend_to_group", fake_score) + monkeypatch.setattr("paperforge.worker.ocr_figures.score_figure_caption", lambda *a, **k: {"score": 0.9}) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = PrimarySamePagePass().run(state) + + assert len(report.accepted) == 1 + assert report.accepted[0].figure_no == 2