From ea6a1f0378dd94f79e93dc0719709bceb4e02dc6 Mon Sep 17 00:00:00 2001
From: LLLin000 <809867916@qq.com>
Date: Fri, 3 Jul 2026 20:07:51 +0800
Subject: [PATCH] refactor(ocr): add table domain objects for Plan B
---
paperforge/worker/ocr_table_domain.py | 144 ++++++++++++++++++++++
tests/test_ocr_table_pairing_framework.py | 105 ++++++++++++++++
2 files changed, 249 insertions(+)
create mode 100644 paperforge/worker/ocr_table_domain.py
create mode 100644 tests/test_ocr_table_pairing_framework.py
diff --git a/paperforge/worker/ocr_table_domain.py b/paperforge/worker/ocr_table_domain.py
new file mode 100644
index 00000000..6072718d
--- /dev/null
+++ b/paperforge/worker/ocr_table_domain.py
@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class TableCorpus:
+ blocks: list[dict]
+ raw_captions: list[dict] = field(default_factory=list)
+ raw_assets: list[dict] = field(default_factory=list)
+ page_footnote_prior: dict[int, float] = field(default_factory=dict)
+ page_max_y: dict[int, float] = field(default_factory=dict)
+
+ @classmethod
+ def from_blocks(cls, blocks: list[dict]) -> TableCorpus:
+ from . import ocr_tables
+
+ raw_captions = [
+ b
+ for b in blocks
+ if b.get("role") in {"table_caption", "table_caption_candidate"}
+ or ocr_tables._is_validation_first_table_candidate(b)
+ ]
+ raw_assets = []
+ for block in blocks:
+ role = block.get("role", "")
+ raw_label = str(block.get("raw_label", "") or "").strip()
+ if role not in ("table_asset", "table_html", "media_asset", "figure_asset"):
+ continue
+ if role == "figure_asset" and raw_label != "table":
+ continue
+ raw_assets.append(block)
+
+ page_footnote_prior = ocr_tables._collect_page_footnote_prior(blocks)
+ page_max_y: dict[int, float] = {}
+ for block in blocks:
+ bbox = block.get("bbox") or [0, 0, 0, 0]
+ if len(bbox) >= 4:
+ page = int(block.get("page", 0) or 0)
+ page_max_y[page] = max(page_max_y.get(page, 0.0), float(bbox[3]))
+
+ return cls(
+ blocks=blocks,
+ raw_captions=raw_captions,
+ raw_assets=raw_assets,
+ page_footnote_prior=page_footnote_prior,
+ page_max_y=page_max_y,
+ )
+
+
+@dataclass
+class TableCandidateIndex:
+ caption_records: list[dict]
+ assets_by_page: dict[int, list[tuple[int, dict]]]
+
+ @classmethod
+ def from_corpus(cls, corpus: TableCorpus) -> TableCandidateIndex:
+ from . import ocr_tables
+
+ captions = sorted(
+ corpus.raw_captions,
+ key=lambda block: (
+ int(block.get("page", 0) or 0),
+ str(block.get("block_id", "")),
+ ),
+ )
+ caption_records = []
+ for caption in captions:
+ text = str(caption.get("text", "") or "")
+ caption_records.append(
+ {
+ "caption": caption,
+ "caption_block_id": str(caption.get("block_id", "")),
+ "caption_text": text,
+ "table_number": ocr_tables._extract_table_number(text),
+ "formal_table_number": ocr_tables._extract_base_table_number(text),
+ "is_continuation": ocr_tables._is_continuation_caption(text),
+ "is_validation_first_candidate": ocr_tables._is_validation_first_table_candidate(
+ caption
+ ),
+ "is_weak_truncated": ocr_tables._is_insufficient_table_caption_evidence(
+ caption
+ ),
+ "is_weak_explicit_caption": ocr_tables._is_weak_explicit_table_caption(
+ caption
+ ),
+ "continuation_ids": [],
+ "status": "pending",
+ "candidate_assets": [],
+ }
+ )
+
+ assets_by_page: dict[int, list[tuple[int, dict]]] = {}
+ for idx, asset in enumerate(corpus.raw_assets):
+ page = int(asset.get("page", 0) or 0)
+ assets_by_page.setdefault(page, []).append((idx, asset))
+
+ return cls(caption_records=caption_records, assets_by_page=assets_by_page)
+
+
+def assemble_table_inventory(
+ state, candidate_index: TableCandidateIndex
+) -> dict[str, Any]:
+ matched_caption_ids = {
+ t.get("caption_block_id", "")
+ for t in state.matches
+ if t.get("has_asset")
+ }
+ used_asset_ids = {
+ str(t.get("asset_block_id", ""))
+ for t in state.matches
+ if t.get("has_asset") and t.get("asset_block_id")
+ }
+ held_tables = [
+ record["caption"]
+ for record in candidate_index.caption_records
+ if record.get("status") == "held"
+ ]
+ unmatched_captions = [
+ record["caption"]
+ for record in candidate_index.caption_records
+ if record["caption_block_id"] not in matched_caption_ids
+ and record.get("status") != "held"
+ ]
+ unmatched_assets = [
+ asset
+ for page_assets in candidate_index.assets_by_page.values()
+ for _, asset in page_assets
+ if str(asset.get("block_id", "")) not in used_asset_ids
+ ]
+ return {
+ "tables": list(state.matches),
+ "held_tables": held_tables,
+ "unmatched_captions": unmatched_captions,
+ "unmatched_assets": unmatched_assets,
+ "official_table_count": len(
+ [
+ t
+ for t in state.matches
+ if t.get("has_asset") and not t.get("is_continuation")
+ ]
+ ),
+ }
diff --git a/tests/test_ocr_table_pairing_framework.py b/tests/test_ocr_table_pairing_framework.py
new file mode 100644
index 00000000..7e493f7c
--- /dev/null
+++ b/tests/test_ocr_table_pairing_framework.py
@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+
+def test_table_corpus_collects_captions_assets_and_page_context() -> None:
+ from paperforge.worker.ocr_table_domain import TableCorpus
+
+ blocks = [
+ {
+ "block_id": "cap1",
+ "page": 5,
+ "role": "table_caption",
+ "text": "Table 1. Example",
+ "bbox": [100, 100, 700, 140],
+ },
+ {
+ "block_id": "asset1",
+ "page": 5,
+ "role": "table_html",
+ "raw_label": "table",
+ "text": "",
+ "bbox": [100, 160, 700, 500],
+ },
+ {
+ "block_id": "note1",
+ "page": 5,
+ "role": "footnote",
+ "raw_label": "vision_footnote",
+ "text": "* p < 0.05",
+ "bbox": [100, 520, 300, 550],
+ },
+ ]
+
+ corpus = TableCorpus.from_blocks(blocks)
+
+ assert [b["block_id"] for b in corpus.raw_captions] == ["cap1"]
+ assert [b["block_id"] for b in corpus.raw_assets] == ["asset1"]
+ assert 5 in corpus.page_footnote_prior
+ assert 5 in corpus.page_max_y
+
+
+def test_table_candidate_index_materializes_caption_records_and_assets_by_page() -> None:
+ from paperforge.worker.ocr_table_domain import TableCandidateIndex, TableCorpus
+
+ blocks = [
+ {
+ "block_id": "cap2",
+ "page": 6,
+ "role": "table_caption_candidate",
+ "text": "Table 2. (continued)",
+ "bbox": [100, 100, 700, 130],
+ },
+ {
+ "block_id": "cap1",
+ "page": 5,
+ "role": "table_caption",
+ "text": "Table 1. Example",
+ "bbox": [100, 100, 700, 140],
+ },
+ {
+ "block_id": "asset1",
+ "page": 5,
+ "role": "table_html",
+ "raw_label": "table",
+ "text": "",
+ "bbox": [100, 160, 700, 500],
+ },
+ ]
+
+ index = TableCandidateIndex.from_corpus(TableCorpus.from_blocks(blocks))
+
+ assert [r["caption_block_id"] for r in index.caption_records] == ["cap1", "cap2"]
+ assert 5 in index.assets_by_page
+ assert index.caption_records[1]["is_continuation"] is True
+
+
+def test_assemble_table_inventory_preserves_public_shape_for_empty_state() -> None:
+ from paperforge.worker.ocr_pairing_state import OwnershipLedger, PipelineState
+ from paperforge.worker.ocr_table_domain import (
+ TableCandidateIndex,
+ TableCorpus,
+ assemble_table_inventory,
+ )
+
+ blocks = [
+ {
+ "block_id": "cap1",
+ "page": 1,
+ "role": "table_caption",
+ "text": "Table 1. Example",
+ "bbox": [0, 0, 10, 10],
+ }
+ ]
+ corpus = TableCorpus.from_blocks(blocks)
+ index = TableCandidateIndex.from_corpus(corpus)
+ state = PipelineState(corpus=corpus, candidate_index=index, ledger=OwnershipLedger())
+
+ inventory = assemble_table_inventory(state, index)
+
+ assert inventory == {
+ "tables": [],
+ "held_tables": [],
+ "unmatched_captions": [blocks[0]],
+ "unmatched_assets": [],
+ "official_table_count": 0,
+ }