mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
refactor(ocr): add table domain objects for Plan B
This commit is contained in:
LLLin000
parent
db01518b31
commit
ea6a1f0378
2 changed files with 249 additions and 0 deletions
144
paperforge/worker/ocr_table_domain.py
Normal file
144
paperforge/worker/ocr_table_domain.py
Normal file
|
|
@ -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")
|
||||
]
|
||||
),
|
||||
}
|
||||
105
tests/test_ocr_table_pairing_framework.py
Normal file
105
tests/test_ocr_table_pairing_framework.py
Normal file
|
|
@ -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": "<table><tr><td>x</td></tr></table>",
|
||||
"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": "<table></table>",
|
||||
"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,
|
||||
}
|
||||
Loading…
Reference in a new issue