mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
feat(ocr): add vnext same-page figure pass
This commit is contained in:
LLLin000
parent
a0cad9cacd
commit
29fc619d4d
4 changed files with 191 additions and 8 deletions
104
paperforge/worker/ocr_figure_vnext_passes.py
Normal file
104
paperforge/worker/ocr_figure_vnext_passes.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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": [],
|
||||
},
|
||||
}
|
||||
|
|
|
|||
17
tests/test_ocr_figure_vnext_compare.py
Normal file
17
tests/test_ocr_figure_vnext_compare.py
Normal file
|
|
@ -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
|
||||
47
tests/test_ocr_figure_vnext_passes.py
Normal file
47
tests/test_ocr_figure_vnext_passes.py
Normal file
|
|
@ -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
|
||||
Loading…
Reference in a new issue