From 174d2e8ccd1d64ea9e36471c44eae2b4eabf7b8c Mon Sep 17 00:00:00 2001 From: Research Assistant Date: Fri, 5 Jun 2026 15:11:20 +0800 Subject: [PATCH] fix: stabilize OCR table matching and image-first rendering --- paperforge/worker/ocr_tables.py | 125 ++++++++++++----- tests/test_ocr_tables.py | 228 +++++++++++++++++++++++++++++++- 2 files changed, 315 insertions(+), 38 deletions(-) diff --git a/paperforge/worker/ocr_tables.py b/paperforge/worker/ocr_tables.py index d610a78d..5cf8e810 100644 --- a/paperforge/worker/ocr_tables.py +++ b/paperforge/worker/ocr_tables.py @@ -11,6 +11,8 @@ _TABLE_PREFIX_PATTERN = re.compile( flags=re.IGNORECASE, ) +_CONTINUATION_PATTERN = re.compile(r"\(Continued\)|\(cont\.\)|\(cont\)", re.IGNORECASE) + def _extract_table_number(text: str) -> int | None: m = _TABLE_PREFIX_PATTERN.search(text) @@ -22,12 +24,31 @@ def _extract_table_number(text: str) -> int | None: return None +def _compute_overlap_score(a_bbox: list[float], b_bbox: list[float]) -> float: + if not a_bbox or not b_bbox or len(a_bbox) < 4 or len(b_bbox) < 4: + return 0.0 + ax1, ay1, ax2, ay2 = a_bbox + bx1, by1, bx2, by2 = b_bbox + ix1 = max(ax1, bx1) + iy1 = max(ay1, by1) + ix2 = min(ax2, bx2) + iy2 = min(ay2, by2) + iw = max(0, ix2 - ix1) + ih = max(0, iy2 - iy1) + intersection = iw * ih + if intersection == 0: + return 0.0 + a_area = (ax2 - ax1) * (ay2 - ay1) + b_area = (bx2 - bx1) * (by2 - by1) + union = a_area + b_area - intersection + if union == 0: + return 0.0 + return intersection / union + + def build_table_inventory(structured_blocks: list[dict]) -> dict[str, Any]: - tables: list[dict] = [] captions: list[dict] = [] assets: list[dict] = [] - unmatched_captions: list[dict] = [] - unmatched_assets: list[dict] = [] for block in structured_blocks: role = block.get("role", "") @@ -40,52 +61,90 @@ def build_table_inventory(structured_blocks: list[dict]) -> dict[str, Any]: assets.append(block) used_asset_indices: set[int] = set() + tables: list[dict] = [] + unmatched_captions: list[dict] = [] + for caption in captions: caption_page = caption.get("page", 0) - caption_text = caption.get("text", "") + caption_text = caption.get("text", "") or "" + caption_bbox = caption.get("bbox", [0, 0, 0, 0]) table_num = _extract_table_number(caption_text) + is_continuation = bool(_CONTINUATION_PATTERN.search(caption_text)) - candidate_pages = [caption_page, caption_page + 1, caption_page - 1] - matched_asset = None + # Continuations should only match same-page assets + candidate_pages = [caption_page] + if not is_continuation: + candidate_pages.extend([caption_page + 1, caption_page - 1]) + best_match: tuple[int | None, float] = (None, -999) for page in candidate_pages: if page < 1: continue - for i, asset in enumerate(assets): - if asset.get("page", 0) == page and i not in used_asset_indices: - matched_asset = asset - used_asset_indices.add(i) - break - if matched_asset: - break + page_assets = [ + (i, a) for i, a in enumerate(assets) + if a.get("page", 0) == page and i not in used_asset_indices + ] + for i, asset in page_assets: + asset_bbox = asset.get("bbox", [0, 0, 0, 0]) + overlap = _compute_overlap_score(caption_bbox, asset_bbox) + if is_continuation: + same_page = 3.0 + h_align = 0.0 + if len(caption_bbox) >= 4 and len(asset_bbox) >= 4: + ccx = (caption_bbox[0] + caption_bbox[2]) / 2 + acx = (asset_bbox[0] + asset_bbox[2]) / 2 + h_align = max(0, 1.0 - abs(ccx - acx) / max(1, caption_bbox[2] - caption_bbox[0])) + score = same_page + h_align * 3 + overlap * 5 + else: + dist_y = 0.0 + if len(caption_bbox) >= 4 and len(asset_bbox) >= 4: + dist_y = abs((caption_bbox[1] + caption_bbox[3]) / 2 - (asset_bbox[1] + asset_bbox[3]) / 2) + same_page = 2.0 if page == caption_page else 0.0 + direction = 1.0 if (len(asset_bbox) >= 4 and len(caption_bbox) >= 4 + and (asset_bbox[1] + asset_bbox[3]) / 2 < (caption_bbox[1] + caption_bbox[3]) / 2) else 0.5 + h_align = 0.0 + if len(caption_bbox) >= 4 and len(asset_bbox) >= 4: + ccx = (caption_bbox[0] + caption_bbox[2]) / 2 + acx = (asset_bbox[0] + asset_bbox[2]) / 2 + h_align = max(0, 1.0 - abs(ccx - acx) / max(1, caption_bbox[2] - caption_bbox[0])) + a_w = asset_bbox[2] - asset_bbox[0] if len(asset_bbox) >= 4 else 0 + a_h = asset_bbox[3] - asset_bbox[1] if len(asset_bbox) >= 4 else 0 + size_score = min(2.0, (a_w * a_h) / 50000) + score = overlap * 10 + same_page + direction + h_align * 2 + size_score * 0.5 - dist_y * 0.01 + if score > best_match[1]: + best_match = (i, score) - tables.append( - { - "caption_block_id": caption.get("block_id", ""), - "page": caption_page, - "caption_text": caption_text, - "table_number": table_num, - "asset_block_id": matched_asset.get("block_id", "") if matched_asset else "", - "asset_bbox": matched_asset.get("bbox", [0, 0, 0, 0]) if matched_asset else [], - "assistive_text": (matched_asset.get("text", "") or "") if matched_asset else "", - "truth_source": "image", - "has_asset": matched_asset is not None, - } - ) + matched_asset = assets[best_match[0]] if best_match[0] is not None else None + if best_match[0] is not None: + used_asset_indices.add(best_match[0]) - for caption in captions: - if not any(t["caption_block_id"] == caption.get("block_id", "") for t in tables if t["has_asset"]): + if matched_asset is None: unmatched_captions.append(caption) - for i, asset in enumerate(assets): - if i not in used_asset_indices: - unmatched_assets.append(asset) + tables.append({ + "caption_block_id": caption.get("block_id", ""), + "page": caption_page, + "caption_text": caption_text, + "table_number": table_num, + "is_continuation": is_continuation, + "asset_block_id": matched_asset.get("block_id", "") if matched_asset else "", + "asset_bbox": matched_asset.get("bbox", [0, 0, 0, 0]) if matched_asset else [], + "assistive_text": (matched_asset.get("text", "") or "") if matched_asset else "", + "truth_source": "image", + "has_asset": matched_asset is not None, + }) + + unmatched_assets_list = [ + a for i, a in enumerate(assets) if i not in used_asset_indices + ] + + official_count = len([t for t in tables if t["has_asset"] and not t["is_continuation"]]) return { "tables": tables, "unmatched_captions": unmatched_captions, - "unmatched_assets": unmatched_assets, - "official_table_count": len([t for t in tables if t["has_asset"]]), + "unmatched_assets": unmatched_assets_list, + "official_table_count": official_count, } diff --git a/tests/test_ocr_tables.py b/tests/test_ocr_tables.py index 6edd493c..e108fe8d 100644 --- a/tests/test_ocr_tables.py +++ b/tests/test_ocr_tables.py @@ -1,8 +1,4 @@ -"""Phase 2 contract tests for table inventory. - -paperforge.worker.ocr_tables does not exist yet -- tests will fail until -Task 7 implements the module. -""" +"""Task 5 tests for hardened table matching.""" from __future__ import annotations @@ -63,3 +59,225 @@ def test_table_without_asset_is_tracked_as_unmatched_caption() -> None: assert inventory["official_table_count"] == 0 assert len(inventory["unmatched_captions"]) == 1 + + +def test_continuation_table_matches_same_page_asset() -> None: + from paperforge.worker.ocr_tables import build_table_inventory + + structured_blocks = [ + { + "paper_id": "KEY001", + "page": 10, + "block_id": "p10_a1", + "role": "table_asset", + "text": "continued data", + "bbox": [100, 100, 600, 400], + }, + { + "paper_id": "KEY001", + "page": 10, + "block_id": "p10_c1", + "role": "table_caption", + "text": "Table 1 (Continued)", + "bbox": [100, 420, 600, 460], + }, + ] + + inventory = build_table_inventory(structured_blocks) + + assert len(inventory["tables"]) == 1 + t = inventory["tables"][0] + assert t["is_continuation"] is True + assert t["has_asset"] is True + assert t["asset_block_id"] == "p10_a1" + + +def test_continuation_does_not_increment_official_count() -> None: + from paperforge.worker.ocr_tables import build_table_inventory + + structured_blocks = [ + { + "paper_id": "KEY001", + "page": 5, + "block_id": "p5_a1", + "role": "table_asset", + "text": "table data", + "bbox": [100, 100, 600, 400], + }, + { + "paper_id": "KEY001", + "page": 5, + "block_id": "p5_c1", + "role": "table_caption", + "text": "Table 1. Main data", + "bbox": [100, 420, 600, 460], + }, + { + "paper_id": "KEY001", + "page": 6, + "block_id": "p6_a1", + "role": "table_asset", + "text": "continued data", + "bbox": [100, 100, 600, 300], + }, + { + "paper_id": "KEY001", + "page": 6, + "block_id": "p6_c1", + "role": "table_caption", + "text": "Table 1 (Continued)", + "bbox": [100, 320, 600, 360], + }, + ] + + inventory = build_table_inventory(structured_blocks) + + assert inventory["official_table_count"] == 1 + normal = [t for t in inventory["tables"] if not t["is_continuation"]] + continued = [t for t in inventory["tables"] if t["is_continuation"]] + assert len(normal) == 1 + assert len(continued) == 1 + assert normal[0]["has_asset"] is True + assert continued[0]["has_asset"] is True + + +def test_continuation_without_asset_has_has_asset_false() -> None: + from paperforge.worker.ocr_tables import build_table_inventory + + structured_blocks = [ + { + "paper_id": "KEY001", + "page": 10, + "block_id": "p10_c1", + "role": "table_caption", + "text": "Table 1 (Continued)", + "bbox": [100, 100, 600, 140], + }, + ] + + inventory = build_table_inventory(structured_blocks) + + assert len(inventory["tables"]) == 1 + assert inventory["tables"][0]["is_continuation"] is True + assert inventory["tables"][0]["has_asset"] is False + assert len(inventory["unmatched_captions"]) == 1 + + +def test_multi_signal_scoring_prefers_better_asset() -> None: + from paperforge.worker.ocr_tables import build_table_inventory + + structured_blocks = [ + { + "paper_id": "KEY001", + "page": 3, + "block_id": "p3_a1", + "role": "table_asset", + "text": "far table", + "bbox": [50, 600, 300, 800], + }, + { + "paper_id": "KEY001", + "page": 3, + "block_id": "p3_a2", + "role": "table_asset", + "text": "near table", + "bbox": [50, 50, 550, 300], + }, + { + "paper_id": "KEY001", + "page": 3, + "block_id": "p3_c1", + "role": "table_caption", + "text": "Table 3. Nearby data", + "bbox": [50, 310, 550, 350], + }, + ] + + inventory = build_table_inventory(structured_blocks) + + assert inventory["official_table_count"] == 1 + t = inventory["tables"][0] + assert t["asset_block_id"] == "p3_a2" + + +def test_continuation_matches_only_same_page_not_adjacent() -> None: + from paperforge.worker.ocr_tables import build_table_inventory + + structured_blocks = [ + { + "paper_id": "KEY001", + "page": 4, + "block_id": "p4_a1", + "role": "table_asset", + "text": "wrong page", + "bbox": [100, 100, 600, 400], + }, + { + "paper_id": "KEY001", + "page": 5, + "block_id": "p5_a1", + "role": "table_asset", + "text": "same page", + "bbox": [100, 100, 600, 400], + }, + { + "paper_id": "KEY001", + "page": 5, + "block_id": "p5_c1", + "role": "table_caption", + "text": "Table 2 (cont.)", + "bbox": [100, 420, 600, 460], + }, + ] + + inventory = build_table_inventory(structured_blocks) + + assert len(inventory["tables"]) == 1 + t = inventory["tables"][0] + assert t["is_continuation"] is True + assert t["asset_block_id"] == "p5_a1" + + +def test_multiple_captions_match_correct_assets_in_order() -> None: + from paperforge.worker.ocr_tables import build_table_inventory + + structured_blocks = [ + { + "paper_id": "KEY001", + "page": 5, + "block_id": "a5b", + "role": "table_asset", + "text": "table 1 body", + "bbox": [100, 100, 600, 400], + }, + { + "paper_id": "KEY001", + "page": 5, + "block_id": "c5a", + "role": "table_caption", + "text": "Table 1. First", + "bbox": [100, 420, 600, 460], + }, + { + "paper_id": "KEY001", + "page": 6, + "block_id": "a6b", + "role": "table_asset", + "text": "table 2 body", + "bbox": [100, 100, 600, 400], + }, + { + "paper_id": "KEY001", + "page": 6, + "block_id": "c6a", + "role": "table_caption", + "text": "Table 2. Second", + "bbox": [100, 420, 600, 460], + }, + ] + + inventory = build_table_inventory(structured_blocks) + + assert inventory["official_table_count"] == 2 + assert inventory["tables"][0]["asset_block_id"] == "a5b" + assert inventory["tables"][1]["asset_block_id"] == "a6b"