mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
47 lines
2.4 KiB
Python
47 lines
2.4 KiB
Python
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
|