From 4d78452acef1eabe66101eef5ae244ca60f40a18 Mon Sep 17 00:00:00 2001 From: LLLin000 <809867916@qq.com> Date: Fri, 3 Jul 2026 16:08:04 +0800 Subject: [PATCH] feat(ocr): populate sidecar_candidates and bundle_source_legend_ids in vnext index --- paperforge/worker/ocr_figure_vnext_corpus.py | 20 +++- tests/test_ocr_figure_vnext_corpus.py | 114 +++++++++++++++++++ 2 files changed, 132 insertions(+), 2 deletions(-) diff --git a/paperforge/worker/ocr_figure_vnext_corpus.py b/paperforge/worker/ocr_figure_vnext_corpus.py index ae6cb114..960efa8e 100644 --- a/paperforge/worker/ocr_figure_vnext_corpus.py +++ b/paperforge/worker/ocr_figure_vnext_corpus.py @@ -87,6 +87,22 @@ class FigureCandidateIndex: and ocr_figures._extract_figure_number(str(leg.get("text", ""))) is not None } + # Populate sidecar_candidates: pages with ≥2 aligned narrow captions. + sidecar_candidates: dict[int, list[dict]] = {} + _legends_by_page: dict[int, list[dict]] = {} + for leg in formal_legends: + _legends_by_page.setdefault(int(leg.get("page", 0) or 0), []).append(leg) + for sp, spl in _legends_by_page.items(): + narrow_set = ocr_figures._same_page_narrow_caption_column(spl, corpus.page_width) + if len(narrow_set) >= 2: + sidecar_candidates[sp] = narrow_set + + # 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( + formal_legends, corpus.raw_assets + ) + return cls( formal_legends=formal_legends, held_legends=[], @@ -94,7 +110,7 @@ class FigureCandidateIndex: deduped_legends=formal_legends, candidate_groups=candidate_groups, competing_caption_pages=competing_caption_pages, - sidecar_candidates={}, - bundle_source_legend_ids=set(), + sidecar_candidates=sidecar_candidates, + bundle_source_legend_ids=bundle_source_legend_ids, locator_candidates=corpus.locator_candidates, ) diff --git a/tests/test_ocr_figure_vnext_corpus.py b/tests/test_ocr_figure_vnext_corpus.py index 1eb2a623..c0582b2e 100644 --- a/tests/test_ocr_figure_vnext_corpus.py +++ b/tests/test_ocr_figure_vnext_corpus.py @@ -86,3 +86,117 @@ def test_rejected_legends_excluded_from_candidate_group_semantics() -> None: assert group.get("page_legend_count") == 1, ( f"Expected page_legend_count=1 (formal legends only), got {group.get('page_legend_count')}" ) + + +def test_candidate_index_populates_sidecar_candidates_for_narrow_caption_page() -> None: + """Three narrow numbered captions on one page → sidecar_candidates has that page.""" + blocks = [ + # Three narrow formal legends on page 5 — width=200, aligned at x-center=200 + { + "block_id": "leg1", + "page": 5, + "role": "figure_caption", + "text": "Figure 1. A caption", + "bbox": [100, 100, 300, 120], + }, + { + "block_id": "leg2", + "page": 5, + "role": "figure_caption", + "text": "Figure 2. Another caption", + "bbox": [100, 130, 300, 150], + }, + { + "block_id": "leg3", + "page": 5, + "role": "figure_caption", + "text": "Figure 3. Yet another caption", + "bbox": [100, 160, 300, 180], + }, + # One asset on page 5 so page has assets + { + "block_id": "ast1", + "page": 5, + "role": "figure_asset", + "bbox": [0, 200, 400, 500], + }, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + assert 5 in index.sidecar_candidates + assert len(index.sidecar_candidates[5]) >= 2 + + +def test_candidate_index_populates_bundle_source_legend_ids() -> None: + """Three numbered legends on a page with zero assets → bundle source ids populated.""" + blocks = [ + # Three numbered legends on page 3 — no assets on page 3 + { + "block_id": "leg1", + "page": 3, + "role": "figure_caption", + "text": "Figure 1. Caption A", + "bbox": [100, 100, 500, 120], + }, + { + "block_id": "leg2", + "page": 3, + "role": "figure_caption", + "text": "Figure 2. Caption B", + "bbox": [100, 130, 500, 150], + }, + { + "block_id": "leg3", + "page": 3, + "role": "figure_caption", + "text": "Figure 3. Caption C", + "bbox": [100, 160, 500, 180], + }, + # Assets on other pages + { + "block_id": "ast1", + "page": 4, + "role": "figure_asset", + "bbox": [0, 0, 400, 300], + }, + { + "block_id": "ast2", + "page": 5, + "role": "figure_asset", + "bbox": [0, 0, 400, 300], + }, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + assert len(index.bundle_source_legend_ids) >= 3 + + +def test_candidate_index_sidecar_candidates_empty_for_wide_captions() -> None: + """Wide captions (width > 720 on page_width=1200) → no sidecar candidates.""" + blocks = [ + # Three wide captions on page 5 — width=1000 > 720 + { + "block_id": "leg1", + "page": 5, + "role": "figure_caption", + "text": "Figure 1. Widest caption", + "bbox": [0, 100, 1000, 120], + }, + { + "block_id": "leg2", + "page": 5, + "role": "figure_caption", + "text": "Figure 2. Also wide", + "bbox": [0, 130, 1000, 150], + }, + { + "block_id": "leg3", + "page": 5, + "role": "figure_caption", + "text": "Figure 3. Still wide", + "bbox": [0, 160, 1000, 180], + }, + ] + corpus = FigureCorpus.from_blocks(blocks, page_width=1200) + index = FigureCandidateIndex.from_corpus(corpus) + assert index.sidecar_candidates == {}