feat(ocr): add vnext CompositeParentPass

This commit is contained in:
LLLin000 2026-07-03 16:39:32 +08:00
parent f3e99e4183
commit 602947b243
2 changed files with 411 additions and 0 deletions

View file

@ -0,0 +1,211 @@
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
def _is_legend_matched(legend_block_id: str, matches: list[dict]) -> bool:
return any(m.get("legend_block_id") == legend_block_id for m in matches)
def _numbered_legends_on_page(page: int, legends: list[dict], extract_fn) -> list[dict]:
"""Return legends on *page* that have a numeric figure number."""
result = []
for leg in legends:
leg_page = _resource_page(leg)
if leg_page != page:
continue
num = extract_fn(str(leg.get("text", "")))
if num is not None:
result.append(leg)
return result
class CompositeParentPass:
"""For each unmatched numbered legend, if a composite_parent_candidate
exists on the same page with sufficient confidence (>=0.60) and >=2 child
groups, claim all child group assets as a single multi-panel figure.
"""
name = "composite_parent"
def run(self, state):
from . import ocr_figures
report = PassReport(pass_name=self.name)
consumed_parent_keys: set[str] = set() # track by group_id
for legend in state.candidate_index.deduped_legends:
legend_page = _resource_page(legend)
if legend_page is None:
continue
legend_text = str(legend.get("text", ""))
figure_no = ocr_figures._extract_figure_number(legend_text)
if figure_no is None:
continue # only numbered legends
legend_block_id = str(legend.get("block_id", ""))
if _is_legend_matched(legend_block_id, state.matches):
continue
# --- find composite parent candidates on same page ---
page_parents = [
cp
for cp in state.candidate_index.composite_parent_candidates
if cp.get("page") == legend_page
and cp.get("group_id", "") not in consumed_parent_keys
]
if not page_parents:
continue
page_parents.sort(key=lambda cp: cp.get("parent_confidence", 0), reverse=True)
best_parent = page_parents[0]
parent_conf = float(best_parent.get("parent_confidence", 0.0))
if parent_conf < 0.60:
continue
# --- competing caption veto ---
same_page_numbered = _numbered_legends_on_page(
legend_page, state.candidate_index.deduped_legends,
ocr_figures._extract_figure_number,
)
has_competing = len(same_page_numbered) > 1
# --- resolve child groups ---
child_group_ids = set(best_parent.get("child_group_ids", []))
all_child_groups = [
g for g in state.candidate_index.candidate_groups
if str(g.get("group_id", "")) in child_group_ids
]
# Band-scoping: when competing captions exist, only keep child
# groups whose best_caption_band_id matches this legend.
if has_competing:
matching_band_groups = [
g for g in all_child_groups
if str((g.get("assist") or {}).get("best_caption_band_id") or "")
== legend_block_id
]
if len(matching_band_groups) >= 2:
effective_groups = matching_band_groups
else:
continue # competing captions, not scoped -> skip
else:
effective_groups = all_child_groups
# --- effective child group count check ---
is_dense = best_parent.get("parent_subtype") == "dense_composite"
fragment_count = int(best_parent.get("fragment_count", 0) or 0)
dense_single_group_ok = is_dense and len(effective_groups) == 1 and fragment_count >= 4
if len(effective_groups) < 2 and not dense_single_group_ok:
continue
# --- collect asset block ids (ordered, deduplicated) ---
asset_block_ids = list(dict.fromkeys(
bid
for group in effective_groups
for bid in group.get("asset_block_ids", [])
))
if not asset_block_ids:
continue
# --- build proposal ---
asset_refs = [
ResourceRef(kind="asset", page=legend_page, block_id=bid)
for bid in asset_block_ids
]
group_refs = [
ResourceRef(
kind="group", page=legend_page, block_id=None,
group_id=str(g.get("group_id", "")),
)
for g in effective_groups
]
legend_ref = ResourceRef(
kind="legend", page=legend_page,
block_id=legend_block_id, figure_no=figure_no,
)
proposal = ClaimProposal(
pass_name=self.name,
figure_no=figure_no,
claim_type="composite_parent",
legends=[legend_ref],
assets=asset_refs,
groups=group_refs,
confidence=parent_conf,
evidence_rank=1,
reason="composite_parent",
diagnostics={
"evidence": ["composite_parent"],
"legend_block_id": legend_block_id,
},
)
report.proposals.append(proposal)
# --- claim assets ---
conflict = state.ledger.try_claim_assets(
proposal.assets, owner=legend_ref, reason=proposal.reason,
)
if conflict is not None:
report.conflicts.append(conflict)
report.rejected.append(proposal)
continue
# --- materialise match record and accept ---
match_record = self._build_match_record(state, proposal, figure_no)
state.accept_match(proposal, match_record)
report.accepted.append(proposal)
consumed_parent_keys.add(best_parent.get("group_id", ""))
return report
@staticmethod
def _build_match_record(state, proposal: ClaimProposal, figure_no: int) -> dict:
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)
figure_id = ocr_figures._format_figure_id(namespace, figure_no)
return {
"figure_id": figure_id,
"figure_namespace": namespace,
"figure_number": 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": "composite_parent",
"confidence": proposal.confidence,
"match_score": {
"score": proposal.confidence,
"decision": "matched",
"evidence": ["composite_parent"],
},
"flags": ["composite_parent_match"],
"bridge_block_ids": [],
}

View file

@ -0,0 +1,200 @@
from __future__ import annotations
from paperforge.worker.ocr_figure_vnext_composite_pass import CompositeParentPass
from paperforge.worker.ocr_figure_vnext_corpus import FigureCandidateIndex, FigureCorpus
from paperforge.worker.ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger
def _cgroup(group_id: str, page: int, *asset_ids: str) -> dict:
"""Build a synthetic candidate group dict."""
return {
"group_id": group_id,
"page": page,
"asset_block_ids": list(asset_ids),
"media_blocks": [{"block_id": aid} for aid in asset_ids],
"group_type": "multi_asset" if len(asset_ids) > 1 else "single_asset",
"cluster_bbox": [0, 0, 100, 100],
}
def _composite_parent(
group_id: str, page: int, child_group_ids: list[str], confidence: float = 0.85,
subtype: str = "composite",
) -> dict:
"""Build a synthetic composite parent candidate dict."""
return {
"group_id": group_id,
"page": page,
"child_group_ids": list(child_group_ids),
"asset_block_ids": [],
"cluster_bbox": [0, 0, 200, 200],
"parent_confidence": confidence,
"parent_subtype": subtype,
}
def _make_index(
legends: list[dict],
candidate_groups: list[dict],
composite_parents: list[dict],
*,
page_width: float = 1200,
) -> tuple[FigureCorpus, FigureCandidateIndex]:
"""Build a FigureCorpus + FigureCandidateIndex from synthetic data.
Assets are auto-extracted from blocks that have role="figure_asset".
"""
assets = [
b for b in legends if b.get("role") == "figure_asset"
]
blocks = list(legends)
for g in candidate_groups:
for aid in g.get("asset_block_ids", []):
if not any(b.get("block_id") == aid for b in blocks):
blocks.append({
"block_id": aid,
"page": g.get("page", 1),
"role": "figure_asset",
"bbox": [0, 0, 100, 100],
})
formal_legends = [b for b in legends if b.get("role") in {"figure_caption", "figure_caption_candidate"}]
corpus = FigureCorpus.from_blocks(blocks, page_width=page_width)
index = FigureCandidateIndex(
formal_legends=formal_legends,
held_legends=[],
rejected_legends=[b for b in legends if b not in formal_legends],
deduped_legends=formal_legends,
candidate_groups=candidate_groups,
competing_caption_pages={
int(b.get("page")) for b in formal_legends
if b.get("page") is not None
},
sidecar_candidates={},
bundle_source_legend_ids=set(),
locator_candidates=[],
composite_parent_candidates=composite_parents,
)
return corpus, index
def test_composite_parent_happy_path():
"""2x2 grid of assets + 1 numbered caption on same page -> composite parent match."""
legends = [
{"block_id": "c1", "page": 1, "role": "figure_caption",
"text": "Figure 1. Caption", "bbox": [0, 500, 300, 550]},
]
candidate_groups = [
_cgroup("g1", 1, "a1", "a2"),
_cgroup("g2", 1, "a3", "a4"),
]
composite_parents = [
_composite_parent("cp1", 1, ["g1", "g2"], confidence=0.85),
]
corpus, index = _make_index(legends, candidate_groups, composite_parents)
state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger())
report = CompositeParentPass().run(state)
assert len(report.accepted) == 1, "expected one accepted proposal"
assert report.accepted[0].claim_type == "composite_parent"
assert len(state.matches) == 1
match = state.matches[0]
assert match["settlement_type"] == "composite_parent"
assert match["flags"] == ["composite_parent_match"]
assert match["figure_number"] == 1
assert set(match["asset_block_ids"]) == {"a1", "a2", "a3", "a4"}
assert match["confidence"] == 0.85
def test_composite_parent_competing_captions_skipped():
"""Two numbered legends on same page -> competing caption veto -> no match."""
legends = [
{"block_id": "c1", "page": 1, "role": "figure_caption",
"text": "Figure 1. First", "bbox": [0, 500, 300, 550]},
{"block_id": "c2", "page": 1, "role": "figure_caption",
"text": "Figure 2. Second", "bbox": [0, 560, 300, 610]},
]
candidate_groups = [
_cgroup("g1", 1, "a1", "a2"),
_cgroup("g2", 1, "a3", "a4"),
]
composite_parents = [
_composite_parent("cp1", 1, ["g1", "g2"], confidence=0.85),
]
corpus, index = _make_index(legends, candidate_groups, composite_parents)
state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger())
report = CompositeParentPass().run(state)
assert len(report.accepted) == 0, "expected no accepted with competing captions"
assert len(state.matches) == 0
def test_composite_parent_insufficient_child_groups_skipped():
"""Composite parent with only 1 child group and not dense_composite -> skip."""
legends = [
{"block_id": "c1", "page": 1, "role": "figure_caption",
"text": "Figure 1. Caption", "bbox": [0, 500, 300, 550]},
]
candidate_groups = [
_cgroup("g1", 1, "a1", "a2"),
]
composite_parents = [
_composite_parent("cp1", 1, ["g1"], confidence=0.85, subtype="composite"),
]
corpus, index = _make_index(legends, candidate_groups, composite_parents)
state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger())
report = CompositeParentPass().run(state)
assert len(report.accepted) == 0, "expected no accepted with 1 child group"
assert len(state.matches) == 0
def test_composite_parent_low_confidence_skipped():
"""Composite parent with confidence below 0.60 -> skip."""
legends = [
{"block_id": "c1", "page": 1, "role": "figure_caption",
"text": "Figure 1. Caption", "bbox": [0, 500, 300, 550]},
]
candidate_groups = [
_cgroup("g1", 1, "a1", "a2"),
_cgroup("g2", 1, "a3", "a4"),
]
composite_parents = [
_composite_parent("cp1", 1, ["g1", "g2"], confidence=0.50),
]
corpus, index = _make_index(legends, candidate_groups, composite_parents)
state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger())
report = CompositeParentPass().run(state)
assert len(report.accepted) == 0, "expected no accepted with low confidence"
assert len(state.matches) == 0
def test_composite_parent_skips_already_matched_legend():
"""Legend already in state.matches -> skip composite pass for it."""
legends = [
{"block_id": "c1", "page": 1, "role": "figure_caption",
"text": "Figure 1. Caption", "bbox": [0, 500, 300, 550]},
]
candidate_groups = [
_cgroup("g1", 1, "a1", "a2"),
_cgroup("g2", 1, "a3", "a4"),
]
composite_parents = [
_composite_parent("cp1", 1, ["g1", "g2"], confidence=0.85),
]
corpus, index = _make_index(legends, candidate_groups, composite_parents)
state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger())
# Pre-seed a match for this legend
state.matches.append({
"legend_block_id": "c1",
"settlement_type": "same_page",
"figure_number": 1,
})
report = CompositeParentPass().run(state)
assert len(report.accepted) == 0, "expected no accepted for already-matched legend"