diff --git a/paperforge/worker/ocr_figure_domain.py b/paperforge/worker/ocr_figure_domain.py index 0906544b..10cb7761 100644 --- a/paperforge/worker/ocr_figure_domain.py +++ b/paperforge/worker/ocr_figure_domain.py @@ -105,6 +105,19 @@ class FigureCandidateIndex: if len(narrow_set) >= 2: sidecar_candidates[sp] = narrow_set + # Single-caption 37-like sidecar rescue: pages with one narrow caption + # that has a same-row right-side image (left-caption/right-image pattern). + # Primary same-page matching misses these — the sidecar pass can rescue them. + for sp, spl in _legends_by_page.items(): + if sp in sidecar_candidates: + continue + for leg in spl: + if ocr_figures._is_single_narrow_sidecar_caption( + leg, corpus.raw_assets, page_width=corpus.page_width + ): + sidecar_candidates[sp] = [leg] + break + # Populate bundle_source_legend_ids: legends on pages with >=3 numbered legends # and zero assets on that page. bundle_source_legend_ids = ocr_figures._identify_bundle_source_legend_ids( diff --git a/paperforge/worker/ocr_figure_vnext_group_seq_pass.py b/paperforge/worker/ocr_figure_vnext_group_seq_pass.py index c7a00d1b..c30db7e1 100644 --- a/paperforge/worker/ocr_figure_vnext_group_seq_pass.py +++ b/paperforge/worker/ocr_figure_vnext_group_seq_pass.py @@ -158,6 +158,23 @@ class GroupSequentialPass: scored.sort(key=lambda x: x[1], reverse=True) best_group = scored[0][0] + # (a') Multi-legend multi-asset fallback: if same-page scoring + # rejected but the page has a multi-asset distance_cluster group + # with multiple short numbered legends, still select it so + # section 6b can do per-asset splitting. + if best_group is None and same_page: + _page_short_legends = [ + l for l in deduped_legends + if _resource_page(l) == lg_page + and ocr_figures._extract_figure_number(str(l.get("text", "") or "")) is not None + and ocr_figures._is_short_numbered_figure_caption(l) + ] + if len(_page_short_legends) >= 2: + for sg in same_page: + if sg.get("group_type") == "distance_cluster" and len(sg.get("media_blocks", [])) >= 2: + best_group = sg + break + # (b) Next-page fallback: first group on next page if best_group is None and next_page: best_group = next_page[0] @@ -178,6 +195,10 @@ class GroupSequentialPass: if best_group is None: continue + # Flag: when set, the group stays in pool for sibling legends + # (multi-legend multi-asset restriction). + _keep_group_in_pool = False + # ----- 6. Collect group assets ----- group_page = int(best_group.get("page", 0) or 0) group_assets = [] @@ -191,6 +212,45 @@ class GroupSequentialPass: if not group_assets: continue + # ----- 6a. Filter already-owned assets for idempotent group re-access ----- + group_assets = [ + a for a in group_assets + if state.ledger.owner_of_asset( + page=group_page, + block_id=str(a.get("block_id", "")), + ) is None + ] + if not group_assets: + continue + + # ----- 6b. Multi-legend multi-asset veto ----- + # When n explicit numbered legends share the same page as a group + # with n unowned assets, block the whole-group claim and restrict + # this legend to the geometrically closest single asset, keeping + # the group in the pool for sibling legends. + if len(group_assets) >= 2: + _page_legends = [ + l for l in deduped_legends + if str(l.get("block_id", "")) not in matched_legend_ids + and _resource_page(l) == lg_page + and ocr_figures._extract_figure_number(str(l.get("text", "") or "")) is not None + ] + if len(_page_legends) >= len(group_assets): + _all_explicit = all( + ocr_figures._is_short_numbered_figure_caption(l) + for l in _page_legends + ) + if _all_explicit: + # Manhattan center distance to find the single closest asset + l_bbox = legend.get("bbox") or legend.get("block_bbox") or [0, 0, 0, 0] + l_cx = (l_bbox[0] + l_bbox[2]) / 2 + l_cy = (l_bbox[1] + l_bbox[3]) / 2 + def _mdist(a): + ab = a.get("bbox") or a.get("block_bbox") or [0, 0, 0, 0] + return abs(l_cx - (ab[0] + ab[2]) / 2) + abs(l_cy - (ab[1] + ab[3]) / 2) + group_assets = [min(group_assets, key=_mdist)] + _keep_group_in_pool = True + # ----- 7. Build proposal and match record ----- asset_refs = [ ResourceRef(kind="asset", page=group_page, block_id=str(a.get("block_id", ""))) @@ -277,6 +337,9 @@ class GroupSequentialPass: report.accepted.append(proposal) # Remove claimed group from pool so subsequent legends don't compete - unmatched_groups = [g for g in unmatched_groups if g is not best_group] + if not _keep_group_in_pool: + # When we restricted a multi-asset group to a single asset, keep + # the group in the pool so sibling legends can claim remaining assets. + unmatched_groups = [g for g in unmatched_groups if g is not best_group] return report diff --git a/paperforge/worker/ocr_figure_vnext_passes.py b/paperforge/worker/ocr_figure_vnext_passes.py index cd1ed6fa..8624e598 100644 --- a/paperforge/worker/ocr_figure_vnext_passes.py +++ b/paperforge/worker/ocr_figure_vnext_passes.py @@ -18,14 +18,17 @@ class PrimarySamePagePass: proposals = [] - # Precompute per-page context for short-caption geometry scoring _numbered_count_by_page: dict[int, int] = {} + _short_caption_count_by_page: dict[int, int] = {} _page_blocks_map: dict[int, list[dict]] = {} _page_height_map: dict[int, float] = {} for _leg in state.candidate_index.deduped_legends: _lp = _resource_page(_leg) - if _lp is not None and ocr_figures._extract_figure_number(str(_leg.get("text", ""))) is not None: - _numbered_count_by_page[_lp] = _numbered_count_by_page.get(_lp, 0) + 1 + if _lp is not None: + if ocr_figures._extract_figure_number(str(_leg.get("text", ""))) is not None: + _numbered_count_by_page[_lp] = _numbered_count_by_page.get(_lp, 0) + 1 + if ocr_figures._is_short_numbered_figure_caption(_leg): + _short_caption_count_by_page[_lp] = _short_caption_count_by_page.get(_lp, 0) + 1 for _b in state.corpus.blocks: _bp = _resource_page(_b) if _bp is not None: @@ -43,6 +46,19 @@ class PrimarySamePagePass: continue page_groups = [g for g in state.candidate_index.candidate_groups if _resource_page(g) == page] for group in page_groups: + # VETO: When multiple explicit short numbered legends share a page, + # no single legend may claim an entire multi-asset distance_cluster + # group. Such groups should be distributed by downstream passes + # (GroupSequentialPass) which handle per-asset splitting. + _gt = group.get("group_type", "") + _n_assets = len(group.get("media_blocks", [])) + if ( + _gt == "distance_cluster" + and _n_assets >= 2 + and _short_caption_count_by_page.get(page, 0) > 1 + and ocr_figures._is_short_numbered_figure_caption(legend) + ): + continue score = ocr_figures._score_legend_to_group( legend, group, diff --git a/paperforge/worker/ocr_figure_vnext_sidecar_pass.py b/paperforge/worker/ocr_figure_vnext_sidecar_pass.py index 2c5bb87d..2e80a999 100644 --- a/paperforge/worker/ocr_figure_vnext_sidecar_pass.py +++ b/paperforge/worker/ocr_figure_vnext_sidecar_pass.py @@ -43,23 +43,36 @@ class SidecarPass: ) bands = ocr_figures._partition_assets_by_caption_bands( - narrow_captions, page_assets, page_height + narrow_captions, page_assets, page_height, + page_width=state.corpus.page_width, ) for caption in narrow_captions: cid = str(caption.get("block_id", "")) - # Skip captions already protected by an earlier pass - protected = False - for match in state.matches: - if ( - str(match.get("legend_block_id", "")) == cid - and ocr_figures._has_protected_figure_ownership(match) - ): - protected = True - break - if protected: - continue + # For single-caption entries (37-like rescue), skip if already + # matched by any earlier pass — this rescue is only a last resort + # for captions the primary pass entirely missed. + if len(narrow_captions) == 1: + already_matched = any( + str(m.get("legend_block_id", "")) == cid + for m in state.matches + ) + if already_matched: + continue + else: + # Multi-caption path: skip only protected matches + # (allows refinement of multi-asset figures). + protected = False + for match in state.matches: + if ( + str(match.get("legend_block_id", "")) == cid + and ocr_figures._has_protected_figure_ownership(match) + ): + protected = True + break + if protected: + continue band_assets = bands.get(cid, []) if not band_assets: diff --git a/paperforge/worker/ocr_figures.py b/paperforge/worker/ocr_figures.py index cbf3d7af..7d5c0be3 100644 --- a/paperforge/worker/ocr_figures.py +++ b/paperforge/worker/ocr_figures.py @@ -1196,10 +1196,17 @@ def _score_legend_to_group( # in the same column as the group with a tight vertical gap, and is # page-locally unique (no competing numbered legend in the same column). if _is_short_numbered_figure_caption(legend): - # Do not override page_assets groups (have their own gate) or - # safe_auto_match groups (already strongly matched). + # Do not override page_assets groups (have their own gate). + # safe_auto_match groups: skip geometry bypass when multiple short + # numbered legends share the page, so each short caption is scored + # against its column-local subset rather than one claiming everything. _sc_gt = group.get("group_type", "") - if _sc_gt != "page_assets" and not (group.get("safe_auto_match") and len(group.get("media_blocks", [])) >= 2): + _sc_geom_bypass = ( + group.get("safe_auto_match") + and len(group.get("media_blocks", [])) >= 2 + and not (page_numbered_legend_count > 1 and _is_short_numbered_figure_caption(legend)) + ) + if _sc_gt != "page_assets" and not _sc_geom_bypass: _leg_bbox = legend.get("bbox") or legend.get("block_bbox") or [0, 0, 0, 0] _grp_bbox = group.get("cluster_bbox", [0, 0, 0, 0]) _cap_top = _leg_bbox[1] @@ -1248,11 +1255,15 @@ def _score_legend_to_group( if gt == "distance_cluster": num_assets = len(group.get("media_blocks", [])) if group.get("safe_auto_match") and num_assets >= 2: - return { - "score": 0.85, - "decision": "matched", - "evidence": ["same_page", "distance_clustered", "safe_auto_match"], - } + # When multiple short numbered legends share a page, a single short + # caption must not auto-claim the whole group. Fall through to + # normal distance_cluster scoring (geometry + orientation). + if not (page_numbered_legend_count > 1 and _is_short_numbered_figure_caption(legend)): + return { + "score": 0.85, + "decision": "matched", + "evidence": ["same_page", "distance_clustered", "safe_auto_match"], + } cluster_bbox = group.get("cluster_bbox", [0, 0, 0, 0]) match_score = _score_legend_to_asset_with_orientation( @@ -1905,6 +1916,25 @@ def _same_page_narrow_caption_column(page_captions: list[dict], page_width: floa return sorted(best, key=lambda cap: (cap.get("bbox") or [0, 0, 0, 0])[1]) +def _is_single_narrow_sidecar_caption( + caption: dict, page_assets: list[dict], *, page_width: float = 1200 +) -> bool: + """Check if a single narrow caption qualifies for sidecar rescue (37-like pattern). + + The 37-like invariant: one formal narrow caption on the page, an image + to the right at the same row (no x-overlap, bounded x-gap, strong y-overlap). + Primary same-page matching misses these because it expects overlap-style + evidence. We rescue only when the geometry is strong. + """ + if _extract_figure_number(str(caption.get("text", ""))) is None: + return False + if _caption_width_ratio(caption, page_width) >= 0.6: + return False + if _PANEL_LABEL_PATTERN.match(str(caption.get("text", "")).strip()): + return False + coupled = _caption_row_coupled_assets(caption, page_assets, page_width=page_width) + return len(coupled) > 0 + def _caption_row_coupled_assets(caption: dict, assets: list[dict], *, page_width: float = 1200) -> list[dict]: caption_bbox = caption.get("bbox") or [0, 0, 0, 0] if len(caption_bbox) < 4: @@ -2362,9 +2392,15 @@ def _allow_previous_page_sequential_match(cap: dict, asset: dict) -> bool: def _partition_assets_by_caption_bands( - page_captions: list[dict], page_assets: list[dict], page_height: float + page_captions: list[dict], page_assets: list[dict], page_height: float, + page_width: float = 1200, ) -> dict[str, list[dict]]: if len(page_captions) < 2: + # Single-caption sidecar rescue: use row-coupled assets (37-like pattern). + if len(page_captions) == 1: + caption = page_captions[0] + coupled = _caption_row_coupled_assets(caption, page_assets, page_width=page_width) + return {str(caption.get("block_id", "")): coupled} return {} ordered = sorted(page_captions, key=lambda cap: (cap.get("bbox") or [0, 0, 0, 0])[1]) diff --git a/tests/test_ocr_figure_vnext_group_seq_pass.py b/tests/test_ocr_figure_vnext_group_seq_pass.py index 5f39a0cb..d0ca401a 100644 --- a/tests/test_ocr_figure_vnext_group_seq_pass.py +++ b/tests/test_ocr_figure_vnext_group_seq_pass.py @@ -168,3 +168,180 @@ def test_group_sequential_prev_page_guard_denied(monkeypatch): assert len(report.accepted) == 0, f"expected 0 accepted, got {len(report.accepted)}" assert len(state.matches) == 0, f"expected 0 matches, got {len(state.matches)}" + + +def test_multi_legend_veto_two_legends_two_assets(): + """Two short numbered explicit legends, distance_cluster with 2 assets on same page. + -> whole-group vetoed, each legend claims one asset by position.""" + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 2", "bbox": [0, 100, 200, 130], "zone": "display_zone"}, + {"block_id": "c2", "page": 5, "role": "figure_caption", + "text": "Figure 3", "bbox": [300, 100, 500, 130], "zone": "display_zone"}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [0, 200, 250, 400], "raw_label": "image"}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [300, 200, 550, 400], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 5, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2"], + "cluster_bbox": [0, 200, 550, 400], + "media_blocks": [{"block_id": "a1"}, {"block_id": "a2"}], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = GroupSequentialPass().run(state) + + # Both legends should have been matched (one each, no whole-group claim) + assert len(report.accepted) == 2, f"expected 2 accepted, got {len(report.accepted)}" + assert len(state.matches) == 2, f"expected 2 matches, got {len(state.matches)}" + + # Each match should have exactly one asset + for m in state.matches: + assert m["settlement_type"] == "group_sequential" + assert "group_sequential_match" in m["flags"] + assert len(m.get("matched_assets", [])) == 1, \ + f"expected 1 asset per match, got {len(m.get('matched_assets', []))}" + + # c1 (Figure 2) should get a1 (left asset), c2 (Figure 3) should get a2 (right asset) + match_by_legend = {m["legend_block_id"]: m for m in state.matches} + assert match_by_legend["c1"]["asset_block_ids"] == ["a1"], \ + f"expected c1 -> a1, got {match_by_legend['c1']['asset_block_ids']}" + assert match_by_legend["c2"]["asset_block_ids"] == ["a2"], \ + f"expected c2 -> a2, got {match_by_legend['c2']['asset_block_ids']}" + + +def test_multi_legend_veto_grouped_figure_no_veto(): + """Single descriptive legend + multi-asset distance_cluster. + -> normal grouped figure, one legend claims all assets (no veto).""" + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 2. Overview of system architecture. (A) Input (B) Processing.", + "bbox": [0, 100, 500, 150]}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [0, 200, 250, 400], "raw_label": "image"}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [300, 200, 550, 400], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 5, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2"], + "cluster_bbox": [0, 200, 550, 400], + "media_blocks": [{"block_id": "a1"}, {"block_id": "a2"}], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = GroupSequentialPass().run(state) + + # Single descriptive legend → normal grouped figure, both assets claimed + assert len(report.accepted) == 1, f"expected 1 accepted, got {len(report.accepted)}" + assert len(state.matches) == 1, f"expected 1 match, got {len(state.matches)}" + + match = state.matches[0] + assert match["settlement_type"] == "group_sequential" + assert "group_sequential_match" in match["flags"] + assert match["legend_block_id"] == "c1" + + matched_bids = {a["block_id"] for a in match.get("matched_assets", [])} + assert matched_bids == {"a1", "a2"}, f"expected a1,a2, got {matched_bids}" + + +def test_multi_legend_veto_ambiguous_no_veto(): + """Two short legends but group has 3 assets (n_legends < n_assets). + -> no veto; first legend claims the whole group.""" + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 2", "bbox": [0, 100, 200, 130], "zone": "display_zone"}, + {"block_id": "c2", "page": 5, "role": "figure_caption", + "text": "Figure 3", "bbox": [300, 100, 500, 130], "zone": "display_zone"}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [0, 200, 180, 300], "raw_label": "image"}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [200, 200, 380, 300], "raw_label": "image"}, + {"block_id": "a3", "page": 5, "role": "figure_asset", + "bbox": [400, 200, 580, 300], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 5, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2", "a3"], + "cluster_bbox": [0, 200, 580, 300], + "media_blocks": [{"block_id": "a1"}, {"block_id": "a2"}, {"block_id": "a3"}], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = GroupSequentialPass().run(state) + + # 2 legends < 3 assets → no veto, first legend takes all + assert len(report.accepted) == 1, f"expected 1 accepted, got {len(report.accepted)}" + assert len(state.matches) == 1, f"expected 1 match, got {len(state.matches)}" + + match = state.matches[0] + assert match["settlement_type"] == "group_sequential" + assert match["legend_block_id"] == "c1" + + matched_bids = {a["block_id"] for a in match.get("matched_assets", [])} + assert matched_bids == {"a1", "a2", "a3"}, f"expected a1,a2,a3, got {matched_bids}" diff --git a/tests/test_ocr_figure_vnext_passes.py b/tests/test_ocr_figure_vnext_passes.py index 52ba5206..1e5d434f 100644 --- a/tests/test_ocr_figure_vnext_passes.py +++ b/tests/test_ocr_figure_vnext_passes.py @@ -493,3 +493,253 @@ def test_short_caption_does_not_override_safe_auto_match(): evidence = report.accepted[0].diagnostics.get("evidence", []) # safe_auto_match evidence includes "safe_auto_match", not "short_caption_geometry" assert not any("short_caption_geometry" in e for e in evidence), f"Unexpected short_caption_geometry: {evidence}" + + +# --- PrimarySamePagePass multi-legend multi-asset veto tests (PR-2a follow-up) --- + + +def test_primary_same_page_veto_two_short_legends_two_assets(): + """Two short numbered explicit legends, distance_cluster with 2 assets on same page. + -> PrimarySamePagePass must veto whole-group claim.""" + from paperforge.worker.ocr_figure_vnext_corpus import FigureCorpus, FigureCandidateIndex + from paperforge.worker.ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger + + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 2", "bbox": [0, 100, 200, 130], "zone": "display_zone"}, + {"block_id": "c2", "page": 5, "role": "figure_caption", + "text": "Figure 3", "bbox": [300, 100, 500, 130], "zone": "display_zone"}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [0, 200, 250, 400], "raw_label": "image"}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [300, 200, 550, 400], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 5, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2"], + "cluster_bbox": [0, 200, 550, 400], + "media_blocks": [{"block_id": "a1"}, {"block_id": "a2"}], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = PrimarySamePagePass().run(state) + + # PrimarySamePagePass should not make a whole-group claim + assert len(report.accepted) == 0, \ + f"expected 0 accepted (veto), got {len(report.accepted)}" + + +def test_primary_same_page_veto_four_legends_four_assets(): + """Four short numbered explicit legends, distance_cluster with 4 assets on same page. + -> PrimarySamePagePass must veto; no single legend claims the whole group.""" + from paperforge.worker.ocr_figure_vnext_corpus import FigureCorpus, FigureCandidateIndex + from paperforge.worker.ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger + + blocks = [ + {"block_id": "c1", "page": 3, "role": "figure_caption", + "text": "Figure 4", "bbox": [0, 100, 200, 130], "zone": "display_zone"}, + {"block_id": "c2", "page": 3, "role": "figure_caption", + "text": "Figure 5", "bbox": [300, 100, 500, 130], "zone": "display_zone"}, + {"block_id": "c3", "page": 3, "role": "figure_caption", + "text": "Figure 6", "bbox": [0, 160, 200, 190], "zone": "display_zone"}, + {"block_id": "c4", "page": 3, "role": "figure_caption", + "text": "Figure 7", "bbox": [300, 160, 500, 190], "zone": "display_zone"}, + {"block_id": "a1", "page": 3, "role": "figure_asset", + "bbox": [0, 200, 280, 350], "raw_label": "image"}, + {"block_id": "a2", "page": 3, "role": "figure_asset", + "bbox": [320, 200, 600, 350], "raw_label": "image"}, + {"block_id": "a3", "page": 3, "role": "figure_asset", + "bbox": [0, 360, 280, 510], "raw_label": "image"}, + {"block_id": "a4", "page": 3, "role": "figure_asset", + "bbox": [320, 360, 600, 510], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 3, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2", "a3", "a4"], + "cluster_bbox": [0, 200, 600, 510], + "media_blocks": [ + {"block_id": "a1"}, {"block_id": "a2"}, + {"block_id": "a3"}, {"block_id": "a4"}, + ], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = PrimarySamePagePass().run(state) + + # PrimarySamePagePass should not make a whole-group claim + assert len(report.accepted) == 0, \ + f"expected 0 accepted (veto), got {len(report.accepted)}" + + +def test_primary_same_page_veto_grouped_figure_unaffected(): + """Single descriptive legend with multi-asset distance_cluster group. + -> Normal grouped figure, one legend claims all assets (no veto).""" + from paperforge.worker.ocr_figure_vnext_corpus import FigureCorpus, FigureCandidateIndex + from paperforge.worker.ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger + + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 2. Overview of system architecture. (A) Input (B) Processing.", + "bbox": [0, 100, 500, 150]}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [0, 200, 250, 400], "raw_label": "image"}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [300, 200, 550, 400], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 5, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2"], + "cluster_bbox": [0, 200, 550, 400], + "media_blocks": [{"block_id": "a1"}, {"block_id": "a2"}], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = PrimarySamePagePass().run(state) + + # Single descriptive legend with multi-asset group — should still match normally + # (not affected by the veto, which only applies to short numbered captions) + assert len(report.accepted) >= 1, \ + f"expected at least 1 accepted (grouped figure), got {len(report.accepted)}" + if report.accepted: + proposal = report.accepted[0] + # Should claim both assets + assert len(proposal.assets) == 2, \ + f"expected 2 assets in grouped figure, got {len(proposal.assets)}" + # Should be a same_page primary match + assert proposal.reason == "same_page_primary" + + +def test_primary_same_page_veto_two_legends_then_group_seq_splits(monkeypatch): + """Full pipeline: PrimarySamePagePass vetoes whole-group; GroupSequentialPass splits. + 2 short legends + 2 assets -> each legend gets 1 asset.""" + from paperforge.worker.ocr_figure_vnext_corpus import FigureCorpus, FigureCandidateIndex + from paperforge.worker.ocr_figure_vnext_state import FigurePipelineState, OwnershipLedger + from paperforge.worker.ocr_figure_vnext_group_seq_pass import GroupSequentialPass + + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 2", "bbox": [0, 100, 200, 130], "zone": "display_zone"}, + {"block_id": "c2", "page": 5, "role": "figure_caption", + "text": "Figure 3", "bbox": [300, 100, 500, 130], "zone": "display_zone"}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [0, 200, 250, 400], "raw_label": "image"}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [300, 200, 550, 400], "raw_label": "image"}, + {"block_id": "filler", "page": 1, "role": "paper_title", + "text": "Title", "bbox": [0, 0, 500, 50]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + + groups = [ + { + "group_id": "g1", + "page": 5, + "group_type": "distance_cluster", + "asset_block_ids": ["a1", "a2"], + "cluster_bbox": [0, 200, 550, 400], + "media_blocks": [{"block_id": "a1"}, {"block_id": "a2"}], + "safe_auto_match": True, + } + ] + + index = FigureCandidateIndex( + formal_legends=[b for b in blocks if b.get("role") == "figure_caption"], + held_legends=[], + rejected_legends=[], + deduped_legends=[b for b in blocks if b.get("role") == "figure_caption"], + candidate_groups=groups, + competing_caption_pages=set(), + sidecar_candidates={}, + bundle_source_legend_ids=set(), + locator_candidates=[], + ) + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + + # Run PrimarySamePagePass — should veto + report1 = PrimarySamePagePass().run(state) + assert len(report1.accepted) == 0, \ + f"PrimarySamePagePass should veto, got {len(report1.accepted)} accepted" + + # Run GroupSequentialPass — should split + report2 = GroupSequentialPass().run(state) + assert len(report2.accepted) == 2, \ + f"GroupSequentialPass should split, got {len(report2.accepted)} accepted" + assert len(state.matches) == 2, \ + f"expected 2 matches overall, got {len(state.matches)}" + + # Each match should have exactly one asset + for m in state.matches: + assert m["settlement_type"] == "group_sequential" + assert len(m.get("matched_assets", [])) == 1, \ + f"expected 1 asset per match, got {len(m.get('matched_assets', []))}" + + # c1 (Figure 2) should get a1 (left asset), c2 (Figure 3) should get a2 (right asset) + match_by_legend = {m["legend_block_id"]: m for m in state.matches} + assert match_by_legend["c1"]["asset_block_ids"] == ["a1"], \ + f"expected c1 -> a1, got {match_by_legend['c1']['asset_block_ids']}" + assert match_by_legend["c2"]["asset_block_ids"] == ["a2"], \ + f"expected c2 -> a2, got {match_by_legend['c2']['asset_block_ids']}" diff --git a/tests/test_ocr_figure_vnext_sidecar_pass.py b/tests/test_ocr_figure_vnext_sidecar_pass.py index de7b09d2..9404c3a4 100644 --- a/tests/test_ocr_figure_vnext_sidecar_pass.py +++ b/tests/test_ocr_figure_vnext_sidecar_pass.py @@ -75,6 +75,101 @@ def test_sidecar_pass_skips_pages_with_fewer_than_two_narrow_captions(): assert len(report.accepted) == 0 + +def test_single_narrow_sidecar_rescue_fires_for_left_caption_right_image(): + """Single narrow left caption with same-row right image → sidecar rescued.""" + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 1. Narrow caption with image to the right", + "bbox": [100, 100, 300, 200]}, # narrow, left column + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [320, 100, 600, 250]}, # right of caption, same-row + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + + # The single narrow caption must now appear in sidecar_candidates + assert 5 in index.sidecar_candidates + assert len(index.sidecar_candidates[5]) == 1 + assert index.sidecar_candidates[5][0]["block_id"] == "c1" + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = SidecarPass().run(state) + + assert len(report.accepted) >= 1 + sidecar_matches = [m for m in state.matches if m.get("settlement_type") == "sidecar"] + assert len(sidecar_matches) == 1 + assert sidecar_matches[0].get("legend_block_id") == "c1" + assert len(sidecar_matches[0].get("matched_assets", [])) >= 1 + assert "sidecar_match" in sidecar_matches[0].get("flags", []) + + +def test_single_narrow_sidecar_no_rescue_when_no_row_coupled_asset(): + """Single narrow caption whose asset is not row-coupled → not rescued.""" + blocks = [ + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 1. Narrow caption", + "bbox": [100, 300, 300, 340]}, # narrow, below asset + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [50, 0, 250, 90]}, # above, no y-overlap + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + + assert 5 not in index.sidecar_candidates + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = SidecarPass().run(state) + assert len(report.accepted) == 0 + + +def test_single_narrow_sidecar_does_not_break_grouped_baseline(): + """Three narrow captions plus single narrow caption — grouped baseline intact.""" + blocks = [ + # Three grouped narrow captions on page 5 (existing multi-caption path) + {"block_id": "c1", "page": 5, "role": "figure_caption", + "text": "Figure 1. First", "bbox": [100, 100, 300, 130]}, + {"block_id": "c2", "page": 5, "role": "figure_caption", + "text": "Figure 2. Second", "bbox": [100, 300, 310, 330]}, + {"block_id": "c3", "page": 5, "role": "figure_caption", + "text": "Figure 3. Third", "bbox": [100, 500, 300, 530]}, + {"block_id": "a1", "page": 5, "role": "figure_asset", + "bbox": [50, 0, 250, 90]}, + {"block_id": "a2", "page": 5, "role": "figure_asset", + "bbox": [50, 180, 260, 270]}, + {"block_id": "a3", "page": 5, "role": "figure_asset", + "bbox": [60, 400, 270, 480]}, + # Single narrow caption on page 6 with row-coupled asset (37-like) + {"block_id": "c4", "page": 6, "role": "figure_caption", + "text": "Figure 4. Side rescued", "bbox": [100, 200, 280, 350]}, + {"block_id": "a4", "page": 6, "role": "figure_asset", + "bbox": [300, 200, 600, 380]}, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + + # Page 5: the three grouped captions still form a sidecar candidate + assert 5 in index.sidecar_candidates + assert len(index.sidecar_candidates[5]) == 3 + assert {b["block_id"] for b in index.sidecar_candidates[5]} == {"c1", "c2", "c3"} + + # Page 6: single caption rescued via the 37-like path + assert 6 in index.sidecar_candidates + assert len(index.sidecar_candidates[6]) == 1 + assert index.sidecar_candidates[6][0]["block_id"] == "c4" + + state = FigurePipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger()) + report = SidecarPass().run(state) + + sidecar_matches = [m for m in state.matches if m.get("settlement_type") == "sidecar"] + assert len(sidecar_matches) == 4 # all four rescued + + page5_matches = [m for m in sidecar_matches if m.get("page") == 5] + page6_matches = [m for m in sidecar_matches if m.get("page") == 6] + assert len(page5_matches) == 3 + assert len(page6_matches) == 1 + assert page6_matches[0].get("legend_block_id") == "c4" + def test_sidecar_pass_does_not_steal_assets_from_protected_same_page_matches(): """Two narrow captions on page 5; c1 already has a protected same-page match.